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

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 {

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(())
     }
 

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<String>) {
         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()

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,
+    }))
+}

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<Vec<(String, String)>> {
+        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,

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 <name>"),
                 },