From cc31b93fa43c04cc706a26bae64daa5d55c8b249 Mon Sep 17 00:00:00 2001 From: Leonard Kugis Date: Sat, 3 Oct 2026 13:55:21 +0200 Subject: Implemented automatic subquery compression --- src/config/mod.rs | 41 +++++++++++++++++++++++++++++++++++++++++ src/config/subquery.rs | 42 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 81 insertions(+), 2 deletions(-) (limited to 'src/config') 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, + pub subquery_compression_threshold: usize, + pub repl_prelude: Option, pub cmd_prelude: Option, pub agent_prelude: Option, @@ -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 = 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::(&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::(&get_env_name("subquery_compression_prompt")) { + self.subquery_compression_prompt = v; + } + if let Some(Some(v)) = + read_env_value::(&get_env_name("subquery_compression_threshold")) + { + self.subquery_compression_threshold = v; + } if let Some(v) = read_env_value::(&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 )); 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 { + 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 +} -- cgit v1.2.3