feat: append new built-in rag__query function to RAG contexts to allow further querying by LLMs

src/function/rag_query.rs

@@ -0,0 +1,101 @@
+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,
+    }))
+}