Implemented automatic subquery compressionmain

1 files changed, 40 insertions, 2 deletions

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