From 5d0f9e58b0f4471998e180f558a200e0536baa7c Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Sat, 27 Sep 2025 23:58:39 +0200 Subject: [PATCH 01/12] feat: add token usage display Signed-off-by: Pierric Buchez --- shai-cli/src/tui/app.rs | 13 +++++++- shai-cli/src/tui/command.rs | 10 ++++++ shai-cli/src/tui/helper.rs | 5 +-- shai-core/src/agent/actions/brain.rs | 12 +++++-- shai-core/src/agent/brain.rs | 28 ++++++++++++++--- shai-core/src/agent/events.rs | 11 +++++++ shai-core/src/agent/output/log.rs | 3 ++ shai-core/src/agent/output/pretty.rs | 4 +++ shai-core/src/runners/coder/coder.rs | 23 +++++++++++--- shai-llm/src/providers/openai_compatible.rs | 33 ++++++++++++++++++- shai-llm/src/providers/ovhcloud.rs | 35 +++++++++++++++++++-- 11 files changed, 160 insertions(+), 17 deletions(-) diff --git a/shai-cli/src/tui/app.rs b/shai-cli/src/tui/app.rs index 6595b02..52d98fd 100644 --- a/shai-cli/src/tui/app.rs +++ b/shai-cli/src/tui/app.rs @@ -61,7 +61,7 @@ pub struct App<'a> { pub(crate) agent: Option, pub(crate) custom_agent: Option>, - + pub(crate) state: AppModalState<'a>, pub(crate) formatter: PrettyFormatter, // streaming log formatter pub(crate) running_tools: HashMap, // (request_id, request) @@ -69,6 +69,9 @@ pub struct App<'a> { pub(crate) commands: HashMap<(String, String),Vec>, pub(crate) exit: bool, pub(crate) permission_queue: VecDeque<(String, PermissionRequest)>, // (request_id, request) + + pub(crate) total_input_tokens: u32, + pub(crate) total_output_tokens: u32, } @@ -150,6 +153,12 @@ impl App<'_> { if let AgentEvent::PermissionRequired { request_id, request } = &event { self.permission_queue.push_back((request_id.clone(), request.clone())); } + + // Handle token usage tracking + if let AgentEvent::TokenUsage { input_tokens, output_tokens } = &event { + self.total_input_tokens += input_tokens; + self.total_output_tokens += output_tokens; + } Ok(()) } @@ -171,6 +180,8 @@ impl App<'_> { exit: false, running_tools: HashMap::new(), permission_queue: VecDeque::new(), + total_input_tokens: 0, + total_output_tokens: 0, } } diff --git a/shai-cli/src/tui/command.rs b/shai-cli/src/tui/command.rs index 990c8c6..7288c80 100644 --- a/shai-cli/src/tui/command.rs +++ b/shai-cli/src/tui/command.rs @@ -9,6 +9,7 @@ impl App<'_> { (("/exit","exit from the tui"), vec![]), (("/auth","select a provider"), vec![]), (("/tc","set the tool call method: [fc | fc2 | so]"), vec!["method"]), + (("/tokens","display token usage (input/output)"), vec![]), ]) .into_iter() .map(|((cmd,desc),args)|((cmd.to_string(),desc.to_string()),args.into_iter().map(|s|s.to_string()).collect())) @@ -55,6 +56,15 @@ impl App<'_> { } } } + "/tokens" => { + let msg = format!( + "Token Usage - Input: {}, Output: {}, Total: {}", + self.total_input_tokens, + self.total_output_tokens, + self.total_input_tokens + self.total_output_tokens + ); + self.input.alert_msg(&msg, Duration::from_secs(5)); + } _ => { self.input.alert_msg("command unknown", Duration::from_secs(1)); } diff --git a/shai-cli/src/tui/helper.rs b/shai-cli/src/tui/helper.rs index e5ae755..79a543a 100644 --- a/shai-cli/src/tui/helper.rs +++ b/shai-cli/src/tui/helper.rs @@ -14,14 +14,15 @@ impl HelpArea { "", " Available Commands:", " /exit exit from the tui", - " /tc set tool call method: [auto | fc | fc2 | so]" + " /tc set tool call method: [auto | fc | fc2 | so]", + " /tokens display token usage" ].join("\n").to_string() } } impl HelpArea { pub fn height(&self) -> u16 { - 7 // content (3 general help lines + 1 blank + 1 header + 2 command lines) + 8 // content (3 general help lines + 1 blank + 1 header + 3 command lines) } pub fn draw(&self, f: &mut Frame, area: Rect) { diff --git a/shai-core/src/agent/actions/brain.rs b/shai-core/src/agent/actions/brain.rs index a9832c6..4b24d2b 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -47,7 +47,7 @@ impl AgentCore { /// Process a brain task result pub async fn process_next_step(&mut self, result: Result) -> Result<(), AgentError> { - let ThinkerDecision{message, flow} = self.handle_brain_error(result).await?; + let ThinkerDecision{message, flow, token_usage} = self.handle_brain_error(result).await?; let ChatMessage::Assistant { content, reasoning_content, tool_calls, .. } = message.clone() else { return self.handle_brain_error::( Err(AgentError::InvalidResponse(format!("ChatMessage::Assistant expected, but got {:?} instead", message)))).await.map(|_| () @@ -60,10 +60,18 @@ impl AgentCore { trace.write().await.push(message.clone()); // Emit event to external consumers - let _ = self.emit_event(AgentEvent::BrainResult { + let _ = self.emit_event(AgentEvent::BrainResult { timestamp: Utc::now(), thought: Ok(message.clone()) }).await; + + // Emit token usage event if available + if let Some((input_tokens, output_tokens)) = token_usage { + let _ = self.emit_event(AgentEvent::TokenUsage { + input_tokens, + output_tokens + }).await; + } // run tool call if any let tool_calls_from_brain = tool_calls.unwrap_or(vec![]); diff --git a/shai-core/src/agent/brain.rs b/shai-core/src/agent/brain.rs index fb6040e..d77a1af 100644 --- a/shai-core/src/agent/brain.rs +++ b/shai-core/src/agent/brain.rs @@ -27,28 +27,48 @@ pub enum ThinkerFlowControl { #[derive(Debug, Clone)] pub struct ThinkerDecision { pub message: ChatMessage, - pub flow: ThinkerFlowControl + pub flow: ThinkerFlowControl, + pub token_usage: Option<(u32, u32)>, // (input_tokens, output_tokens) } impl ThinkerDecision { pub fn new(message: ChatMessage) -> Self { ThinkerDecision{ message, - flow: ThinkerFlowControl::AgentPause + flow: ThinkerFlowControl::AgentPause, + token_usage: None, } } pub fn agent_continue(message: ChatMessage) -> Self { ThinkerDecision{ message, - flow: ThinkerFlowControl::AgentContinue + flow: ThinkerFlowControl::AgentContinue, + token_usage: None, } } pub fn agent_pause(message: ChatMessage) -> Self { ThinkerDecision{ message, - flow: ThinkerFlowControl::AgentPause + flow: ThinkerFlowControl::AgentPause, + token_usage: None, + } + } + + pub fn agent_continue_with_tokens(message: ChatMessage, input_tokens: u32, output_tokens: u32) -> Self { + ThinkerDecision{ + message, + flow: ThinkerFlowControl::AgentContinue, + token_usage: Some((input_tokens, output_tokens)), + } + } + + pub fn agent_pause_with_tokens(message: ChatMessage, input_tokens: u32, output_tokens: u32) -> Self { + ThinkerDecision{ + message, + flow: ThinkerFlowControl::AgentPause, + token_usage: Some((input_tokens, output_tokens)), } } diff --git a/shai-core/src/agent/events.rs b/shai-core/src/agent/events.rs index d4febc6..e0ed5e6 100644 --- a/shai-core/src/agent/events.rs +++ b/shai-core/src/agent/events.rs @@ -96,6 +96,11 @@ pub enum AgentEvent { Error { error: String }, /// Agent execution completed Completed { success: bool, message: String }, + /// Token usage information from LLM response + TokenUsage { + input_tokens: u32, + output_tokens: u32 + }, } /// Types of user input that an agent can request @@ -263,6 +268,12 @@ impl std::fmt::Debug for AgentEvent { .field("message", message) .finish() } + AgentEvent::TokenUsage { input_tokens, output_tokens } => { + f.debug_struct("TokenUsage") + .field("input_tokens", input_tokens) + .field("output_tokens", output_tokens) + .finish() + } } } } diff --git a/shai-core/src/agent/output/log.rs b/shai-core/src/agent/output/log.rs index e80b651..1c4ec18 100644 --- a/shai-core/src/agent/output/log.rs +++ b/shai-core/src/agent/output/log.rs @@ -53,6 +53,9 @@ impl FileEventLogger { AgentEvent::Completed { success, message } => { format!("Completed: success={} - {}", success, message) } + AgentEvent::TokenUsage { input_tokens, output_tokens } => { + format!("Token Usage: input={} output={} total={}", input_tokens, output_tokens, input_tokens + output_tokens) + } }; let log_line = format!("[{}] {}\n", timestamp.format("%Y-%m-%d %H:%M:%S%.3f"), event_str); diff --git a/shai-core/src/agent/output/pretty.rs b/shai-core/src/agent/output/pretty.rs index ccd5470..d06b705 100644 --- a/shai-core/src/agent/output/pretty.rs +++ b/shai-core/src/agent/output/pretty.rs @@ -110,6 +110,10 @@ impl PrettyFormatter { Some(completion_skin.term_text(&markdown).to_string()) }, + AgentEvent::TokenUsage { .. } => { + // Don't display token usage in the main output - it's handled by /tokens command + None + }, }.map(|s| format!("\n{}", s)) } diff --git a/shai-core/src/runners/coder/coder.rs b/shai-core/src/runners/coder/coder.rs index 287ce5a..04931d5 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -77,15 +77,30 @@ impl Brain for CoderBrain { context.method) .await .map_err(|e| AgentError::LlmError(e.to_string()))?; - + + // Extract token usage information + let (input_tokens, output_tokens) = if let Some(usage) = &brain_decision.usage { + let input = usage.prompt_tokens.unwrap_or(0); + let output = usage.completion_tokens.unwrap_or(0); + debug!(target: "brain::coder::tokens", + input_tokens = input, + output_tokens = output, + total_tokens = usage.total_tokens, + "Token usage for LLM call" + ); + (input, output) + } else { + (0, 0) + }; + // stop here if there's no other tool calls let message = brain_decision.choices.into_iter().next().unwrap().message; if let ChatMessage::Assistant { reasoning_content, content, tool_calls, .. } = &message { if tool_calls.as_ref().map_or(true, |calls| calls.is_empty()) { - return Ok(ThinkerDecision::agent_pause(message)); + return Ok(ThinkerDecision::agent_pause_with_tokens(message, input_tokens, output_tokens)); } - } - Ok(ThinkerDecision::agent_continue(message)) + } + Ok(ThinkerDecision::agent_continue_with_tokens(message, input_tokens, output_tokens)) } } diff --git a/shai-llm/src/providers/openai_compatible.rs b/shai-llm/src/providers/openai_compatible.rs index 76b2fc0..5085a1d 100644 --- a/shai-llm/src/providers/openai_compatible.rs +++ b/shai-llm/src/providers/openai_compatible.rs @@ -7,8 +7,10 @@ use openai_dive::v1::{ resources::{ chat::{ChatCompletionParameters, ChatCompletionResponse, ChatCompletionChunkResponse}, model::ListModelResponse, + shared::Usage, }, }; +use serde_json::Value; pub struct OpenAICompatibleProvider { client: Client, @@ -31,6 +33,33 @@ impl OpenAICompatibleProvider { _ => None } } + + fn process_usage_information(&self, mut response: ChatCompletionResponse) -> ChatCompletionResponse { + // Convert response to JSON to extract usage information + if let Ok(response_json) = serde_json::to_value(&response) { + if let Some(usage_obj) = response_json.get("usage") { + let input_tokens = usage_obj.get("prompt_tokens") + .or_else(|| usage_obj.get("input_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + let output_tokens = usage_obj.get("completion_tokens") + .or_else(|| usage_obj.get("output_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + // Update usage with properly extracted token counts + response.usage = Some(Usage { + prompt_tokens: Some(input_tokens), + completion_tokens: Some(output_tokens), + total_tokens: input_tokens + output_tokens, + prompt_tokens_details: None, + completion_tokens_details: None, + }); + } + } + response + } } #[async_trait] @@ -42,8 +71,10 @@ impl LlmProvider for OpenAICompatibleProvider { } async fn chat(&self, request: ChatCompletionParameters) -> Result { - let response = self.client.chat().create(request).await + let mut response = self.client.chat().create(request).await .map_err(|e| Box::new(e) as LlmError)?; + + response = self.process_usage_information(response); Ok(response) } diff --git a/shai-llm/src/providers/ovhcloud.rs b/shai-llm/src/providers/ovhcloud.rs index 71980e7..5909875 100644 --- a/shai-llm/src/providers/ovhcloud.rs +++ b/shai-llm/src/providers/ovhcloud.rs @@ -7,9 +7,11 @@ use openai_dive::v1::{ resources::{ chat::{ChatCompletionParameters, ChatCompletionResponse, ChatCompletionChunkResponse}, model::ListModelResponse, + shared::Usage, }, error::APIError }; +use serde_json::Value; const OVH_API_BASE: &str = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1"; @@ -34,15 +36,40 @@ impl OvhCloudProvider { }) } - fn sanitize_request(&self, mut request: ChatCompletionParameters) -> ChatCompletionParameters { + fn sanitize_request(&self, mut request: ChatCompletionParameters) -> ChatCompletionParameters { // OVH uses max_tokens instead of max_completion_tokens if request.max_completion_tokens.is_some() { request.max_tokens = request.max_completion_tokens; request.max_completion_tokens = None; } - + request } + + fn process_usage_information(&self, mut response: ChatCompletionResponse) -> ChatCompletionResponse { + if let Ok(response_json) = serde_json::to_value(&response) { + if let Some(usage_obj) = response_json.get("usage") { + let input_tokens = usage_obj.get("prompt_tokens") + .or_else(|| usage_obj.get("input_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + let output_tokens = usage_obj.get("completion_tokens") + .or_else(|| usage_obj.get("output_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + response.usage = Some(Usage { + prompt_tokens: Some(input_tokens), + completion_tokens: Some(output_tokens), + total_tokens: input_tokens + output_tokens, + prompt_tokens_details: None, + completion_tokens_details: None, + }); + } + } + response + } } #[async_trait] @@ -65,8 +92,10 @@ impl LlmProvider for OvhCloudProvider { async fn chat(&self, request: ChatCompletionParameters) -> Result { let sanitized_request = self.sanitize_request(request); - let response = self.client.chat().create(sanitized_request).await + let mut response = self.client.chat().create(sanitized_request).await .map_err(|e| Box::new(e) as LlmError)?; + + response = self.process_usage_information(response); Ok(response) } From cd24d70c0656bb1b5080976bc899035ec225babc Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Sun, 28 Sep 2025 00:58:52 +0200 Subject: [PATCH 02/12] =?UTF-8?q?feat:=20token=20compression=20at=2080?= =?UTF-8?q?=E2=80=AF%=20of=20max=5Fcontext=5Ftokens?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Pierric Buchez --- README.md | 44 +++ shai-core/src/agent/actions/brain.rs | 71 +++- shai-core/src/agent/brain.rs | 27 +- shai-core/src/agent/events.rs | 19 ++ shai-core/src/agent/output/log.rs | 16 + shai-core/src/agent/output/pretty.rs | 41 +++ shai-core/src/config/config.rs | 25 +- shai-core/src/runners/coder/coder.rs | 61 +++- shai-core/src/runners/compacter/compact.rs | 375 +++++++++++++++++++++ shai-core/src/runners/compacter/mod.rs | 5 +- shai-core/src/runners/compacter/prompt.rs | 20 ++ 11 files changed, 691 insertions(+), 13 deletions(-) create mode 100644 shai-core/src/runners/compacter/prompt.rs diff --git a/README.md b/README.md index d767ceb..aa10e20 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,50 @@ the `shai` binary will be installed in `$HOME/.local/bin` ## Configure a provider and Run! +### Configuration files + +Shai can be configured via **configuration files** written in JSON. By default, the configuration file is `auth.config` located in `~/.config/shai/`. The file defines the list of LLM providers, the selected provider, model, and tool call method. + +#### Example `.shai.config` +```json +{ + "providers": [ + { + "provider": "ovhcloud", + "env_vars": { + "OVH_BASE_URL": "https://gpt-oss-120b.endpoints.kepler.ai.cloud.ovh.net/api/openai_compat/v1" + }, + "model": "gpt-oss-120b", + "tool_method": "FunctionCall", + "max_context_tokens": 8192 + } + ], + "selected_provider": 0 +} +``` + +- **providers**: an array of provider definitions. Each provider can specify environment variables (`env_vars`), the model name, the tool call method (`FunctionCall` or `Chat`), and optionally `max_context_tokens` to limit the context size. +- **selected_provider**: the index of the provider to use (starting at `0`). +- **max_context_tokens** (optional, per provider): maximum number of tokens that can be sent in the context to the LLM. If omitted, the default for the model is used. + +You can create multiple configuration files for different agents (see the *Custom Agent* section). To use a specific configuration, place the file in `~/.config/shai/agents/` and run the agent by its filename (without the `.config` extension): +``` +shai my_custom_agent +``` + +Shai will automatically load the configuration, set the required environment variables, and use the selected provider for all subsequent interactions. + +### Using the configuration + +- **Automatic loading**: If a `.shai.config` file is present in the current directory, Shai will load it automatically. +- **Explicit loading**: Use the `--config ` flag to specify a custom configuration file: +``` +shai --config ~/.config/shai/agents/example.config +``` + +The configuration system allows you to switch providers, models, or tool call methods without recompiling the binary. + + By default `shai` uses OVHcloud as an anonymous user meaning you will be rate limited! If you want to sign in with your account or select another provider, run: ``` diff --git a/shai-core/src/agent/actions/brain.rs b/shai-core/src/agent/actions/brain.rs index 4b24d2b..1ba7fb0 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -47,7 +47,7 @@ impl AgentCore { /// Process a brain task result pub async fn process_next_step(&mut self, result: Result) -> Result<(), AgentError> { - let ThinkerDecision{message, flow, token_usage} = self.handle_brain_error(result).await?; + let ThinkerDecision{message, flow, token_usage, compression_info} = self.handle_brain_error(result).await?; let ChatMessage::Assistant { content, reasoning_content, tool_calls, .. } = message.clone() else { return self.handle_brain_error::( Err(AgentError::InvalidResponse(format!("ChatMessage::Assistant expected, but got {:?} instead", message)))).await.map(|_| () @@ -72,6 +72,18 @@ impl AgentCore { output_tokens }).await; } + + // Emit context compression event if available + if let Some(compression_info) = compression_info { + let _ = self.emit_event(AgentEvent::ContextCompressed { + original_message_count: compression_info.original_message_count, + compressed_message_count: compression_info.compressed_message_count, + tokens_before: compression_info.tokens_before, + current_tokens: compression_info.current_tokens, + max_tokens: compression_info.max_tokens, + ai_summary: compression_info.ai_summary, + }).await; + } // run tool call if any let tool_calls_from_brain = tool_calls.unwrap_or(vec![]); @@ -86,19 +98,74 @@ impl AgentCore { self.set_state(InternalAgentState::Running).await; } ThinkerFlowControl::AgentPause => { + // Check if we need to compress context when task is complete + self.check_and_compress_context().await?; self.set_state(InternalAgentState::Paused).await; } } Ok(()) } + /// Check if context compression is needed and apply it when task is complete + async fn check_and_compress_context(&mut self) -> Result<(), AgentError> { + // Extract compression logic from the brain if it's a CoderBrain + let brain = self.brain.clone(); + let brain_read = brain.read().await; + + // This is a bit hacky but we need to check if the brain has a compressor + // We'll use Any trait to downcast to CoderBrain + use std::any::Any; + + if let Some(coder_brain) = (&**brain_read as &dyn Any).downcast_ref::() { + if let Some(compressor) = &coder_brain.context_compressor { + let compressor_clone = compressor.clone(); + drop(brain_read); // Release the read lock + + let trace = self.trace.read().await.clone(); + let mut compressor_clone = compressor_clone; + + if compressor_clone.should_compress_conversation(&trace) { + let (compressed_trace, compression_info) = compressor_clone.compress_messages(trace).await; + + // Update the trace with compressed version + { + let mut trace_write = self.trace.write().await; + *trace_write = compressed_trace; + } + + // Update the compressor in the brain + { + let mut brain_write = brain.write().await; + if let Some(coder_brain_mut) = (&mut **brain_write as &mut dyn Any).downcast_mut::() { + coder_brain_mut.context_compressor = Some(compressor_clone); + } + } + + // Emit compression event if compression occurred + if let Some(compression_info) = compression_info { + let _ = self.emit_event(AgentEvent::ContextCompressed { + original_message_count: compression_info.original_message_count, + compressed_message_count: compression_info.compressed_message_count, + tokens_before: compression_info.tokens_before, + current_tokens: compression_info.current_tokens, + max_tokens: compression_info.max_tokens, + ai_summary: compression_info.ai_summary, + }).await; + } + } + } + } + + Ok(()) + } + // Helper method that emits error events before returning the error async fn handle_brain_error(&mut self, result: Result) -> Result { match result { Ok(value) => Ok(value), Err(error) => { self.set_state(InternalAgentState::Paused).await; - let _ = self.emit_event(AgentEvent::BrainResult { + let _ = self.emit_event(AgentEvent::BrainResult { timestamp: Utc::now(), thought: Err(error.clone()) }).await; diff --git a/shai-core/src/agent/brain.rs b/shai-core/src/agent/brain.rs index d77a1af..7d9a6b2 100644 --- a/shai-core/src/agent/brain.rs +++ b/shai-core/src/agent/brain.rs @@ -4,6 +4,7 @@ use shai_llm::{ChatMessage, ToolCallMethod}; use tokio::sync::RwLock; use crate::tools::types::AnyToolBox; +use crate::runners::compacter::CompressionInfo; use super::error::AgentError; @@ -29,6 +30,7 @@ pub struct ThinkerDecision { pub message: ChatMessage, pub flow: ThinkerFlowControl, pub token_usage: Option<(u32, u32)>, // (input_tokens, output_tokens) + pub compression_info: Option, } impl ThinkerDecision { @@ -37,6 +39,7 @@ impl ThinkerDecision { message, flow: ThinkerFlowControl::AgentPause, token_usage: None, + compression_info: None, } } @@ -45,6 +48,7 @@ impl ThinkerDecision { message, flow: ThinkerFlowControl::AgentContinue, token_usage: None, + compression_info: None, } } @@ -53,6 +57,7 @@ impl ThinkerDecision { message, flow: ThinkerFlowControl::AgentPause, token_usage: None, + compression_info: None, } } @@ -61,6 +66,7 @@ impl ThinkerDecision { message, flow: ThinkerFlowControl::AgentContinue, token_usage: Some((input_tokens, output_tokens)), + compression_info: None, } } @@ -69,6 +75,25 @@ impl ThinkerDecision { message, flow: ThinkerFlowControl::AgentPause, token_usage: Some((input_tokens, output_tokens)), + compression_info: None, + } + } + + pub fn agent_continue_with_compression(message: ChatMessage, input_tokens: u32, output_tokens: u32, compression_info: CompressionInfo) -> Self { + ThinkerDecision{ + message, + flow: ThinkerFlowControl::AgentContinue, + token_usage: Some((input_tokens, output_tokens)), + compression_info: Some(compression_info), + } + } + + pub fn agent_pause_with_compression(message: ChatMessage, input_tokens: u32, output_tokens: u32, compression_info: CompressionInfo) -> Self { + ThinkerDecision{ + message, + flow: ThinkerFlowControl::AgentPause, + token_usage: Some((input_tokens, output_tokens)), + compression_info: Some(compression_info), } } @@ -79,7 +104,7 @@ impl ThinkerDecision { /// Core thinking interface - pure decision making #[async_trait] -pub trait Brain: Send + Sync { +pub trait Brain: Send + Sync + std::any::Any { /// This method is called at every step of the agent to decide next step /// note that if the message contains toolcall, it will always continue async fn next_step(&mut self, context: ThinkerContext) -> Result; diff --git a/shai-core/src/agent/events.rs b/shai-core/src/agent/events.rs index e0ed5e6..35582a4 100644 --- a/shai-core/src/agent/events.rs +++ b/shai-core/src/agent/events.rs @@ -101,6 +101,15 @@ pub enum AgentEvent { input_tokens: u32, output_tokens: u32 }, + /// Context compression notification + ContextCompressed { + original_message_count: usize, + compressed_message_count: usize, + tokens_before: Option, + current_tokens: Option, + max_tokens: u32, + ai_summary: Option, + }, } /// Types of user input that an agent can request @@ -274,6 +283,16 @@ impl std::fmt::Debug for AgentEvent { .field("output_tokens", output_tokens) .finish() } + AgentEvent::ContextCompressed { original_message_count, compressed_message_count, tokens_before, current_tokens, max_tokens, ai_summary } => { + f.debug_struct("ContextCompressed") + .field("original_message_count", original_message_count) + .field("compressed_message_count", compressed_message_count) + .field("tokens_before", tokens_before) + .field("current_tokens", current_tokens) + .field("max_tokens", max_tokens) + .field("ai_summary", ai_summary) + .finish() + } } } } diff --git a/shai-core/src/agent/output/log.rs b/shai-core/src/agent/output/log.rs index 1c4ec18..447dba7 100644 --- a/shai-core/src/agent/output/log.rs +++ b/shai-core/src/agent/output/log.rs @@ -56,6 +56,22 @@ impl FileEventLogger { AgentEvent::TokenUsage { input_tokens, output_tokens } => { format!("Token Usage: input={} output={} total={}", input_tokens, output_tokens, input_tokens + output_tokens) } + AgentEvent::ContextCompressed { original_message_count, compressed_message_count, tokens_before, current_tokens, max_tokens, ai_summary } => { + let summary_text = if let Some(summary) = ai_summary { + format!(" | Summary: {}", summary) + } else { + "".to_string() + }; + + let token_info = match (tokens_before, current_tokens) { + (Some(before), Some(after)) => format!(", tokens: {} → {}", before, after), + (Some(before), None) => format!(", tokens before: {}", before), + _ => "".to_string(), + }; + + format!("Context Compressed with AI Summary: {} → {} messages{}{}", + original_message_count, compressed_message_count, token_info, summary_text) + } }; let log_line = format!("[{}] {}\n", timestamp.format("%Y-%m-%d %H:%M:%S%.3f"), event_str); diff --git a/shai-core/src/agent/output/pretty.rs b/shai-core/src/agent/output/pretty.rs index d06b705..b370e3e 100644 --- a/shai-core/src/agent/output/pretty.rs +++ b/shai-core/src/agent/output/pretty.rs @@ -114,6 +114,47 @@ impl PrettyFormatter { // Don't display token usage in the main output - it's handled by /tokens command None }, + AgentEvent::ContextCompressed { original_message_count, compressed_message_count, tokens_before, current_tokens, max_tokens, ai_summary } => { + let net_change = if original_message_count > compressed_message_count { + original_message_count - compressed_message_count + } else { + 0 + }; + + let markdown = match (tokens_before, current_tokens) { + (Some(before), Some(after)) => { + if net_change > 0 { + format!( + "● **Context Compressed with AI Summary** - Summarized {} message(s) to stay within token limits ({} → {} tokens)", + net_change, before, after + ) + } else { + format!( + "● **Context Compression Applied** - Added AI summary to optimize token usage ({} → {} tokens)", + before, after + ) + } + } + _ => { + if net_change > 0 { + format!( + "● **Context Compressed with AI Summary** - Summarized {} message(s) to stay within token limits", + net_change + ) + } else { + format!( + "● **Context Compression Applied** - Added AI summary to optimize token usage" + ) + } + } + }; + + let mut compression_skin = self.skin.clone(); + compression_skin.paragraph.set_fg(rgb(100, 200, 255)); // Blue for AI compression + compression_skin.bold.set_fg(rgb(120, 220, 255)); // Light blue for bold + + Some(compression_skin.term_text(&markdown).to_string()) + }, }.map(|s| format!("\n{}", s)) } diff --git a/shai-core/src/config/config.rs b/shai-core/src/config/config.rs index 63b0caf..88a99be 100644 --- a/shai-core/src/config/config.rs +++ b/shai-core/src/config/config.rs @@ -12,7 +12,9 @@ pub struct ProviderConfig { pub provider: String, pub env_vars: std::collections::HashMap, pub model: String, - pub tool_method: ToolCallMethod + pub tool_method: ToolCallMethod, + #[serde(default)] + pub max_context_tokens: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -36,9 +38,23 @@ impl ShaiConfig { provider, env_vars, model, - tool_method: ToolCallMethod::FunctionCall + tool_method: ToolCallMethod::FunctionCall, + max_context_tokens: None, }; - + + self.providers.push(provider_config); + self.providers.len() - 1 + } + + pub fn add_provider_with_context(&mut self, provider: String, env_vars: std::collections::HashMap, model: String, max_context_tokens: Option) -> usize { + let provider_config = ProviderConfig { + provider, + env_vars, + model, + tool_method: ToolCallMethod::FunctionCall, + max_context_tokens, + }; + self.providers.push(provider_config); self.providers.len() - 1 } @@ -228,7 +244,8 @@ impl Default for ShaiConfig { (String::from("OVH_BASE_URL"), String::from("https://qwen-3-32b.endpoints.kepler.ai.cloud.ovh.net/api/openai_compat/v1")) ]), model: "Qwen3-32B".to_string(), - tool_method: ToolCallMethod::FunctionCall + tool_method: ToolCallMethod::FunctionCall, + max_context_tokens: Some(32768), // Qwen3-32B has 32k context }], selected_provider: 0, mcp_configs: HashMap::new(), diff --git a/shai-core/src/runners/coder/coder.rs b/shai-core/src/runners/coder/coder.rs index 04931d5..655bec2 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -10,6 +10,8 @@ use crate::agent::{Agent, AgentBuilder, AgentError, Brain, ThinkerContext}; use crate::tools::types::{ContainsAnyTool, IntoToolBox}; use shai_llm::tool::LlmToolCall; use crate::tools::{AnyTool, BashTool, EditTool, FetchTool, FindTool, LsTool, MultiEditTool, ReadTool, TodoReadTool, TodoWriteTool, WriteTool, TodoStorage, FsOperationLog}; +use crate::runners::compacter::ContextCompressor; +use crate::config::config::ShaiConfig; use super::prompt::{render_system_prompt_template, get_todo_read}; @@ -19,26 +21,59 @@ pub struct CoderBrain { pub model: String, pub system_prompt_template: String, pub temperature: f32, + pub context_compressor: Option, } impl CoderBrain { pub fn new(llm: Arc, model: String) -> Self { debug!(target: "brain::coder", provider =?llm.provider_name(), model = ?model); - Self { - llm, + + // Try to get context limit from configuration + let context_compressor = if let Ok(config) = ShaiConfig::load() { + if let Some(provider_config) = config.get_selected_provider() { + provider_config.max_context_tokens.map(|max_tokens| { + debug!(target: "brain::coder", max_context_tokens = max_tokens, "Initializing context compressor with AI summarization"); + ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone()) + }) + } else { + None + } + } else { + None + }; + + Self { + llm, model, system_prompt_template: "{{CODER_BASE_PROMPT}}".to_string(), temperature: 0.3, + context_compressor, } } pub fn with_custom_prompt(llm: Arc, model: String, system_prompt_template: String, temperature: f32) -> Self { debug!(target: "brain::coder", provider =?llm.provider_name(), model = ?model); - Self { - llm, + + // Try to get context limit from configuration + let context_compressor = if let Ok(config) = ShaiConfig::load() { + if let Some(provider_config) = config.get_selected_provider() { + provider_config.max_context_tokens.map(|max_tokens| { + debug!(target: "brain::coder", max_context_tokens = max_tokens, "Initializing context compressor with AI summarization"); + ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone()) + }) + } else { + None + } + } else { + None + }; + + Self { + llm, model, system_prompt_template, temperature, + context_compressor, } } } @@ -51,7 +86,7 @@ impl Brain for CoderBrain { // Render the user's system prompt template let mut system_prompt = render_system_prompt_template(&self.system_prompt_template); - + // Add todo status if available if let Some(tool) = context.available_tools.get_tool("todo_read") { let todo_status = get_todo_read(&tool).await; @@ -63,6 +98,8 @@ impl Brain for CoderBrain { name: None, }); + // Note: Context compression is now handled when tasks complete, not during thinking + // get next step with custom temperature let request = ChatCompletionParametersBuilder::default() .model(&self.model) @@ -88,6 +125,19 @@ impl Brain for CoderBrain { total_tokens = usage.total_tokens, "Token usage for LLM call" ); + + // Update context compressor with token usage + if let Some(compressor) = &mut self.context_compressor { + compressor.update_token_count(input, output); + if compressor.is_near_limit() { + debug!(target: "brain::coder::context", + current_tokens = compressor.get_current_tokens(), + max_tokens = compressor.get_max_tokens(), + "Approaching context limit" + ); + } + } + (input, output) } else { (0, 0) @@ -100,6 +150,7 @@ impl Brain for CoderBrain { return Ok(ThinkerDecision::agent_pause_with_tokens(message, input_tokens, output_tokens)); } } + Ok(ThinkerDecision::agent_continue_with_tokens(message, input_tokens, output_tokens)) } } diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index e69de29..4ec6795 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -0,0 +1,375 @@ +use shai_llm::{ChatMessage, ChatMessageContent, client::LlmClient}; +use tracing::{debug, info, warn}; +use std::sync::Arc; +use openai_dive::v1::resources::chat::ChatCompletionParametersBuilder; + +use super::prompt::get_compression_summary_prompt; + +/// Information about a compression operation +#[derive(Debug, Clone)] +pub struct CompressionInfo { + pub original_message_count: usize, + pub compressed_message_count: usize, + pub tokens_before: Option, + pub current_tokens: Option, + pub max_tokens: u32, + pub ai_summary: Option, +} + +/// Context compression utilities for managing conversation history within token limits +#[derive(Clone)] +pub struct ContextCompressor { + max_tokens: u32, + current_tokens: u32, + llm_client: Option>, + model: Option, +} + +impl ContextCompressor { + pub fn new(max_tokens: u32) -> Self { + Self { + max_tokens, + current_tokens: 0, + llm_client: None, + model: None, + } + } + + pub fn new_with_llm(max_tokens: u32, llm_client: Arc, model: String) -> Self { + Self { + max_tokens, + current_tokens: 0, + llm_client: Some(llm_client), + model: Some(model), + } + } + + /// Update the current token count + pub fn update_token_count(&mut self, input_tokens: u32, output_tokens: u32) { + self.current_tokens += input_tokens + output_tokens; + debug!( + target: "context_compression", + current_tokens = self.current_tokens, + max_tokens = self.max_tokens, + "Updated token count" + ); + } + + /// Check if we're approaching the context limit with dynamic threshold + /// Uses different thresholds based on context size to avoid premature compression + pub fn should_compress(&self) -> bool { + let threshold_percentage = 0.80; + let threshold = (self.max_tokens as f64 * threshold_percentage) as u32; + self.current_tokens >= threshold + } + + /// Check if compression is actually beneficial given the current conversation + pub fn should_compress_conversation(&self, messages: &[ChatMessage]) -> bool { + // Don't compress if we're not near the limit + if !self.should_compress() { + return false; + } + + // Count non-system messages + let non_system_count = messages.iter() + .filter(|msg| !matches!(msg, ChatMessage::System { .. })) + .count(); + + // Only compress if we have at least 2 messages (1 user-assistant pair) + // This ensures we have something meaningful to summarize + let should_compress = non_system_count > 2; + + let threshold_percentage = 80; + + debug!( + target: "context_compression", + non_system_count = non_system_count, + current_tokens = self.current_tokens, + max_tokens = self.max_tokens, + should_compress = should_compress, + threshold_percentage = threshold_percentage, + "Compression decision" + ); + + should_compress + } + + /// Compress the conversation history by removing older messages and replacing with AI summary + /// Keeps the system message and recent messages while summarizing middle conversation + /// Returns (compressed_messages, compression_info) + pub async fn compress_messages(&mut self, messages: Vec) -> (Vec, Option) { + if !self.should_compress_conversation(&messages) { + return (messages, None); + } + + let original_count = messages.len(); + let tokens_before_compression = self.current_tokens; + + info!( + target: "context_compression", + total_tokens_before = self.current_tokens, + max_tokens = self.max_tokens, + original_message_count = messages.len(), + "Compressing context due to token limit" + ); + + let mut compressed = Vec::new(); + let mut system_messages = Vec::new(); + let mut middle_messages = Vec::new(); + let mut recent_messages = Vec::new(); + + // Separate messages into categories + for (index, message) in messages.iter().enumerate() { + match message { + ChatMessage::System { .. } => { + system_messages.push(message.clone()); + } + _ => { + // Keep only the last 2 messages (1 user-assistant pair) as recent + // Everything else goes to middle_messages for summarization + if index >= messages.len().saturating_sub(2) { + recent_messages.push(message.clone()); + } else { + middle_messages.push(message.clone()); + } + } + } + } + + // Add system messages first + compressed.extend(system_messages); + + // Try to generate AI summary of middle conversation + let (ai_summary, summary_tokens) = if !middle_messages.is_empty() { + match self.summarize_conversation(&middle_messages).await { + Ok((summary, tokens)) => { + info!(target: "context_compression", "Successfully generated AI summary"); + compressed.push(ChatMessage::System { + content: ChatMessageContent::Text(format!( + "Previous conversation summary: {}", + summary + )), + name: Some("summary".to_string()), + }); + (Some(summary), tokens) + } + Err(e) => { + warn!(target: "context_compression", error = e, "Failed to generate AI summary, using fallback"); + compressed.push(ChatMessage::System { + content: ChatMessageContent::Text( + "[Previous conversation history compressed - AI summary unavailable]".to_string() + ), + name: Some("system".to_string()), + }); + (None, 50) // Estimate for fallback message + } + } + } else { + (None, 0) + }; + + // Add recent messages + compressed.extend(recent_messages); + + // Safely create compression info with validation + let compression_info = CompressionInfo { + original_message_count: original_count, + compressed_message_count: compressed.len(), + tokens_before: Some(tokens_before_compression), + // Only include token info if we have valid data (summary_tokens > 0) + current_tokens: if summary_tokens > 0 { Some(summary_tokens) } else { None }, + max_tokens: self.max_tokens, + ai_summary: ai_summary.clone(), + }; + + info!( + target: "context_compression", + compressed_message_count = compressed.len(), + estimated_tokens_after_compression = self.current_tokens, + output_tokens = self.current_tokens, + middle_messages_summarized = middle_messages.len(), + "Context compression with AI summary completed" + ); + + (compressed, Some(compression_info)) + } + + /// Get the current token count + pub fn get_current_tokens(&self) -> u32 { + self.current_tokens + } + + /// Get the maximum token limit + pub fn get_max_tokens(&self) -> u32 { + self.max_tokens + } + + /// Check if we're near the absolute limit (95% threshold) + pub fn is_near_limit(&self) -> bool { + let threshold = (self.max_tokens as f64 * 0.95) as u32; + self.current_tokens >= threshold + } + + /// Create a summary of the conversation history using AI + /// Returns (summary_text, summary_tokens_used) + async fn summarize_conversation(&mut self, messages: &[ChatMessage]) -> Result<(String, u32), String> { + let Some(ref llm_client) = self.llm_client else { + return Err("No LLM client available for summarization".to_string()); + }; + + let Some(ref model) = self.model else { + return Err("No model specified for summarization".to_string()); + }; + + // Create a conversation string from messages + let mut conversation_text = String::new(); + for message in messages { + match message { + ChatMessage::User { content, .. } => { + if let ChatMessageContent::Text(text) = content { + conversation_text.push_str(&format!("User: {}\n", text)); + } + } + ChatMessage::Assistant { content: Some(content), .. } => { + if let ChatMessageContent::Text(text) = content { + conversation_text.push_str(&format!("Assistant: {}\n", text)); + } + } + ChatMessage::System { content, .. } => { + if let ChatMessageContent::Text(text) = content { + // Skip system prompts in summary, only include actual system messages + if !text.contains("CODER_BASE_PROMPT") && !text.contains("You are") { + conversation_text.push_str(&format!("System: {}\n", text)); + } + } + } + _ => {} // Skip other message types + } + } + + let summary_prompt = get_compression_summary_prompt(&conversation_text); + + let summary_request = ChatCompletionParametersBuilder::default() + .model(model) + .messages(vec![ChatMessage::User { + content: ChatMessageContent::Text(summary_prompt), + name: None, + }]) + .temperature(0.1) // Low temperature for consistent summaries + .build() + .map_err(|e| format!("Failed to build summary request: {}", e))?; + + match llm_client.chat(summary_request).await { + Ok(response) => { + // Safely extract token usage from the summary generation + let summary_tokens = if let Some(usage) = &response.usage { + // Use completion_tokens if available, otherwise fall back to 0 + usage.completion_tokens.unwrap_or(0) + } else { + // No usage information available - return error instead of crashing + warn!(target: "context_compression", "No usage information available in LLM response"); + return Err("No token usage information available from LLM".to_string()); + }; + + // Safely extract the summary content + if let Some(choice) = response.choices.first() { + if let ChatMessage::Assistant { content: Some(content), .. } = &choice.message { + if let ChatMessageContent::Text(summary) = content { + // Only proceed if we have both valid summary and token count + if !summary.trim().is_empty() && summary_tokens > 0 { + debug!(target: "context_compression", + summary_length = summary.len(), + summary_tokens = summary_tokens, + "Successfully generated AI summary with token count" + ); + return Ok((summary.clone(), summary_tokens)); + } else { + warn!(target: "context_compression", "Received empty summary or zero tokens"); + return Err("Received empty summary or invalid token count".to_string()); + } + } + } + } + Err("No valid summary content found in LLM response".to_string()) + } + Err(e) => { + warn!(target: "context_compression", error = ?e, "Failed to generate summary, falling back to simple compression"); + Err(format!("LLM summarization failed: {}", e)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_compression_threshold() { + let compressor = ContextCompressor::new(1000); + assert!(!compressor.should_compress()); + + let mut compressor = ContextCompressor::new(1000); + compressor.update_token_count(800, 0); + assert!(compressor.should_compress()); + } + + #[test] + fn test_message_compression() { + let mut compressor = ContextCompressor::new(1000); + compressor.current_tokens = 850; // Above 80% threshold + + let messages = vec![ + ChatMessage::System { + content: ChatMessageContent::Text("System prompt".to_string()), + name: None, + }, + ChatMessage::User { + content: ChatMessageContent::Text("Old message 1".to_string()), + name: None, + }, + ChatMessage::Assistant { + content: Some(ChatMessageContent::Text("Old response 1".to_string())), + reasoning_content: None, + tool_calls: None, + refusal: None, + name: None, + audio: None, + }, + ChatMessage::User { + content: ChatMessageContent::Text("Recent message".to_string()), + name: None, + }, + ChatMessage::Assistant { + content: Some(ChatMessageContent::Text("Recent response".to_string())), + reasoning_content: None, + tool_calls: None, + refusal: None, + name: None, + audio: None, + }, + ]; + + let compressed = compressor.compress_messages(messages); + + // Should contain: system message, compression notice, recent messages + assert!(compressed.len() >= 4); + + // First message should be system + assert!(matches!(compressed[0], ChatMessage::System { .. })); + + // Should contain compression notice + let has_compression_notice = compressed.iter().any(|msg| { + if let ChatMessage::System { content, .. } = msg { + if let ChatMessageContent::Text(text) = content { + text.contains("compressed") + } else { + false + } + } else { + false + } + }); + assert!(has_compression_notice); + } +} \ No newline at end of file diff --git a/shai-core/src/runners/compacter/mod.rs b/shai-core/src/runners/compacter/mod.rs index 9e86497..ef6ceef 100644 --- a/shai-core/src/runners/compacter/mod.rs +++ b/shai-core/src/runners/compacter/mod.rs @@ -1 +1,4 @@ -pub mod compact; \ No newline at end of file +pub mod compact; +pub mod prompt; + +pub use compact::{ContextCompressor, CompressionInfo}; \ No newline at end of file diff --git a/shai-core/src/runners/compacter/prompt.rs b/shai-core/src/runners/compacter/prompt.rs new file mode 100644 index 0000000..fab0887 --- /dev/null +++ b/shai-core/src/runners/compacter/prompt.rs @@ -0,0 +1,20 @@ +static COMPRESSION_SUMMARY_PROMPT: &str = r#"Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. +This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing with the conversation and supporting any continuing tasks. + +Your summary should be structured as follows: +Context: The context to continue the conversation with. If applicable based on the current task, this should include: + 1. Previous Conversation: High level details about what was discussed throughout the entire conversation with the user. This should be written to allow someone to be able to follow the general overarching conversation flow. + 2. Current Work: Describe in detail what was being worked on prior to this request to summarize the conversation. Pay special attention to the more recent messages in the conversation. + 3. Key Technical Concepts: List all important technical concepts, technologies, coding conventions, and frameworks discussed, which might be relevant for continuing with this work. + 4. Relevant Files and Code: If applicable, enumerate specific files and code sections examined, modified, or created for the task continuation. Pay special attention to the most recent messages and changes. + 5. Problem Solving: Document problems solved thus far and any ongoing troubleshooting efforts. + 6. Pending Tasks and Next Steps: Outline all pending tasks that you have explicitly been asked to work on, as well as list the next steps you will take for all outstanding work, if applicable. Include code snippets where they add clarity. For any next steps, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no information loss in context between tasks. + +Output only the summary of the conversation so far, without any additional commentary or explanation. + +Conversation to summarize: +{}"#; + +pub fn get_compression_summary_prompt(conversation_text: &str) -> String { + COMPRESSION_SUMMARY_PROMPT.replace("{}", conversation_text) +} \ No newline at end of file From e99bc17109e699092fe955e53c2ab6d27a1e8897 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Sun, 28 Sep 2025 16:11:46 +0200 Subject: [PATCH 03/12] feat: update context compressor with token usage tracking --- shai-core/src/agent/events.rs | 3 --- shai-core/src/runners/coder/coder.rs | 12 ++++++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/shai-core/src/agent/events.rs b/shai-core/src/agent/events.rs index cc49729..35582a4 100644 --- a/shai-core/src/agent/events.rs +++ b/shai-core/src/agent/events.rs @@ -283,7 +283,6 @@ impl std::fmt::Debug for AgentEvent { .field("output_tokens", output_tokens) .finish() } -<<<<<<< HEAD AgentEvent::ContextCompressed { original_message_count, compressed_message_count, tokens_before, current_tokens, max_tokens, ai_summary } => { f.debug_struct("ContextCompressed") .field("original_message_count", original_message_count) @@ -294,8 +293,6 @@ impl std::fmt::Debug for AgentEvent { .field("ai_summary", ai_summary) .finish() } -======= ->>>>>>> dev/pbuchez/tokens-usage } } } diff --git a/shai-core/src/runners/coder/coder.rs b/shai-core/src/runners/coder/coder.rs index e607b14..f20d0cd 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -125,6 +125,18 @@ impl Brain for CoderBrain { total_tokens = usage.total_tokens, "Token usage for LLM call" ); + // Update context compressor with token usage + if let Some(compressor) = &mut self.context_compressor { + compressor.update_token_count(input, output); + if compressor.is_near_limit() { + debug!(target: "brain::coder::context", + current_tokens = compressor.get_current_tokens(), + max_tokens = compressor.get_max_tokens(), + "Approaching context limit" + ); + } + } + Some((input, output)) } else { None From 6006184b1c6f603a5a9d97cc76e7d4a834dfbe34 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 10:19:34 +0000 Subject: [PATCH 04/12] feat: Increase threshold compression, switch manual config to automatic detection, and refactor prompt --- shai-core/src/config/config.rs | 17 +---- shai-core/src/runners/coder/coder.rs | 16 +++-- shai-core/src/runners/compacter/compact.rs | 2 +- shai-core/src/runners/compacter/prompt.rs | 53 ++++++++++---- shai-llm/src/tool/max_context.rs | 83 ++++++++++++++++++++++ shai-llm/src/tool/mod.rs | 5 +- 6 files changed, 139 insertions(+), 37 deletions(-) create mode 100644 shai-llm/src/tool/max_context.rs diff --git a/shai-core/src/config/config.rs b/shai-core/src/config/config.rs index 88a99be..8689845 100644 --- a/shai-core/src/config/config.rs +++ b/shai-core/src/config/config.rs @@ -13,8 +13,7 @@ pub struct ProviderConfig { pub env_vars: std::collections::HashMap, pub model: String, pub tool_method: ToolCallMethod, - #[serde(default)] - pub max_context_tokens: Option, + } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -39,25 +38,13 @@ impl ShaiConfig { env_vars, model, tool_method: ToolCallMethod::FunctionCall, - max_context_tokens: None, }; self.providers.push(provider_config); self.providers.len() - 1 } - pub fn add_provider_with_context(&mut self, provider: String, env_vars: std::collections::HashMap, model: String, max_context_tokens: Option) -> usize { - let provider_config = ProviderConfig { - provider, - env_vars, - model, - tool_method: ToolCallMethod::FunctionCall, - max_context_tokens, - }; - self.providers.push(provider_config); - self.providers.len() - 1 - } pub fn is_duplicate_config(&self, provider_name: &str, env_vars: &std::collections::HashMap, model: &str) -> bool { self.providers.iter().any(|provider_config| { @@ -245,7 +232,7 @@ impl Default for ShaiConfig { ]), model: "Qwen3-32B".to_string(), tool_method: ToolCallMethod::FunctionCall, - max_context_tokens: Some(32768), // Qwen3-32B has 32k context + }], selected_provider: 0, mcp_configs: HashMap::new(), diff --git a/shai-core/src/runners/coder/coder.rs b/shai-core/src/runners/coder/coder.rs index f20d0cd..41d6b93 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -8,7 +8,7 @@ use tracing::debug; use crate::agent::brain::ThinkerDecision; use crate::agent::{Agent, AgentBuilder, AgentError, Brain, ThinkerContext}; use crate::tools::types::{ContainsAnyTool, IntoToolBox}; -use shai_llm::tool::LlmToolCall; +use shai_llm::tool::{LlmToolCall, get_max_context}; use crate::tools::{AnyTool, BashTool, EditTool, FetchTool, FindTool, LsTool, MultiEditTool, ReadTool, TodoReadTool, TodoWriteTool, WriteTool, TodoStorage, FsOperationLog}; use crate::runners::compacter::ContextCompressor; use crate::config::config::ShaiConfig; @@ -31,10 +31,11 @@ impl CoderBrain { // Try to get context limit from configuration let context_compressor = if let Ok(config) = ShaiConfig::load() { if let Some(provider_config) = config.get_selected_provider() { - provider_config.max_context_tokens.map(|max_tokens| { + { + let max_tokens = get_max_context(&provider_config.model) as u32; debug!(target: "brain::coder", max_context_tokens = max_tokens, "Initializing context compressor with AI summarization"); - ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone()) - }) + Some(ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone())) + } } else { None } @@ -57,10 +58,11 @@ impl CoderBrain { // Try to get context limit from configuration let context_compressor = if let Ok(config) = ShaiConfig::load() { if let Some(provider_config) = config.get_selected_provider() { - provider_config.max_context_tokens.map(|max_tokens| { + { + let max_tokens = get_max_context(&provider_config.model) as u32; debug!(target: "brain::coder", max_context_tokens = max_tokens, "Initializing context compressor with AI summarization"); - ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone()) - }) + Some(ContextCompressor::new_with_llm(max_tokens, llm.clone(), model.clone())) + } } else { None } diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index 4ec6795..4e44f5a 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -58,7 +58,7 @@ impl ContextCompressor { /// Check if we're approaching the context limit with dynamic threshold /// Uses different thresholds based on context size to avoid premature compression pub fn should_compress(&self) -> bool { - let threshold_percentage = 0.80; + let threshold_percentage = 0.90; let threshold = (self.max_tokens as f64 * threshold_percentage) as u32; self.current_tokens >= threshold } diff --git a/shai-core/src/runners/compacter/prompt.rs b/shai-core/src/runners/compacter/prompt.rs index fab0887..8961015 100644 --- a/shai-core/src/runners/compacter/prompt.rs +++ b/shai-core/src/runners/compacter/prompt.rs @@ -1,16 +1,43 @@ -static COMPRESSION_SUMMARY_PROMPT: &str = r#"Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions. -This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing with the conversation and supporting any continuing tasks. - -Your summary should be structured as follows: -Context: The context to continue the conversation with. If applicable based on the current task, this should include: - 1. Previous Conversation: High level details about what was discussed throughout the entire conversation with the user. This should be written to allow someone to be able to follow the general overarching conversation flow. - 2. Current Work: Describe in detail what was being worked on prior to this request to summarize the conversation. Pay special attention to the more recent messages in the conversation. - 3. Key Technical Concepts: List all important technical concepts, technologies, coding conventions, and frameworks discussed, which might be relevant for continuing with this work. - 4. Relevant Files and Code: If applicable, enumerate specific files and code sections examined, modified, or created for the task continuation. Pay special attention to the most recent messages and changes. - 5. Problem Solving: Document problems solved thus far and any ongoing troubleshooting efforts. - 6. Pending Tasks and Next Steps: Outline all pending tasks that you have explicitly been asked to work on, as well as list the next steps you will take for all outstanding work, if applicable. Include code snippets where they add clarity. For any next steps, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no information loss in context between tasks. - -Output only the summary of the conversation so far, without any additional commentary or explanation. +static COMPRESSION_SUMMARY_PROMPT: &str = r#"Generate a comprehensive summary of this conversation, focusing on capturing every detail necessary to seamlessly continue the work. + +Your summary must preserve: +- All user requests and instructions (both completed and pending) +- Technical decisions, code patterns, and architectural choices +- The progression of work and problem-solving approaches + +Structure your summary as follows: + +**Context** + +1. **Conversation Overview**: Provide a high-level narrative of the entire discussion, capturing the flow from initial objectives through all major topics and pivots. + +2. **Current Work Status**: Describe in detail the most recent task being addressed. Focus particularly on the latest messages to capture the immediate context before this summary request. + +3. **Technical Details**: Document all relevant technical elements including: + - Technologies, frameworks, and libraries used + - Coding conventions and patterns established + - Architectural decisions and design principles + - Configuration settings and environment details + +4. **Files and Code References**: List all files, code sections, or resources that were: + - Examined or reviewed + - Modified or updated + - Created from scratch + Prioritize the most recently touched items. + +5. **Problem-Solving History**: Summarize: + - Issues that were identified and resolved + - Solutions that were implemented + - Debugging approaches that were attempted + - Any ongoing troubleshooting efforts + +6. **Outstanding Work and Next Actions**: + - List all pending tasks explicitly requested by the user + - Outline the planned next steps for incomplete work + - Include verbatim quotes from recent messages showing the exact task in progress and where it was paused + - Add relevant code snippets where they provide clarity + +Output only the summary itself, with no preamble or meta-commentary. Conversation to summarize: {}"#; diff --git a/shai-llm/src/tool/max_context.rs b/shai-llm/src/tool/max_context.rs new file mode 100644 index 0000000..8cc5ab7 --- /dev/null +++ b/shai-llm/src/tool/max_context.rs @@ -0,0 +1,83 @@ +//! Utility for retrieving the maximum context length for known LLM models. + +const MIN_SIMILARITY_THRESHOLD: f64 = 0.6; + +/// Calculate similarity score between two strings using Jaro-Winkler distance +fn similarity_score(s1: &str, s2: &str) -> f64 { + let s1_lower = s1.to_lowercase(); + let s2_lower = s2.to_lowercase(); + + let s1_chars: Vec = s1_lower.chars().collect(); + let s2_chars: Vec = s2_lower.chars().collect(); + + let max_len = s1_chars.len().max(s2_chars.len()); + if max_len == 0 { + return 1.0; + } + + let mut matches = 0; + let min_len = s1_chars.len().min(s2_chars.len()); + + for i in 0..min_len { + if s1_chars[i] == s2_chars[i] { + matches += 1; + } + } + + if s1_lower.contains(&s2_lower) || s2_lower.contains(&s1_lower) { + matches += min_len / 2; + } + + matches as f64 / max_len as f64 +} + +pub fn get_max_context(model_name: &str) -> usize { + let models = [ + // OpenAI models + ("gpt-oss", 131_000), + + // Mistral models + ("mistral-small-3-2", 128_000), + ("mistral-7b", 32_000), + ("mistral-nemo", 32_000), + ("mixtral-8x7b", 32_000), + + // Qwen models + ("qwen3", 32_000), + ("qwen-2-5", 32_000), + + // Llama models + ("llama-3-1", 131_000), + ("llama-3_3", 131_000), + ("meta-llama-3_3", 131_000), + ("meta-llama-3_1", 131_000), + + // Deepseek models + ("deepseek-r1", 128_000), + ]; + + // Try exact match first + for (model, context) in models.iter() { + if *model == model_name { + return *context; + } + } + + // Fuzzy matching with minimum threshold + let mut best_match: Option<(f64, usize)> = None; + + for (model, context) in models.iter() { + let score = similarity_score(model_name, model); + if score >= MIN_SIMILARITY_THRESHOLD { + if let Some((best_score, _)) = best_match { + if score > best_score { + best_match = Some((score, *context)); + } + } else { + best_match = Some((score, *context)); + } + } + } + + best_match.map(|(_, context)| context).unwrap_or(30_096) +} diff --git a/shai-llm/src/tool/mod.rs b/shai-llm/src/tool/mod.rs index 6a3eaaf..69ff3d7 100644 --- a/shai-llm/src/tool/mod.rs +++ b/shai-llm/src/tool/mod.rs @@ -11,4 +11,7 @@ pub use tool::{ToolDescription, ToolCallMethod, ToolBox, ContainsTool}; pub use call::{LlmToolCall,ToolCallAuto}; pub use call_structured_output::{AssistantResponse, StructuredOutputBuilder, IntoChatMessage}; pub use call_fc_auto::FunctionCallingAutoBuilder; -pub use call_fc_required::FunctionCallingRequiredBuilder; \ No newline at end of file +pub use call_fc_required::FunctionCallingRequiredBuilder; + +pub mod max_context; +pub use max_context::get_max_context; \ No newline at end of file From 38d1043807fb42c81fda43fbd86fc19811d24ad6 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 11:47:00 +0000 Subject: [PATCH 05/12] feat: context remaining --- shai-cli/src/tui/app.rs | 33 +++++++++++++++++++++++++++----- shai-llm/src/tool/max_context.rs | 2 +- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/shai-cli/src/tui/app.rs b/shai-cli/src/tui/app.rs index 52d98fd..4a0c725 100644 --- a/shai-cli/src/tui/app.rs +++ b/shai-cli/src/tui/app.rs @@ -22,6 +22,7 @@ use shai_core::logging::LoggingConfig; use shai_core::runners::coder::coder::coder; use shai_core::tools::{ToolCall, ToolResult}; use shai_llm::{LlmClient, ToolCallMethod}; +use shai_llm::tool::max_context::get_max_context; use ratatui::{ layout::{Constraint, Direction, Layout}, style::{Color, Style}, @@ -72,6 +73,8 @@ pub struct App<'a> { pub(crate) total_input_tokens: u32, pub(crate) total_output_tokens: u32, + pub(crate) current_tokens: usize, + pub(crate) max_context: usize, } @@ -83,6 +86,7 @@ impl App<'_> { let config = AgentConfig::load(agent_name)?; println!("\x1b[2m░ agent {} - {} on {}\x1b[0m", agent_name, config.llm_provider.model, config.llm_provider.provider); + self.max_context = get_max_context(&config.llm_provider.model); // Create agent from config let agent_builder = AgentBuilder::from_config(config).await?; @@ -91,7 +95,7 @@ impl App<'_> { // Use default coder agent let (llm, model) = ShaiConfig::get_llm().await?; println!("\x1b[2m░ {} on {}\x1b[0m", model, llm.provider().name()); - + self.max_context = get_max_context(&model); Box::new(coder(Arc::new(llm), model)) }; @@ -158,6 +162,13 @@ impl App<'_> { if let AgentEvent::TokenUsage { input_tokens, output_tokens } = &event { self.total_input_tokens += input_tokens; self.total_output_tokens += output_tokens; + self.current_tokens += (input_tokens + output_tokens) as usize; + } + // Update current_tokens after context compression + if let AgentEvent::ContextCompressed { current_tokens, .. } = &event { + if let Some(ct) = current_tokens { + self.current_tokens = *ct as usize; + } } Ok(()) @@ -182,6 +193,8 @@ impl App<'_> { permission_queue: VecDeque::new(), total_input_tokens: 0, total_output_tokens: 0, + current_tokens: 0, + max_context: 0, } } @@ -385,7 +398,7 @@ impl App<'_> { AppModalState::PermissionModal { widget } => widget.height(), }.max(5); let height = modal_height - + 1 + + 2 + self.running_tools.len() as u16; if let Some(ref mut terminal) = self.terminal { @@ -395,11 +408,13 @@ impl App<'_> { } terminal.draw(|frame| { - let [_, inprogress, modal] = Layout::vertical([ + let [_, inprogress, modal, ctx_line] = Layout::vertical([ Constraint::Length(1), // padding Constraint::Length(self.running_tools.len() as u16 + 1), // running tool (if any) - Constraint::Length(modal_height)]) // input or modal - .areas(frame.area()); + Constraint::Length(modal_height), + Constraint::Length(1) + ]).areas(frame.area()); + // draw running tool if !self.running_tools.is_empty() { @@ -418,6 +433,14 @@ impl App<'_> { widget.draw(frame, modal) } } + // Render context usage line + if self.max_context > 0 { + let used = self.current_tokens; + let remaining = if used > self.max_context { 0 } else { self.max_context - used }; + let percent = (remaining as f64 / self.max_context as f64) * 100.0; + let ctx_text = format!("Context: {:.1}% remaining", percent); + frame.render_widget(Span::styled(ctx_text, Style::default().fg(Color::DarkGray)), ctx_line); + } })?; } Ok(()) diff --git a/shai-llm/src/tool/max_context.rs b/shai-llm/src/tool/max_context.rs index 8cc5ab7..13f97fe 100644 --- a/shai-llm/src/tool/max_context.rs +++ b/shai-llm/src/tool/max_context.rs @@ -34,7 +34,7 @@ fn similarity_score(s1: &str, s2: &str) -> f64 { pub fn get_max_context(model_name: &str) -> usize { let models = [ // OpenAI models - ("gpt-oss", 131_000), + ("gpt-oss", 10_000), // Mistral models ("mistral-small-3-2", 128_000), From 800b699aab617eb243f37a00a0b77c6e0925085f Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 12:43:10 +0000 Subject: [PATCH 06/12] feat: manual compact trigger --- shai-cli/src/tui/command.rs | 7 +++ shai-cli/src/tui/helper.rs | 5 +- shai-core/src/agent/actions/brain.rs | 65 ++++++++++++++++++++++ shai-core/src/agent/agent.rs | 8 ++- shai-core/src/agent/events.rs | 6 +- shai-core/src/agent/protocol.rs | 7 +++ shai-core/src/agent/states/pause.rs | 4 ++ shai-core/src/runners/compacter/compact.rs | 24 ++++++++ 8 files changed, 121 insertions(+), 5 deletions(-) diff --git a/shai-cli/src/tui/command.rs b/shai-cli/src/tui/command.rs index 7288c80..fbbee5e 100644 --- a/shai-cli/src/tui/command.rs +++ b/shai-cli/src/tui/command.rs @@ -10,6 +10,7 @@ impl App<'_> { (("/auth","select a provider"), vec![]), (("/tc","set the tool call method: [fc | fc2 | so]"), vec!["method"]), (("/tokens","display token usage (input/output)"), vec![]), + (("/compact","trigger context compaction"), vec![]), ]) .into_iter() .map(|((cmd,desc),args)|((cmd.to_string(),desc.to_string()),args.into_iter().map(|s|s.to_string()).collect())) @@ -65,6 +66,12 @@ impl App<'_> { ); self.input.alert_msg(&msg, Duration::from_secs(5)); } + "/compact" => { + if let Some(ref agent) = self.agent { + let _ = agent.controller.trigger_context_compression().await; + self.input.alert_msg("Context compression triggered", Duration::from_secs(2)); + } + } _ => { self.input.alert_msg("command unknown", Duration::from_secs(1)); } diff --git a/shai-cli/src/tui/helper.rs b/shai-cli/src/tui/helper.rs index 79a543a..532f92b 100644 --- a/shai-cli/src/tui/helper.rs +++ b/shai-cli/src/tui/helper.rs @@ -15,14 +15,15 @@ impl HelpArea { " Available Commands:", " /exit exit from the tui", " /tc set tool call method: [auto | fc | fc2 | so]", - " /tokens display token usage" + " /tokens display token usage", + " /compact trigger context compaction" ].join("\n").to_string() } } impl HelpArea { pub fn height(&self) -> u16 { - 8 // content (3 general help lines + 1 blank + 1 header + 3 command lines) + 9 // content (3 general help lines + 1 blank + 1 header + 4 command lines) } pub fn draw(&self, f: &mut Frame, area: Rect) { diff --git a/shai-core/src/agent/actions/brain.rs b/shai-core/src/agent/actions/brain.rs index 1ba7fb0..577cdcb 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -106,6 +106,64 @@ impl AgentCore { Ok(()) } + /// Trigger manual context compression regardless of threshold + pub async fn check_and_compress_context_manual(&mut self) -> Result<(), AgentError> { + // Set state to Processing to block new messages + self.set_state(InternalAgentState::Processing { + task_name: "context_compression".to_string(), + tools_exec_at: Utc::now(), + cancellation_token: CancellationToken::new(), + }).await; + + let brain = self.brain.clone(); + let brain_read = brain.read().await; + + use std::any::Any; + + if let Some(coder_brain) = (&**brain_read as &dyn Any).downcast_ref::() { + if let Some(compressor) = &coder_brain.context_compressor { + let compressor_clone = compressor.clone(); + drop(brain_read); + + let trace = self.trace.read().await.clone(); + let mut compressor_clone = compressor_clone; + + // Force compression - manually call compress_messages_force + let (compressed_trace, compression_info) = compressor_clone.compress_messages_force(trace).await; + + // Update the trace with compressed version + { + let mut trace_write = self.trace.write().await; + *trace_write = compressed_trace; + } + + // Update the compressor in the brain + { + let mut brain_write = brain.write().await; + if let Some(coder_brain_mut) = (&mut **brain_write as &mut dyn Any).downcast_mut::() { + coder_brain_mut.context_compressor = Some(compressor_clone); + } + } + + // Emit compression event if compression occurred + if let Some(compression_info) = compression_info { + let _ = self.emit_event(AgentEvent::ContextCompressed { + original_message_count: compression_info.original_message_count, + compressed_message_count: compression_info.compressed_message_count, + tokens_before: compression_info.tokens_before, + current_tokens: compression_info.current_tokens, + max_tokens: compression_info.max_tokens, + ai_summary: compression_info.ai_summary, + }).await; + } + } + } + + // Return to Paused state after compression + self.set_state(InternalAgentState::Paused).await; + Ok(()) + } + /// Check if context compression is needed and apply it when task is complete async fn check_and_compress_context(&mut self) -> Result<(), AgentError> { // Extract compression logic from the brain if it's a CoderBrain @@ -125,6 +183,13 @@ impl AgentCore { let mut compressor_clone = compressor_clone; if compressor_clone.should_compress_conversation(&trace) { + // Set state to Processing to block new messages during compression + self.set_state(InternalAgentState::Processing { + task_name: "context_compression".to_string(), + tools_exec_at: Utc::now(), + cancellation_token: CancellationToken::new(), + }).await; + let (compressed_trace, compression_info) = compressor_clone.compress_messages(trace).await; // Update the trace with compressed version diff --git a/shai-core/src/agent/agent.rs b/shai-core/src/agent/agent.rs index 695f5ed..098fcd9 100644 --- a/shai-core/src/agent/agent.rs +++ b/shai-core/src/agent/agent.rs @@ -432,7 +432,13 @@ impl AgentCore { AgentRequest::WaitTurn => { self.handle_wait_turn(backchannel).await; return Ok(()); // We handle the response in the spawned task - } + } + AgentRequest::TriggerContextCompression => { + // Trigger context compression by calling the check_and_compress_context method + // We'll send a special internal event to trigger the compression + let _ = self.internal_tx.send(InternalAgentEvent::ManualCompressionRequested).map_err(|_| AgentError::SessionClosed)?; + Ok(AgentResponse::Ack) + } }.unwrap_or_else(|e| AgentResponse::Error { error: e.to_string() }); // ignore if channel is closed diff --git a/shai-core/src/agent/events.rs b/shai-core/src/agent/events.rs index 35582a4..828d321 100644 --- a/shai-core/src/agent/events.rs +++ b/shai-core/src/agent/events.rs @@ -45,10 +45,12 @@ pub enum InternalAgentEvent { response: UserResponse }, /// Permission response received from controller - PermissionResponseReceived { + PermissionResponseReceived { request_id: String, response: PermissionResponse - } + }, + /// Manual context compression requested by user + ManualCompressionRequested } /// Public events emitted to external controllers/UI diff --git a/shai-core/src/agent/protocol.rs b/shai-core/src/agent/protocol.rs index 9bfcce5..81fd8cb 100644 --- a/shai-core/src/agent/protocol.rs +++ b/shai-core/src/agent/protocol.rs @@ -40,6 +40,8 @@ pub enum AgentRequest { /// Drop controller IO, this closes it for all controller. /// Once this is done, it cannot be reopen! Droping, + /// Trigger context compression manually + TriggerContextCompression, } /// Commands that can be sent to a running agent @@ -174,4 +176,9 @@ impl AgentController { _ => Err(AgentError::InvalidResponse("Expected SudoStatus response".to_string())) } } + + /// Trigger context compression manually + pub async fn trigger_context_compression(&self) -> Result<(), AgentError> { + self.send(AgentRequest::TriggerContextCompression).await.map(|_| Ok(()))? + } } \ No newline at end of file diff --git a/shai-core/src/agent/states/pause.rs b/shai-core/src/agent/states/pause.rs index 89946a5..93ab120 100644 --- a/shai-core/src/agent/states/pause.rs +++ b/shai-core/src/agent/states/pause.rs @@ -8,6 +8,10 @@ impl AgentCore { // Silently ignore Ok(()) } + InternalAgentEvent::ManualCompressionRequested => { + // Trigger manual context compression + self.check_and_compress_context_manual().await + } _ => { // Paused state: All other events are illegal until user send something // ignore all events but log error diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index 4e44f5a..9fe3778 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -94,6 +94,24 @@ impl ContextCompressor { should_compress } + /// Force compress the conversation history regardless of thresholds + /// Keeps the system message and recent messages while summarizing middle conversation + /// Returns (compressed_messages, compression_info) + pub async fn compress_messages_force(&mut self, messages: Vec) -> (Vec, Option) { + // Count non-system messages to ensure we have something to compress + let non_system_count = messages.iter() + .filter(|msg| !matches!(msg, ChatMessage::System { .. })) + .count(); + + // Only compress if we have at least 2 messages (1 user-assistant pair) + if non_system_count <= 2 { + info!(target: "context_compression", "Not enough messages to compress (need > 2 non-system messages)"); + return (messages, None); + } + + self.compress_messages_internal(messages).await + } + /// Compress the conversation history by removing older messages and replacing with AI summary /// Keeps the system message and recent messages while summarizing middle conversation /// Returns (compressed_messages, compression_info) @@ -102,6 +120,12 @@ impl ContextCompressor { return (messages, None); } + self.compress_messages_internal(messages).await + } + + /// Internal method that performs the actual compression + async fn compress_messages_internal(&mut self, messages: Vec) -> (Vec, Option) { + let original_count = messages.len(); let tokens_before_compression = self.current_tokens; From b42fe6a3d41368b13ba362f7863336489a1e073f Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 12:47:56 +0000 Subject: [PATCH 07/12] fix: gpt-oss context --- shai-llm/src/tool/max_context.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shai-llm/src/tool/max_context.rs b/shai-llm/src/tool/max_context.rs index 13f97fe..8cc5ab7 100644 --- a/shai-llm/src/tool/max_context.rs +++ b/shai-llm/src/tool/max_context.rs @@ -34,7 +34,7 @@ fn similarity_score(s1: &str, s2: &str) -> f64 { pub fn get_max_context(model_name: &str) -> usize { let models = [ // OpenAI models - ("gpt-oss", 10_000), + ("gpt-oss", 131_000), // Mistral models ("mistral-small-3-2", 128_000), From 736b5d184af27e48360b162ed22b832f5c18065c Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 15:27:24 +0000 Subject: [PATCH 08/12] dev in progress --- shai-cli/src/headless/app.rs | 10 ++- shai-cli/src/headless/tools.rs | 4 +- shai-core/src/agent/actions/brain.rs | 9 ++- shai-core/src/agent/actions/tools.rs | 11 ++- shai-core/src/agent/agent.rs | 34 +++++---- shai-core/src/agent/builder.rs | 4 +- shai-core/src/agent/output/pretty.rs | 4 +- shai-core/src/runners/compacter/compact.rs | 37 ++++++--- shai-core/src/runners/compacter/prompt.rs | 89 ++++++++++++++-------- shai-core/src/tools/fs/find/find.rs | 2 +- shai-core/src/tools/fs/find/tests.rs | 2 +- shai-llm/src/tool/max_context.rs | 2 +- 12 files changed, 137 insertions(+), 71 deletions(-) diff --git a/shai-cli/src/headless/app.rs b/shai-cli/src/headless/app.rs index 0a826ef..365f47c 100644 --- a/shai-cli/src/headless/app.rs +++ b/shai-cli/src/headless/app.rs @@ -96,11 +96,15 @@ impl AppHeadless { .run().await; match result { - Ok(AgentResult { success, message, trace: agent_trace }) => { + Ok(AgentResult { success, message, trace: full_trace, compressed_trace }) => { if trace { - println!("{}", serde_json::to_string_pretty(&agent_trace)?); + // Display comparison between compressed and full trace + println!("========== COMPRESSED TRACE ({} messages) ==========", compressed_trace.len()); + println!("{}", serde_json::to_string_pretty(&compressed_trace)?); + println!("\n========== FULL TRACE ({} messages) ==========", full_trace.len()); + println!("{}", serde_json::to_string_pretty(&full_trace)?); } else { - if let Some(message) = agent_trace.last() { + if let Some(message) = full_trace.last() { match message { ChatMessage::Assistant { content: Some(ChatMessageContent::Text(content)), .. } => { println!("{}",content); diff --git a/shai-cli/src/headless/tools.rs b/shai-cli/src/headless/tools.rs index 220ab22..46dec06 100644 --- a/shai-cli/src/headless/tools.rs +++ b/shai-cli/src/headless/tools.rs @@ -39,7 +39,7 @@ impl ToolName { ToolName::Bash => "bash", ToolName::Edit => "edit", ToolName::Fetch => "fetch", - ToolName::Find => "find", + ToolName::Find => "search", ToolName::Ls => "ls", ToolName::MultiEdit => "multiedit", ToolName::Read => "read", @@ -54,7 +54,7 @@ impl ToolName { "bash" => Some(ToolName::Bash), "edit" => Some(ToolName::Edit), "fetch" => Some(ToolName::Fetch), - "find" => Some(ToolName::Find), + "search" => Some(ToolName::Find), "ls" => Some(ToolName::Ls), "multiedit" => Some(ToolName::MultiEdit), "read" => Some(ToolName::Read), diff --git a/shai-core/src/agent/actions/brain.rs b/shai-core/src/agent/actions/brain.rs index 577cdcb..70720e5 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -47,6 +47,7 @@ impl AgentCore { /// Process a brain task result pub async fn process_next_step(&mut self, result: Result) -> Result<(), AgentError> { + let _ = self.check_and_compress_context().await?; let ThinkerDecision{message, flow, token_usage, compression_info} = self.handle_brain_error(result).await?; let ChatMessage::Assistant { content, reasoning_content, tool_calls, .. } = message.clone() else { return self.handle_brain_error::( @@ -57,8 +58,10 @@ impl AgentCore { // Add the message to trace info!(target: "agent::think", reasoning_content = ?reasoning_content, content = ?content); let trace = self.trace.clone(); + let full_trace = self.full_trace.clone(); trace.write().await.push(message.clone()); - + full_trace.write().await.push(message.clone()); + // Emit event to external consumers let _ = self.emit_event(AgentEvent::BrainResult { timestamp: Utc::now(), @@ -97,9 +100,7 @@ impl AgentCore { ThinkerFlowControl::AgentContinue => { self.set_state(InternalAgentState::Running).await; } - ThinkerFlowControl::AgentPause => { - // Check if we need to compress context when task is complete - self.check_and_compress_context().await?; + ThinkerFlowControl::AgentPause => { self.set_state(InternalAgentState::Paused).await; } } diff --git a/shai-core/src/agent/actions/tools.rs b/shai-core/src/agent/actions/tools.rs index 067031a..0a8cad2 100644 --- a/shai-core/src/agent/actions/tools.rs +++ b/shai-core/src/agent/actions/tools.rs @@ -25,10 +25,11 @@ impl AgentCore { let available_tools = self.available_tools.clone(); let claims = self.permissions.clone(); let trace = self.trace.clone(); + let full_trace = self.full_trace.clone(); // Spawn a task to wait for all tool executions let mut join_handles = Vec::new(); - + // Spawn all tool executions for tc in tool_calls { let handle = Self::spawn_tool_static( @@ -39,6 +40,7 @@ impl AgentCore { claims.clone(), internal_tx.clone(), trace.clone(), + full_trace.clone(), ); join_handles.push(handle); } @@ -83,6 +85,7 @@ impl AgentCore { claims: Arc>, internal_tx: broadcast::Sender, trace: Arc>>, + full_trace: Arc>>, ) -> tokio::task::JoinHandle { tokio::spawn(async move { let tc_for_error = tc.clone(); @@ -144,10 +147,12 @@ impl AgentCore { // let's first add tool result to trace let _ = { - trace.write().await.push(ChatMessage::Tool { + let tool_message = ChatMessage::Tool { tool_call_id: call.tool_call_id.clone(), content: result.to_string() - }); + }; + trace.write().await.push(tool_message.clone()); + full_trace.write().await.push(tool_message); }; // Emit tool call finish event diff --git a/shai-core/src/agent/agent.rs b/shai-core/src/agent/agent.rs index 098fcd9..8c47a0c 100644 --- a/shai-core/src/agent/agent.rs +++ b/shai-core/src/agent/agent.rs @@ -49,7 +49,8 @@ pub trait Agent: Send + Sync { pub struct AgentResult { pub success: bool, pub message: String, - pub trace: Vec, + pub trace: Vec, // Full uncompressed trace + pub compressed_trace: Vec, // Compressed trace (after compression) } /// Core agent implementation that orchestrates any Thinker implementation @@ -65,6 +66,7 @@ pub struct AgentCore { /// agent state (manipulated by main looper + brain/tool coroutines) pub trace: Arc>>, + pub full_trace: Arc>>, // Complete history before compression pub available_tools: Vec>, pub permissions: Arc>, pub state: InternalAgentState, @@ -100,7 +102,8 @@ impl AgentCore { }, brain: Arc::new(RwLock::new(brain)), method: ToolCallMethod::FunctionCall, - trace: Arc::new(RwLock::new(trace)), + trace: Arc::new(RwLock::new(trace.clone())), + full_trace: Arc::new(RwLock::new(trace)), available_tools: available_tools.into_iter().map(|t| Arc::from(t) as Arc).collect(), permissions: Arc::new(RwLock::new(permissions)), state: InternalAgentState::Starting, @@ -287,12 +290,15 @@ impl AgentCore { match &self.state { InternalAgentState::Completed { success } => { debug!(target: "agent::terminated", "completed"); - let trace = self.trace.clone(); - let guard = trace.read().await; + let full_trace = self.full_trace.clone(); + let compressed_trace = self.trace.clone(); + let full_guard = full_trace.read().await; + let compressed_guard = compressed_trace.read().await; return Ok(AgentResult { success: success.clone(), message: "Agent completed".to_string(), - trace: guard.clone(), + trace: full_guard.clone(), + compressed_trace: compressed_guard.clone(), }); }, InternalAgentState::Failed { error } => { @@ -400,15 +406,17 @@ impl AgentCore { self.handle_event(InternalAgentEvent::CancelTask).await .and({ // Emit UserInput event - let _ = self.emit_event(AgentEvent::UserInput { - input: input.clone() + let _ = self.emit_event(AgentEvent::UserInput { + input: input.clone() }).await; - - self.trace.write().await.push(ChatMessage::User { - content: ChatMessageContent::Text(input), - name: None - }); - + + let user_message = ChatMessage::User { + content: ChatMessageContent::Text(input), + name: None + }; + self.trace.write().await.push(user_message.clone()); + self.full_trace.write().await.push(user_message); + self.set_state(InternalAgentState::Running).await; Ok(AgentResponse::Ack) }) diff --git a/shai-core/src/agent/builder.rs b/shai-core/src/agent/builder.rs index e0df1a2..93b4d19 100644 --- a/shai-core/src/agent/builder.rs +++ b/shai-core/src/agent/builder.rs @@ -146,7 +146,7 @@ impl AgentBuilder { // Add builtin tools based on config let builtin_tools_to_add = if config.tools.builtin.contains(&"*".to_string()) { // Add all builtin tools - vec!["bash", "edit", "multiedit", "fetch", "find", "ls", "read", "todo_read", "todo_write", "write"] + vec!["bash", "edit", "multiedit", "fetch", "search", "ls", "read", "todo_read", "todo_write", "write"] } else { // Add only specified tools config.tools.builtin.iter().map(|s| s.as_str()).collect() @@ -163,7 +163,7 @@ impl AgentBuilder { "edit" => tools.push(Box::new(EditTool::new(fs_log.clone()))), "multiedit" => tools.push(Box::new(MultiEditTool::new(fs_log.clone()))), "fetch" => tools.push(Box::new(FetchTool::new())), - "find" => tools.push(Box::new(FindTool::new())), + "search" => tools.push(Box::new(FindTool::new())), "ls" => tools.push(Box::new(LsTool::new())), "read" => tools.push(Box::new(ReadTool::new(fs_log.clone()))), "todo_read" => tools.push(Box::new(TodoReadTool::new(todo_storage.clone()))), diff --git a/shai-core/src/agent/output/pretty.rs b/shai-core/src/agent/output/pretty.rs index b370e3e..1c2f524 100644 --- a/shai-core/src/agent/output/pretty.rs +++ b/shai-core/src/agent/output/pretty.rs @@ -252,7 +252,7 @@ impl PrettyFormatter { } // Show first N lines for user display only for specific tools - if matches!(call.tool_name.as_str(), "ls" | "bash" | "edit" | "multiedit" | "find" | "todo_read" | "todo_write") { + if matches!(call.tool_name.as_str(), "ls" | "bash" | "edit" | "multiedit" | "search" | "todo_read" | "todo_write") { let preview_lines: Vec<&str> = tool_output.lines().take(self.max_preview_lines).collect(); if !preview_lines.is_empty() { let mut markdown_content = String::new(); @@ -291,7 +291,7 @@ impl PrettyFormatter { let param_names = match tool_name { "read" | "write" | "edit" | "multiedit" => vec!["file_path", "path"], "ls" | "glob" => vec!["path", "pattern"], - "find" | "grep" => vec!["pattern", "path"], + "search" | "grep" => vec!["pattern", "path"], "bash" => vec!["command"], _ => vec!["path", "file_path", "pattern", "command", "query", "input"] }; diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index 9fe3778..682631e 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -142,25 +142,42 @@ impl ContextCompressor { let mut middle_messages = Vec::new(); let mut recent_messages = Vec::new(); - // Separate messages into categories - for (index, message) in messages.iter().enumerate() { + // First pass: filter out old summary messages and collect non-system messages + let non_summary_messages: Vec = messages.iter() + .filter(|msg| { + // Filter out old summary messages + !matches!(msg, ChatMessage::System { name: Some(name), .. } if name == "summary") + }) + .cloned() + .collect(); + + // Second pass: categorize messages + let non_system_count = non_summary_messages.iter() + .filter(|msg| !matches!(msg, ChatMessage::System { .. })) + .count(); + + let mut non_system_index = 0; + for message in non_summary_messages { match message { ChatMessage::System { .. } => { + // Keep non-summary system messages (like the original system prompt) system_messages.push(message.clone()); } _ => { - // Keep only the last 2 messages (1 user-assistant pair) as recent - // Everything else goes to middle_messages for summarization - if index >= messages.len().saturating_sub(2) { + // Keep the last 6 non-system messages (2-3 complete interaction cycles) as recent + // This ensures we preserve enough context for the agent to understand + // what it was doing and avoid repeating actions + if non_system_index >= non_system_count.saturating_sub(6) { recent_messages.push(message.clone()); } else { middle_messages.push(message.clone()); } + non_system_index += 1; } } } - // Add system messages first + // Add system messages first (excluding old summaries) compressed.extend(system_messages); // Try to generate AI summary of middle conversation @@ -195,6 +212,8 @@ impl ContextCompressor { // Add recent messages compressed.extend(recent_messages); + self.current_tokens = summary_tokens; + // Safely create compression info with validation let compression_info = CompressionInfo { original_message_count: original_count, @@ -338,8 +357,8 @@ mod tests { assert!(compressor.should_compress()); } - #[test] - fn test_message_compression() { + #[tokio::test] + async fn test_message_compression() { let mut compressor = ContextCompressor::new(1000); compressor.current_tokens = 850; // Above 80% threshold @@ -374,7 +393,7 @@ mod tests { }, ]; - let compressed = compressor.compress_messages(messages); + let (compressed, _info) = compressor.compress_messages(messages).await; // Should contain: system message, compression notice, recent messages assert!(compressed.len() >= 4); diff --git a/shai-core/src/runners/compacter/prompt.rs b/shai-core/src/runners/compacter/prompt.rs index 8961015..87795af 100644 --- a/shai-core/src/runners/compacter/prompt.rs +++ b/shai-core/src/runners/compacter/prompt.rs @@ -1,43 +1,72 @@ -static COMPRESSION_SUMMARY_PROMPT: &str = r#"Generate a comprehensive summary of this conversation, focusing on capturing every detail necessary to seamlessly continue the work. +static COMPRESSION_SUMMARY_PROMPT: &str = r#"Compress this conversation by eliminating ONLY redundant information while preserving every unique piece of data needed to continue the work. -Your summary must preserve: -- All user requests and instructions (both completed and pending) -- Technical decisions, code patterns, and architectural choices -- The progression of work and problem-solving approaches +## ORIGINAL OBJECTIVE +Reproduce the FIRST user message VERBATIM - do not summarize or modify it: +""" +[Insert complete first user message here] +""" -Structure your summary as follows: +## CONVERSATION FACTS +Extract and organize ALL unique information from the conversation. If something was mentioned multiple times, include it only once. If it's unique information, include it even if it seems minor. -**Context** +### Technical Stack & Architecture +[Every technology, library, framework, pattern, or architectural decision mentioned - list each once with version if specified] -1. **Conversation Overview**: Provide a high-level narrative of the entire discussion, capturing the flow from initial objectives through all major topics and pivots. +### Files & Code +For each file mentioned: +- `filepath`: [what it does] | [changes made] | [current state] | [remaining work] -2. **Current Work Status**: Describe in detail the most recent task being addressed. Focus particularly on the latest messages to capture the immediate context before this summary request. +Include code snippets where they contain decisions, patterns, or solutions that need to be preserved. -3. **Technical Details**: Document all relevant technical elements including: - - Technologies, frameworks, and libraries used - - Coding conventions and patterns established - - Architectural decisions and design principles - - Configuration settings and environment details +### User Requests & Instructions +List chronologically, one per line, using EXACT quotes: +1. "[exact user request]" → Status: [completed/in-progress/pending] → [deliverable or progress made] +2. "[exact user request]" → Status: [completed/in-progress/pending] → [deliverable or progress made] +[continue for all requests] -4. **Files and Code References**: List all files, code sections, or resources that were: - - Examined or reviewed - - Modified or updated - - Created from scratch - Prioritize the most recently touched items. +### Technical Decisions & Solutions +[Every problem solved, decision made, or approach chosen - include the reasoning if it was discussed] +- [Decision/Solution]: [context] → [implementation] → [outcome] -5. **Problem-Solving History**: Summarize: - - Issues that were identified and resolved - - Solutions that were implemented - - Debugging approaches that were attempted - - Any ongoing troubleshooting efforts +### Configuration & Environment +[Everything about setup, environment variables, dependencies, compilation flags, etc. - each item once] -6. **Outstanding Work and Next Actions**: - - List all pending tasks explicitly requested by the user - - Outline the planned next steps for incomplete work - - Include verbatim quotes from recent messages showing the exact task in progress and where it was paused - - Add relevant code snippets where they provide clarity +### Coding Conventions & Patterns +[Any established patterns for naming, structure, error handling, testing, etc. that were agreed upon] -Output only the summary itself, with no preamble or meta-commentary. +### Current State & Progress +**Most recent exchange**: +- Last user message: "[exact quote]" +- What was being done: [precise description] +- Exact stopping point: [file, function, line of code, or specific action] + +**State of work**: +[For each active work item: what's done, what's in progress, what remains] + +### Outstanding Work (In Order) +**Next immediate action**: [specific next step with file/function names] + +**Remaining tasks**: +1. [task] - [any context needed] +2. [task] - [any context needed] +[ordered by priority or logical sequence] + +**Deferred/Future**: [anything explicitly postponed or lower priority] + +### Important Constraints & Notes +[User preferences, requirements, known bugs, warnings, gotchas - anything that affects how to proceed] + +--- + +**Compression instructions**: +- Keep the FIRST user message completely intact +- Include information ONCE - if repeated in conversation, write it only once +- Use exact quotes for user requests - never paraphrase what the user asked for +- Include ALL unique technical details, decisions, and code patterns +- Preserve ALL code snippets that show solutions or patterns +- Be dense but complete - no fluff, but don't omit facts +- Organize logically but don't create redundant categories +- Most recent context is most critical - ensure it's captured precisely Conversation to summarize: {}"#; diff --git a/shai-core/src/tools/fs/find/find.rs b/shai-core/src/tools/fs/find/find.rs index ab5c21f..f623662 100644 --- a/shai-core/src/tools/fs/find/find.rs +++ b/shai-core/src/tools/fs/find/find.rs @@ -110,7 +110,7 @@ impl FindTool { } } -#[tool(name = "find", description = r#"A high-performance search utility for locating files or specific text within files across the project. +#[tool(name = "search", description = r#"A high-performance search utility for locating files or specific text within files across the project. **Core Functionality:** - Employs regular expressions for powerful content searches, allowing for complex pattern matching. diff --git a/shai-core/src/tools/fs/find/tests.rs b/shai-core/src/tools/fs/find/tests.rs index 486aee4..e664a41 100644 --- a/shai-core/src/tools/fs/find/tests.rs +++ b/shai-core/src/tools/fs/find/tests.rs @@ -9,7 +9,7 @@ use std::fs; #[tokio::test] async fn test_find_tool_creation() { let tool = FindTool::new(); - assert_eq!(&tool.name(), "find"); + assert_eq!(&tool.name(), "search"); assert!(!tool.description().is_empty()); } diff --git a/shai-llm/src/tool/max_context.rs b/shai-llm/src/tool/max_context.rs index 8cc5ab7..0de96d1 100644 --- a/shai-llm/src/tool/max_context.rs +++ b/shai-llm/src/tool/max_context.rs @@ -34,7 +34,7 @@ fn similarity_score(s1: &str, s2: &str) -> f64 { pub fn get_max_context(model_name: &str) -> usize { let models = [ // OpenAI models - ("gpt-oss", 131_000), + ("gpt-oss", 100), // Mistral models ("mistral-small-3-2", 128_000), From 39c02db1e3df2fd0fa3df44d176a1150f450ed90 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Tue, 30 Sep 2025 22:35:48 +0200 Subject: [PATCH 09/12] refactor: prompt, debug log and always use latest user prompt --- shai-cli/src/headless/app.rs | 8 +- shai-core/src/agent/actions/brain.rs | 6 +- shai-core/src/runners/compacter/compact.rs | 50 +++++++---- shai-core/src/runners/compacter/prompt.rs | 96 +++++++++------------- 4 files changed, 80 insertions(+), 80 deletions(-) diff --git a/shai-cli/src/headless/app.rs b/shai-cli/src/headless/app.rs index 365f47c..de90ff1 100644 --- a/shai-cli/src/headless/app.rs +++ b/shai-cli/src/headless/app.rs @@ -98,10 +98,10 @@ impl AppHeadless { match result { Ok(AgentResult { success, message, trace: full_trace, compressed_trace }) => { if trace { - // Display comparison between compressed and full trace - println!("========== COMPRESSED TRACE ({} messages) ==========", compressed_trace.len()); - println!("{}", serde_json::to_string_pretty(&compressed_trace)?); - println!("\n========== FULL TRACE ({} messages) ==========", full_trace.len()); + #[cfg(debug_assertions)] + { + println!("{}", serde_json::to_string_pretty(&compressed_trace)?); + } println!("{}", serde_json::to_string_pretty(&full_trace)?); } else { if let Some(message) = full_trace.last() { diff --git a/shai-core/src/agent/actions/brain.rs b/shai-core/src/agent/actions/brain.rs index 70720e5..0a80370 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -127,10 +127,11 @@ impl AgentCore { drop(brain_read); let trace = self.trace.read().await.clone(); + let full_trace = self.full_trace.read().await.clone(); let mut compressor_clone = compressor_clone; // Force compression - manually call compress_messages_force - let (compressed_trace, compression_info) = compressor_clone.compress_messages_force(trace).await; + let (compressed_trace, compression_info) = compressor_clone.compress_messages_force(trace, full_trace).await; // Update the trace with compressed version { @@ -191,7 +192,8 @@ impl AgentCore { cancellation_token: CancellationToken::new(), }).await; - let (compressed_trace, compression_info) = compressor_clone.compress_messages(trace).await; + let full_trace = self.full_trace.read().await.clone(); + let (compressed_trace, compression_info) = compressor_clone.compress_messages(trace, full_trace).await; // Update the trace with compressed version { diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index 682631e..56b7293 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -97,7 +97,7 @@ impl ContextCompressor { /// Force compress the conversation history regardless of thresholds /// Keeps the system message and recent messages while summarizing middle conversation /// Returns (compressed_messages, compression_info) - pub async fn compress_messages_force(&mut self, messages: Vec) -> (Vec, Option) { + pub async fn compress_messages_force(&mut self, messages: Vec, full_trace: Vec) -> (Vec, Option) { // Count non-system messages to ensure we have something to compress let non_system_count = messages.iter() .filter(|msg| !matches!(msg, ChatMessage::System { .. })) @@ -109,22 +109,22 @@ impl ContextCompressor { return (messages, None); } - self.compress_messages_internal(messages).await + self.compress_messages_internal(messages, full_trace).await } /// Compress the conversation history by removing older messages and replacing with AI summary /// Keeps the system message and recent messages while summarizing middle conversation /// Returns (compressed_messages, compression_info) - pub async fn compress_messages(&mut self, messages: Vec) -> (Vec, Option) { + pub async fn compress_messages(&mut self, messages: Vec, full_trace: Vec) -> (Vec, Option) { if !self.should_compress_conversation(&messages) { return (messages, None); } - self.compress_messages_internal(messages).await + self.compress_messages_internal(messages, full_trace).await } /// Internal method that performs the actual compression - async fn compress_messages_internal(&mut self, messages: Vec) -> (Vec, Option) { + async fn compress_messages_internal(&mut self, messages: Vec, full_trace: Vec) -> (Vec, Option) { let original_count = messages.len(); let tokens_before_compression = self.current_tokens; @@ -137,6 +137,19 @@ impl ContextCompressor { "Compressing context due to token limit" ); + // Extract the most recent user message from the full conversation history (full_trace) + let first_user_message = full_trace.iter() + .rev() + .find_map(|msg| { + if let ChatMessage::User { content, .. } = msg { + if let ChatMessageContent::Text(text) = content { + return Some(text.clone()); + } + } + None + }) + .unwrap_or_else(|| "[No user message found]".to_string()); + let mut compressed = Vec::new(); let mut system_messages = Vec::new(); let mut middle_messages = Vec::new(); @@ -157,7 +170,7 @@ impl ContextCompressor { .count(); let mut non_system_index = 0; - for message in non_summary_messages { + for message in &non_summary_messages { match message { ChatMessage::System { .. } => { // Keep non-summary system messages (like the original system prompt) @@ -181,8 +194,9 @@ impl ContextCompressor { compressed.extend(system_messages); // Try to generate AI summary of middle conversation + // Pass all non-summary messages and the first user message from full_trace let (ai_summary, summary_tokens) = if !middle_messages.is_empty() { - match self.summarize_conversation(&middle_messages).await { + match self.summarize_conversation(&non_summary_messages, &first_user_message).await { Ok((summary, tokens)) => { info!(target: "context_compression", "Successfully generated AI summary"); compressed.push(ChatMessage::System { @@ -255,7 +269,7 @@ impl ContextCompressor { /// Create a summary of the conversation history using AI /// Returns (summary_text, summary_tokens_used) - async fn summarize_conversation(&mut self, messages: &[ChatMessage]) -> Result<(String, u32), String> { + async fn summarize_conversation(&mut self, messages: &[ChatMessage], first_user_message: &str) -> Result<(String, u32), String> { let Some(ref llm_client) = self.llm_client else { return Err("No LLM client available for summarization".to_string()); }; @@ -281,7 +295,7 @@ impl ContextCompressor { ChatMessage::System { content, .. } => { if let ChatMessageContent::Text(text) = content { // Skip system prompts in summary, only include actual system messages - if !text.contains("CODER_BASE_PROMPT") && !text.contains("You are") { + if !text.contains("CODER_BASE_PROMPT") { conversation_text.push_str(&format!("System: {}\n", text)); } } @@ -290,15 +304,21 @@ impl ContextCompressor { } } - let summary_prompt = get_compression_summary_prompt(&conversation_text); + let summary_prompt = get_compression_summary_prompt(); let summary_request = ChatCompletionParametersBuilder::default() .model(model) - .messages(vec![ChatMessage::User { - content: ChatMessageContent::Text(summary_prompt), - name: None, - }]) - .temperature(0.1) // Low temperature for consistent summaries + .messages(vec![ + ChatMessage::System { + content: ChatMessageContent::Text(summary_prompt.to_string()), + name: None, + }, + ChatMessage::User { + content: ChatMessageContent::Text(format!("Original user request: \"{}\"\n\nFull conversation:\n{}", first_user_message, conversation_text)), + name: None, + }, + ]) + .temperature(0.1) .build() .map_err(|e| format!("Failed to build summary request: {}", e))?; diff --git a/shai-core/src/runners/compacter/prompt.rs b/shai-core/src/runners/compacter/prompt.rs index 87795af..7ac4e9e 100644 --- a/shai-core/src/runners/compacter/prompt.rs +++ b/shai-core/src/runners/compacter/prompt.rs @@ -1,76 +1,54 @@ -static COMPRESSION_SUMMARY_PROMPT: &str = r#"Compress this conversation by eliminating ONLY redundant information while preserving every unique piece of data needed to continue the work. +static COMPRESSION_SUMMARY_PROMPT: &str = r#"You are compressing conversation history. Your summary will replace older messages, so it must contain ALL information needed to continue the task. -## ORIGINAL OBJECTIVE -Reproduce the FIRST user message VERBATIM - do not summarize or modify it: -""" -[Insert complete first user message here] -""" +**CRITICAL RULES:** +1. Extract the FIRST user message from the conversation and reproduce it EXACTLY (word for word) +2. List every file that was read, with key data extracted from each file +3. List every action taken (reads, edits, tool calls, reasoning) +4. Identify what the assistant was doing and what remains to be done +5. Be factual and complete - losing information breaks the conversation flow -## CONVERSATION FACTS -Extract and organize ALL unique information from the conversation. If something was mentioned multiple times, include it only once. If it's unique information, include it even if it seems minor. +**FORMAT YOUR SUMMARY LIKE THIS:** -### Technical Stack & Architecture -[Every technology, library, framework, pattern, or architectural decision mentioned - list each once with version if specified] +**Original user request (verbatim):** +"[exact first user message here]" -### Files & Code -For each file mentioned: -- `filepath`: [what it does] | [changes made] | [current state] | [remaining work] +**Actions completed:** +- Read file X: [key content/data from file] +- Read file Y: [key content/data from file] +- [any other actions taken] -Include code snippets where they contain decisions, patterns, or solutions that need to be preserved. +**Key information extracted:** +- [important data point 1] +- [important data point 2] +- [etc.] -### User Requests & Instructions -List chronologically, one per line, using EXACT quotes: -1. "[exact user request]" → Status: [completed/in-progress/pending] → [deliverable or progress made] -2. "[exact user request]" → Status: [completed/in-progress/pending] → [deliverable or progress made] -[continue for all requests] +**Current state:** +[What was being done at the end of this conversation segment] -### Technical Decisions & Solutions -[Every problem solved, decision made, or approach chosen - include the reasoning if it was discussed] -- [Decision/Solution]: [context] → [implementation] → [outcome] +**Next steps:** +[What remains to be done to complete the user's request] -### Configuration & Environment -[Everything about setup, environment variables, dependencies, compilation flags, etc. - each item once] +**EXAMPLE:** +If the user said "read file.txt and summarize it", and the assistant read the file containing "Hello World", your summary should be: -### Coding Conventions & Patterns -[Any established patterns for naming, structure, error handling, testing, etc. that were agreed upon] +**Original user request (verbatim):** +"read file.txt and summarize it" -### Current State & Progress -**Most recent exchange**: -- Last user message: "[exact quote]" -- What was being done: [precise description] -- Exact stopping point: [file, function, line of code, or specific action] +**Actions completed:** +- Read file.txt: Contains "Hello World" -**State of work**: -[For each active work item: what's done, what's in progress, what remains] +**Key information extracted:** +- file.txt content: "Hello World" -### Outstanding Work (In Order) -**Next immediate action**: [specific next step with file/function names] +**Current state:** +File has been read, summary needs to be provided to user -**Remaining tasks**: -1. [task] - [any context needed] -2. [task] - [any context needed] -[ordered by priority or logical sequence] - -**Deferred/Future**: [anything explicitly postponed or lower priority] - -### Important Constraints & Notes -[User preferences, requirements, known bugs, warnings, gotchas - anything that affects how to proceed] +**Next steps:** +Provide summary of file.txt to the user --- +"#; -**Compression instructions**: -- Keep the FIRST user message completely intact -- Include information ONCE - if repeated in conversation, write it only once -- Use exact quotes for user requests - never paraphrase what the user asked for -- Include ALL unique technical details, decisions, and code patterns -- Preserve ALL code snippets that show solutions or patterns -- Be dense but complete - no fluff, but don't omit facts -- Organize logically but don't create redundant categories -- Most recent context is most critical - ensure it's captured precisely - -Conversation to summarize: -{}"#; - -pub fn get_compression_summary_prompt(conversation_text: &str) -> String { - COMPRESSION_SUMMARY_PROMPT.replace("{}", conversation_text) +pub fn get_compression_summary_prompt() -> String { + COMPRESSION_SUMMARY_PROMPT.to_string() } \ No newline at end of file From 5ef91d14c68ba5a66d2b22c43d88435177129283 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Wed, 1 Oct 2025 00:21:31 +0200 Subject: [PATCH 10/12] feat: fix gpt-oss context & try to get ONLY output token without reasoning --- shai-llm/src/providers/openai_compatible.rs | 16 ++++++++++++++-- shai-llm/src/providers/ovhcloud.rs | 16 ++++++++++++++-- shai-llm/src/tool/max_context.rs | 2 +- 3 files changed, 29 insertions(+), 5 deletions(-) diff --git a/shai-llm/src/providers/openai_compatible.rs b/shai-llm/src/providers/openai_compatible.rs index 5253499..83a444f 100644 --- a/shai-llm/src/providers/openai_compatible.rs +++ b/shai-llm/src/providers/openai_compatible.rs @@ -33,6 +33,18 @@ impl OpenAICompatibleProvider { _ => None } } + + fn adjust_response_tokens(&self, mut response: ChatCompletionResponse) -> ChatCompletionResponse { + if let Some(ref mut usage) = response.usage { + if let Some(details) = &usage.completion_tokens_details { + let reasoning_tokens = details.reasoning_tokens; + if let Some(completion) = usage.completion_tokens { + usage.completion_tokens = Some(completion.saturating_sub(reasoning_tokens)); + } + } + } + response + } } #[async_trait] @@ -44,10 +56,10 @@ impl LlmProvider for OpenAICompatibleProvider { } async fn chat(&self, request: ChatCompletionParameters) -> Result { - let mut response = self.client.chat().create(request).await + let response = self.client.chat().create(request).await .map_err(|e| Box::new(e) as LlmError)?; - Ok(response) + Ok(self.adjust_response_tokens(response)) } async fn chat_stream(&self, mut request: ChatCompletionParameters) -> Result { diff --git a/shai-llm/src/providers/ovhcloud.rs b/shai-llm/src/providers/ovhcloud.rs index 8ae6ab0..63e7fa8 100644 --- a/shai-llm/src/providers/ovhcloud.rs +++ b/shai-llm/src/providers/ovhcloud.rs @@ -45,6 +45,18 @@ impl OvhCloudProvider { request } + + fn adjust_response_tokens(&self, mut response: ChatCompletionResponse) -> ChatCompletionResponse { + if let Some(ref mut usage) = response.usage { + if let Some(details) = &usage.completion_tokens_details { + let reasoning_tokens = details.reasoning_tokens; + if let Some(completion) = usage.completion_tokens { + usage.completion_tokens = Some(completion.saturating_sub(reasoning_tokens)); + } + } + } + response + } } #[async_trait] @@ -67,10 +79,10 @@ impl LlmProvider for OvhCloudProvider { async fn chat(&self, request: ChatCompletionParameters) -> Result { let sanitized_request = self.sanitize_request(request); - let mut response = self.client.chat().create(sanitized_request).await + let response = self.client.chat().create(sanitized_request).await .map_err(|e| Box::new(e) as LlmError)?; - Ok(response) + Ok(self.adjust_response_tokens(response)) } async fn chat_stream(&self, mut request: ChatCompletionParameters) -> Result { diff --git a/shai-llm/src/tool/max_context.rs b/shai-llm/src/tool/max_context.rs index 0de96d1..8cc5ab7 100644 --- a/shai-llm/src/tool/max_context.rs +++ b/shai-llm/src/tool/max_context.rs @@ -34,7 +34,7 @@ fn similarity_score(s1: &str, s2: &str) -> f64 { pub fn get_max_context(model_name: &str) -> usize { let models = [ // OpenAI models - ("gpt-oss", 100), + ("gpt-oss", 131_000), // Mistral models ("mistral-small-3-2", 128_000), From a5e478ad8587816eb7cdbba8449426f7747481a4 Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Thu, 2 Oct 2025 07:53:07 +0000 Subject: [PATCH 11/12] refactor(coder): refactor token usage extraction to if-let --- shai-core/src/runners/coder/coder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shai-core/src/runners/coder/coder.rs b/shai-core/src/runners/coder/coder.rs index e82fa12..41d6b93 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -118,7 +118,7 @@ impl Brain for CoderBrain { .map_err(|e| AgentError::LlmError(e.to_string()))?; // Extract token usage information - let token_usage = brain_decision.usage.as_ref().map(|usage| { + let token_usage = if let Some(usage) = &brain_decision.usage { let input = usage.prompt_tokens.unwrap_or(0); let output = usage.completion_tokens.unwrap_or(0); debug!(target: "brain::coder::tokens", From 21650f0816f066e268bbb9c002f1e2752601209f Mon Sep 17 00:00:00 2001 From: Pierric Buchez Date: Thu, 2 Oct 2025 08:12:48 +0000 Subject: [PATCH 12/12] feat(compacter): retain tool and system messages in compressed log --- shai-core/src/runners/compacter/compact.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index 56b7293..0de6de8 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -292,12 +292,12 @@ impl ContextCompressor { conversation_text.push_str(&format!("Assistant: {}\n", text)); } } + ChatMessage::Tool { content, .. } => { + conversation_text.push_str(&format!("Tool: {}\n", content)); + } ChatMessage::System { content, .. } => { if let ChatMessageContent::Text(text) = content { - // Skip system prompts in summary, only include actual system messages - if !text.contains("CODER_BASE_PROMPT") { - conversation_text.push_str(&format!("System: {}\n", text)); - } + conversation_text.push_str(&format!("System: {}\n", text)); } } _ => {} // Skip other message types