diff options
| -rw-r--r-- | config.example.yaml | 4 | ||||
| -rw-r--r-- | src/config/mod.rs | 41 | ||||
| -rw-r--r-- | src/config/subquery.rs | 42 |
3 files changed, 84 insertions, 3 deletions
diff --git a/config.example.yaml b/config.example.yaml index ae249b6..4cac6c0 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -19,7 +19,9 @@ mapping_tools: # Alias for a tool or toolset use_tools: null # Which tools to use by default. (e.g. 'fs,web_search') subquery: false # Enable the built-in subquery tool, letting the model spawn fresh context-limited LLM queries. subquery_max_depth: 3 # Maximum nesting depth of subqueries (the main query counts as depth 1). - +subquery_compression: false # Compress subquery output before handing it to the parent query. +subquery_compression_prompt: null # Prompt used to compress subquery output (null = built-in prompt). +subquery_compression_threshold: 1000 # Token threshold above which subquery output is compressed (0 disables). # ---- prelude ---- repl_prelude: null # Set a default role or session for REPL mode (e.g. role:<name>, session:<name>, <session>:<role>) cmd_prelude: null # Set a default role or session for CMD mode (e.g. role:<name>, session:<name>, <session>:<role>) diff --git a/src/config/mod.rs b/src/config/mod.rs index 693c9f5..ff7a22b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -73,6 +73,8 @@ const SYNC_MODELS_URL: &str = const SUMMARIZE_PROMPT: &str = "Summarize the discussion briefly in 200 words or less to use as a prompt for future context."; const SUMMARY_PROMPT: &str = "This is a summary of the chat history as a recap: "; +const SUBQUERY_COMPRESSION_PROMPT: &str = + "Summarize the subquery thought process by keeping all key facts, numbers, names, identifiers needed to answer follow-up questions. Leave the final result of the subquery unchanged. Reply with the full text directly."; const RAG_TEMPLATE: &str = r#"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags) @@ -122,6 +124,10 @@ pub struct Config { pub subquery: bool, pub subquery_max_depth: usize, + pub subquery_compression: bool, + pub subquery_compression_prompt: Option<String>, + pub subquery_compression_threshold: usize, + pub repl_prelude: Option<String>, pub cmd_prelude: Option<String>, pub agent_prelude: Option<String>, @@ -201,6 +207,10 @@ impl Default for Config { subquery: false, subquery_max_depth: 3, + subquery_compression: false, + subquery_compression_prompt: None, + subquery_compression_threshold: 1000, + repl_prelude: None, cmd_prelude: None, agent_prelude: None, @@ -609,6 +619,14 @@ impl Config { ("function_calling", self.function_calling.to_string()), ("subquery", self.subquery.to_string()), ("subquery_max_depth", self.subquery_max_depth.to_string()), + ( + "subquery_compression", + self.subquery_compression.to_string(), + ), + ( + "subquery_compression_threshold", + self.subquery_compression_threshold.to_string(), + ), ("stream", self.stream.to_string()), ("save", self.save.to_string()), ("keybindings", self.keybindings.clone()), @@ -695,6 +713,18 @@ impl Config { let value = value.parse().with_context(|| "Invalid value")?; config.write().subquery_max_depth = value; } + "subquery_compression" => { + let value = value.parse().with_context(|| "Invalid value")?; + config.write().subquery_compression = value; + } + "subquery_compression_prompt" => { + let value: Option<String> = parse_value(value)?; + config.write().subquery_compression_prompt = value; + } + "subquery_compression_threshold" => { + let value = value.parse().with_context(|| "Invalid value")?; + config.write().subquery_compression_threshold = value; + } "stream" => { let value = value.parse().with_context(|| "Invalid value")?; config.write().stream = value; @@ -2336,6 +2366,17 @@ impl Config { if let Some(v) = read_env_value::<String>(&get_env_name("summary_prompt")) { self.summary_prompt = v; } + if let Some(Some(v)) = read_env_bool(&get_env_name("subquery_compression")) { + self.subquery_compression = v; + } + if let Some(v) = read_env_value::<String>(&get_env_name("subquery_compression_prompt")) { + self.subquery_compression_prompt = v; + } + if let Some(Some(v)) = + read_env_value::<usize>(&get_env_name("subquery_compression_threshold")) + { + self.subquery_compression_threshold = v; + } if let Some(v) = read_env_value::<String>(&get_env_name("rag_embedding_model")) { self.rag_embedding_model = v; diff --git a/src/config/subquery.rs b/src/config/subquery.rs index 2b23376..e3716a9 100644 --- a/src/config/subquery.rs +++ b/src/config/subquery.rs @@ -1,4 +1,4 @@ -use super::{GlobalConfig, Input, RoleLike}; +use super::{GlobalConfig, Input, RoleLike, SUBQUERY_COMPRESSION_PROMPT}; use crate::client::{call_chat_completions, call_chat_completions_streaming, Message, MessageRole}; use crate::function::ToolResult; @@ -168,7 +168,10 @@ pub async fn eval_subquery( }) .await; match result { - Ok(text) => Ok(json!({ "output": text })), + Ok(text) => { + let text = maybe_compress_subquery_output(config, text).await; + Ok(json!({ "output": text })) + } Err(err) => Ok(json!({ "error": format!("{err:?}") })), } } @@ -394,3 +397,38 @@ fn tool_result_messages(text: &str, tool_results: &[ToolResult]) -> Vec<Message> )); messages } + +async fn maybe_compress_subquery_output(config: &GlobalConfig, output: String) -> String { + if !config.read().subquery_compression { + return output; + } + let threshold = config.read().subquery_compression_threshold; + if threshold == 0 || count_tokens(config, &output) <= threshold { + return output; + } + match compress_subquery_output(config, &output).await { + Ok(compressed) if !compressed.trim().is_empty() => compressed, + _ => output, + } +} + +fn count_tokens(config: &GlobalConfig, text: &str) -> usize { + let input = Input::from_str(config, "", None); + let messages = [Message::new( + MessageRole::User, + crate::client::MessageContent::Text(text.to_string()), + )]; + input.role().model().total_tokens(&messages) +} + +async fn compress_subquery_output(config: &GlobalConfig, output: &str) -> Result<String> { + let prompt = config + .read() + .subquery_compression_prompt + .clone() + .unwrap_or_else(|| SUBQUERY_COMPRESSION_PROMPT.to_string()); + let role = config.read().extract_role(); + let text = format!("{prompt}\n\n{output}"); + let input = Input::from_str(config, &text, Some(role)); + input.fetch_chat_text().await +} |