diff --git a/README.md b/README.md index 2a33d2b..8571e7a 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-cli/src/headless/app.rs b/shai-cli/src/headless/app.rs index 0a826ef..de90ff1 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)?); + #[cfg(debug_assertions)] + { + println!("{}", serde_json::to_string_pretty(&compressed_trace)?); + } + 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-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-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 4b24d2b..0a80370 100644 --- a/shai-core/src/agent/actions/brain.rs +++ b/shai-core/src/agent/actions/brain.rs @@ -47,7 +47,8 @@ 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 _ = 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::( Err(AgentError::InvalidResponse(format!("ChatMessage::Assistant expected, but got {:?} instead", message)))).await.map(|_| () @@ -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(), @@ -72,6 +75,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![]); @@ -85,20 +100,140 @@ impl AgentCore { ThinkerFlowControl::AgentContinue => { self.set_state(InternalAgentState::Running).await; } - ThinkerFlowControl::AgentPause => { + ThinkerFlowControl::AgentPause => { self.set_state(InternalAgentState::Paused).await; } } 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 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, full_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 + 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) { + // 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 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 + { + 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/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 695f5ed..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) }) @@ -432,7 +440,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/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/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/events.rs b/shai-core/src/agent/events.rs index e0ed5e6..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 @@ -101,6 +103,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 +285,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..1c2f524 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)) } @@ -211,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(); @@ -250,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/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/config/config.rs b/shai-core/src/config/config.rs index 63b0caf..8689845 100644 --- a/shai-core/src/config/config.rs +++ b/shai-core/src/config/config.rs @@ -12,7 +12,8 @@ pub struct ProviderConfig { pub provider: String, pub env_vars: std::collections::HashMap, pub model: String, - pub tool_method: ToolCallMethod + pub tool_method: ToolCallMethod, + } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -36,13 +37,15 @@ impl ShaiConfig { provider, env_vars, model, - tool_method: ToolCallMethod::FunctionCall + tool_method: ToolCallMethod::FunctionCall, }; - + 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| { provider_config.provider == provider_name && @@ -228,7 +231,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, + }], 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 4c034c9..41d6b93 100644 --- a/shai-core/src/runners/coder/coder.rs +++ b/shai-core/src/runners/coder/coder.rs @@ -8,8 +8,10 @@ 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; use super::prompt::{render_system_prompt_template, get_todo_read}; @@ -19,26 +21,61 @@ 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() { + { + 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"); + Some(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() { + { + 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"); + Some(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 +88,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 +100,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) @@ -79,11 +118,31 @@ 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); - (input, output) - }); + debug!(target: "brain::coder::tokens", + input_tokens = input, + output_tokens = output, + 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 + }; // stop here if there's no other tool calls let message = brain_decision.choices.into_iter().next().unwrap().message; diff --git a/shai-core/src/runners/compacter/compact.rs b/shai-core/src/runners/compacter/compact.rs index e69de29..0de6de8 100644 --- a/shai-core/src/runners/compacter/compact.rs +++ b/shai-core/src/runners/compacter/compact.rs @@ -0,0 +1,438 @@ +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.90; + 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 + } + + /// 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, 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 { .. })) + .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, 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, full_trace: Vec) -> (Vec, Option) { + if !self.should_compress_conversation(&messages) { + return (messages, None); + } + + self.compress_messages_internal(messages, full_trace).await + } + + /// Internal method that performs the actual compression + 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; + + 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" + ); + + // 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(); + let mut recent_messages = Vec::new(); + + // 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 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 (excluding old summaries) + 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(&non_summary_messages, &first_user_message).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); + + self.current_tokens = summary_tokens; + + // 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], 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()); + }; + + 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::Tool { content, .. } => { + conversation_text.push_str(&format!("Tool: {}\n", content)); + } + ChatMessage::System { content, .. } => { + if let ChatMessageContent::Text(text) = content { + conversation_text.push_str(&format!("System: {}\n", text)); + } + } + _ => {} // Skip other message types + } + } + + let summary_prompt = get_compression_summary_prompt(); + + let summary_request = ChatCompletionParametersBuilder::default() + .model(model) + .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))?; + + 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()); + } + + #[tokio::test] + async 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, _info) = compressor.compress_messages(messages).await; + + // 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..7ac4e9e --- /dev/null +++ b/shai-core/src/runners/compacter/prompt.rs @@ -0,0 +1,54 @@ +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. + +**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 + +**FORMAT YOUR SUMMARY LIKE THIS:** + +**Original user request (verbatim):** +"[exact first user message here]" + +**Actions completed:** +- Read file X: [key content/data from file] +- Read file Y: [key content/data from file] +- [any other actions taken] + +**Key information extracted:** +- [important data point 1] +- [important data point 2] +- [etc.] + +**Current state:** +[What was being done at the end of this conversation segment] + +**Next steps:** +[What remains to be done to complete the user's request] + +**EXAMPLE:** +If the user said "read file.txt and summarize it", and the assistant read the file containing "Hello World", your summary should be: + +**Original user request (verbatim):** +"read file.txt and summarize it" + +**Actions completed:** +- Read file.txt: Contains "Hello World" + +**Key information extracted:** +- file.txt content: "Hello World" + +**Current state:** +File has been read, summary needs to be provided to user + +**Next steps:** +Provide summary of file.txt to the user + +--- +"#; + +pub fn get_compression_summary_prompt() -> String { + COMPRESSION_SUMMARY_PROMPT.to_string() +} \ No newline at end of file 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/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 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