102 lines
3.3 KiB
Rust
102 lines
3.3 KiB
Rust
use super::{FunctionDeclaration, JsonSchema};
|
|
use crate::config::RequestContext;
|
|
|
|
use anyhow::{Result, anyhow};
|
|
use indexmap::IndexMap;
|
|
use serde_json::{Value, json};
|
|
|
|
pub const RAG_FUNCTION_PREFIX: &str = "rag__";
|
|
|
|
pub fn rag_query_function_declarations() -> Vec<FunctionDeclaration> {
|
|
vec![FunctionDeclaration {
|
|
name: format!("{RAG_FUNCTION_PREFIX}query"),
|
|
description: "Search the RAG knowledge base attached to this session and return \
|
|
the most relevant text chunks with their source paths. The relevant \
|
|
context has already been injected into the prompt up-front; use this \
|
|
tool to pull additional context on-demand when the initial retrieval \
|
|
does not fully answer the question. Prefer specific, keyword-rich queries."
|
|
.to_string(),
|
|
parameters: JsonSchema {
|
|
type_value: Some("object".to_string()),
|
|
properties: Some(IndexMap::from([
|
|
(
|
|
"query".to_string(),
|
|
JsonSchema {
|
|
type_value: Some("string".to_string()),
|
|
description: Some(
|
|
"Natural language search query used to retrieve relevant chunks."
|
|
.into(),
|
|
),
|
|
..Default::default()
|
|
},
|
|
),
|
|
(
|
|
"top_k".to_string(),
|
|
JsonSchema {
|
|
type_value: Some("integer".to_string()),
|
|
description: Some(
|
|
"Maximum number of chunks to return. Defaults to the RAG's \
|
|
configured top_k when omitted."
|
|
.into(),
|
|
),
|
|
..Default::default()
|
|
},
|
|
),
|
|
])),
|
|
required: Some(vec!["query".to_string()]),
|
|
..Default::default()
|
|
},
|
|
agent: false,
|
|
}]
|
|
}
|
|
|
|
pub async fn handle_rag_tool(
|
|
ctx: &mut RequestContext,
|
|
cmd_name: &str,
|
|
args: &Value,
|
|
) -> Result<Value> {
|
|
let action = cmd_name
|
|
.strip_prefix(RAG_FUNCTION_PREFIX)
|
|
.unwrap_or(cmd_name);
|
|
|
|
match action {
|
|
"query" => handle_query(ctx, args).await,
|
|
_ => Err(anyhow!("Unknown RAG action: {action}")),
|
|
}
|
|
}
|
|
|
|
async fn handle_query(ctx: &RequestContext, args: &Value) -> Result<Value> {
|
|
let rag = ctx
|
|
.rag
|
|
.clone()
|
|
.ok_or_else(|| anyhow!("No RAG is attached to this session"))?;
|
|
|
|
let query = args
|
|
.get("query")
|
|
.and_then(Value::as_str)
|
|
.ok_or_else(|| anyhow!("'query' is required"))?;
|
|
|
|
let top_k = args
|
|
.get("top_k")
|
|
.and_then(Value::as_u64)
|
|
.map(|v| v as usize)
|
|
.unwrap_or_else(|| rag.configured_top_k());
|
|
|
|
let rerank_model = rag.configured_reranker().map(|s| s.to_string());
|
|
|
|
let chunks = rag
|
|
.search_chunks(query, top_k, rerank_model.as_deref())
|
|
.await?;
|
|
|
|
let chunks_json: Vec<Value> = chunks
|
|
.into_iter()
|
|
.map(|(text, source)| json!({ "text": text, "source": source }))
|
|
.collect();
|
|
|
|
Ok(json!({
|
|
"rag_name": rag.name(),
|
|
"count": chunks_json.len(),
|
|
"chunks": chunks_json,
|
|
}))
|
|
}
|