From 68135b97d172565a023d9c6c5b0ae88ffe35a7f7 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 12 Aug 2026 16:25:50 -0600 Subject: [PATCH] feat: append new built-in rag__query function to RAG contexts to allow further querying by LLMs --- src/config/agent.rs | 4 ++ src/config/request_context.rs | 23 ++++++-- src/function/mod.rs | 21 +++++++ src/function/rag_query.rs | 101 ++++++++++++++++++++++++++++++++++ src/rag/mod.rs | 14 +++++ src/repl/mod.rs | 2 +- 6 files changed, 159 insertions(+), 6 deletions(-) create mode 100644 src/function/rag_query.rs diff --git a/src/config/agent.rs b/src/config/agent.rs index 6b993dc..95467ce 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -248,6 +248,10 @@ impl Agent { } } + if rag.is_some() && app.function_calling_support && graph_for_rag.is_none() { + functions.append_rag_query_functions(); + } + agent_config.replace_tools_placeholder(&functions); Ok(Self { diff --git a/src/config/request_context.rs b/src/config/request_context.rs index c0eea43..6f9c4ed 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -16,8 +16,9 @@ use super::{MessageContentToolCalls, prompts}; use crate::client::{Model, ModelType, list_models}; use crate::function::{ FunctionDeclaration, Functions, ToolCallTracker, ToolResult, memory::MEMORY_FUNCTION_PREFIX, - skill::SKILL_FUNCTION_PREFIX, supervisor::SUPERVISOR_FUNCTION_PREFIX, - todo::TODO_FUNCTION_PREFIX, user_interaction::USER_FUNCTION_PREFIX, + rag_query::RAG_FUNCTION_PREFIX, skill::SKILL_FUNCTION_PREFIX, + supervisor::SUPERVISOR_FUNCTION_PREFIX, todo::TODO_FUNCTION_PREFIX, + user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, @@ -715,6 +716,7 @@ impl RequestContext { pub fn exit_rag(&mut self) -> Result<()> { self.rag.take(); + self.tool_scope.functions.remove_rag_query_functions(); Ok(()) } @@ -1137,6 +1139,7 @@ impl RequestContext { && !v.name.starts_with("agent__") && !v.name.starts_with("memory__") && !v.name.starts_with("skill__") + && !v.name.starts_with("rag__") }) .map(|v| v.name.clone()) .collect() @@ -1957,7 +1960,8 @@ impl RequestContext { || (!matches!(role.skills_enabled(), Some(false)) && v.name.starts_with(SKILL_FUNCTION_PREFIX)) || (self.auto_continue_config().enabled - && v.name.starts_with(TODO_FUNCTION_PREFIX))) + && v.name.starts_with(TODO_FUNCTION_PREFIX)) + || v.name.starts_with(RAG_FUNCTION_PREFIX)) && !existing.contains(&v.name) }) .cloned() @@ -1987,6 +1991,7 @@ impl RequestContext { || v.name.starts_with(TODO_FUNCTION_PREFIX) || v.name.starts_with(SUPERVISOR_FUNCTION_PREFIX) || v.name.starts_with(MEMORY_FUNCTION_PREFIX) + || v.name.starts_with(RAG_FUNCTION_PREFIX) }); } @@ -3467,6 +3472,12 @@ impl RequestContext { if self.should_register_memory_tools() { functions.append_memory_functions(); } + if self.rag.is_some() + && app.function_calling_support + && !self.agent.as_ref().is_some_and(|a| a.is_graph()) + { + functions.append_rag_query_functions(); + } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { @@ -4136,7 +4147,7 @@ impl RequestContext { super::TEMP_RAG_NAME, &rag_path, &[], - abort_signal, + abort_signal.clone(), false, ) .await?, @@ -4172,10 +4183,11 @@ impl RequestContext { }; self.rag = Some(rag); self.rag_key = rag_key; + self.refresh_tool_scope(abort_signal).await?; Ok(()) } - pub async fn attach_rag(&mut self, name: &str) -> Result<()> { + pub async fn attach_rag(&mut self, name: &str, abort_signal: AbortSignal) -> Result<()> { let rag_path = self.rag_file(name); if rag_path.exists() { bail!( @@ -4192,6 +4204,7 @@ impl RequestContext { self.rag_cache().insert(key.clone(), &rag); self.rag = Some(rag); self.rag_key = Some(key); + self.refresh_tool_scope(abort_signal).await?; Ok(()) } diff --git a/src/function/mod.rs b/src/function/mod.rs index 19b85cf..2545fb0 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod memory; +pub(crate) mod rag_query; pub(crate) mod skill; pub(crate) mod supervisor; pub(crate) mod todo; @@ -23,6 +24,7 @@ use futures_util::future; use indexmap::IndexMap; use indoc::formatdoc; use memory::MEMORY_FUNCTION_PREFIX; +use rag_query::RAG_FUNCTION_PREFIX; use rust_embed::Embed; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; @@ -495,6 +497,16 @@ impl Functions { .extend(user_interaction::user_interaction_function_declarations()); } + pub fn append_rag_query_functions(&mut self) { + self.declarations + .extend(rag_query::rag_query_function_declarations()); + } + + pub fn remove_rag_query_functions(&mut self) { + self.declarations + .retain(|f| !f.name.starts_with(RAG_FUNCTION_PREFIX)); + } + pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { let mut invoke_function_properties = IndexMap::new(); invoke_function_properties.insert( @@ -1252,6 +1264,15 @@ impl ToolCall { json!({"tool_call_error": error_msg}) }) } + _ if cmd_name.starts_with(RAG_FUNCTION_PREFIX) => { + rag_query::handle_rag_tool(ctx, &cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("RAG query failed: {e}"); + eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) + } _ => match run_llm_function(cmd_name, cmd_args, envs, agent_name) { Ok(Some(contents)) => serde_json::from_str(&contents) .ok() diff --git a/src/function/rag_query.rs b/src/function/rag_query.rs new file mode 100644 index 0000000..970c6d6 --- /dev/null +++ b/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 { + 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 { + 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 { + 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 = 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, + })) +} diff --git a/src/rag/mod.rs b/src/rag/mod.rs index a71d7b7..bf44653 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -856,6 +856,20 @@ impl Rag { Ok((embeddings, sources, ids)) } + pub async fn search_chunks( + &self, + text: &str, + top_k: usize, + rerank_model: Option<&str>, + ) -> Result> { + let results = self.hybrid_search(text, top_k, rerank_model).await?; + + Ok(results + .into_iter() + .map(|(id, content)| (content, self.resolve_source(&id))) + .collect()) + } + pub async fn search_with_template( &self, app: &AppConfig, diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 7ade0a6..13d2c50 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -892,7 +892,7 @@ pub async fn run_repl_command( ".rag" => match split_first_arg(args) { Some(("attach", rest)) => match rest { Some(name) if !name.trim().is_empty() => { - ctx.attach_rag(name.trim()).await?; + ctx.attach_rag(name.trim(), abort_signal.clone()).await?; } _ => println!("Usage: .rag attach "), },