diff --git a/crates/q_cli/src/cli/chat/conversation_state.rs b/crates/q_cli/src/cli/chat/conversation_state.rs index 3bd831fac7..7d053470ac 100644 --- a/crates/q_cli/src/cli/chat/conversation_state.rs +++ b/crates/q_cli/src/cli/chat/conversation_state.rs @@ -5,6 +5,7 @@ use std::collections::{ use std::env; use std::sync::Arc; +use aws_smithy_types::Document; use fig_api_client::model::{ AssistantResponseMessage, ChatMessage, @@ -32,11 +33,11 @@ use tracing::{ warn, }; -use super::chat_state::{ +use super::context::ContextManager; +use super::summarization_state::{ MAX_CHARS, TokenWarningLevel, }; -use super::context::ContextManager; use super::tools::{ QueuedTool, ToolSpec, @@ -123,14 +124,8 @@ impl ConversationState { } } - /// Returns a vector representation of the history (for accessors) - pub fn history_as_vec(&self) -> Vec { - self.history.iter().cloned().collect() - } - - /// Returns the length of the conversation history - pub fn history_len(&self) -> usize { - self.history.len() + pub fn history(&self) -> &VecDeque { + &self.history } /// Clears the conversation history and optionally the summary. @@ -476,10 +471,7 @@ impl ConversationState { /// Calculate the total character count in the conversation pub fn calculate_char_count(&self) -> usize { - // Calculate total character count in all messages let mut total_chars = 0; - - // Count characters in history for message in &self.history { match message { ChatMessage::UserInputMessage(msg) => { @@ -490,14 +482,11 @@ impl ConversationState { for result in results { for content in &result.content { match content { - ToolResultContentBlock::Text(text) => total_chars += text.len(), + ToolResultContentBlock::Text(text) => { + total_chars += text.len(); + }, ToolResultContentBlock::Json(doc) => { - if let Some(s) = doc.as_string() { - total_chars += s.len(); - } else { - // Approximate JSON size - total_chars += 100; - } + total_chars += calculate_document_char_count(doc); }, } } @@ -507,10 +496,12 @@ impl ConversationState { }, ChatMessage::AssistantResponseMessage(msg) => { total_chars += msg.content.len(); - // Add tool uses if any if let Some(tool_uses) = &msg.tool_uses { - // Approximation for tool uses - total_chars += tool_uses.len() * 200; + total_chars += tool_uses + .iter() + .map(|v| calculate_document_char_count(&v.input)) + .reduce(|acc, e| acc + e) + .unwrap_or_default(); } }, } @@ -607,8 +598,22 @@ fn build_shell_state() -> ShellState { } } +fn calculate_document_char_count(document: &Document) -> usize { + match document { + Document::Object(hash_map) => hash_map + .values() + .fold(0, |acc, e| acc + calculate_document_char_count(e)), + Document::Array(vec) => vec.iter().fold(0, |acc, e| acc + calculate_document_char_count(e)), + Document::Number(_) => 1, + Document::String(s) => s.len(), + Document::Bool(_) => 1, + Document::Null => 1, + } +} + #[cfg(test)] mod tests { + use aws_smithy_types::Number; use fig_api_client::model::{ AssistantResponseMessage, ToolResultStatus, @@ -635,6 +640,50 @@ mod tests { println!("{env_state:?}"); } + #[test] + fn test_calculate_document_char_count() { + // Test simple types + assert_eq!(calculate_document_char_count(&Document::String("hello".to_string())), 5); + assert_eq!(calculate_document_char_count(&Document::Number(Number::PosInt(123))), 1); + assert_eq!(calculate_document_char_count(&Document::Bool(true)), 1); + assert_eq!(calculate_document_char_count(&Document::Null), 1); + + // Test array + let array = Document::Array(vec![ + Document::String("test".to_string()), + Document::Number(Number::PosInt(42)), + Document::Bool(false), + ]); + assert_eq!(calculate_document_char_count(&array), 6); // "test" (4) + Number (1) + Bool (1) + + // Test object + let mut obj = HashMap::new(); + obj.insert("key1".to_string(), Document::String("value1".to_string())); + obj.insert("key2".to_string(), Document::Number(Number::PosInt(99))); + let object = Document::Object(obj); + assert_eq!(calculate_document_char_count(&object), 7); // "value1" (6) + Number (1) + + // Test nested structure + let mut nested_obj = HashMap::new(); + let mut inner_obj = HashMap::new(); + inner_obj.insert("inner_key".to_string(), Document::String("inner_value".to_string())); + nested_obj.insert("outer_key".to_string(), Document::Object(inner_obj)); + nested_obj.insert( + "array_key".to_string(), + Document::Array(vec![ + Document::String("item1".to_string()), + Document::String("item2".to_string()), + ]), + ); + + let complex = Document::Object(nested_obj); + assert_eq!(calculate_document_char_count(&complex), 21); // "inner_value" (11) + "item1" (5) + "item2" (5) + + // Test empty structures + assert_eq!(calculate_document_char_count(&Document::Array(vec![])), 0); + assert_eq!(calculate_document_char_count(&Document::Object(HashMap::new())), 0); + } + fn assert_conversation_state_invariants(state: FigConversationState, i: usize) { if let Some(Some(msg)) = state.history.as_ref().map(|h| h.first()) { assert!( diff --git a/crates/q_cli/src/cli/chat/mod.rs b/crates/q_cli/src/cli/chat/mod.rs index 49729eb410..a7060f1c14 100644 --- a/crates/q_cli/src/cli/chat/mod.rs +++ b/crates/q_cli/src/cli/chat/mod.rs @@ -1,4 +1,3 @@ -mod chat_state; mod command; mod context; mod conversation_state; @@ -6,14 +5,13 @@ mod input_source; mod parse; mod parser; mod prompt; +mod summarization_state; mod tools; -// Re-export types for use in this module use std::borrow::Cow; use std::collections::{ HashMap, HashSet, - VecDeque, }; use std::io::{ IsTerminal, @@ -31,10 +29,6 @@ use std::{ fs, }; -use chat_state::{ - SummarizationState, - TokenWarningLevel, -}; use command::{ Command, ToolsSubcommand, @@ -71,6 +65,10 @@ use fig_api_client::model::{ use fig_os_shim::Context; use fig_settings::Settings; use fig_util::CLI_BINARY_NAME; +use summarization_state::{ + SummarizationState, + TokenWarningLevel, +}; /// Help text for the compact command fn compact_help_text() -> String { @@ -79,11 +77,11 @@ fn compact_help_text() -> String { Conversation Compaction The /compact command summarizes the conversation history to free up context space -while preserving the essential information. This is useful for long-running conversations +while preserving essential information. This is useful for long-running conversations that may eventually reach memory constraints. Usage - /compact Summarize conversation and clear history + /compact Summarize the conversation and clear history /compact [prompt] Provide custom guidance for summarization /compact --summary Show the summary after compacting @@ -157,7 +155,7 @@ const WELCOME_TEXT: &str = color_print::cstr! {" /issue Report an issue or make a feature request /profile (Beta) Manage profiles for the chat session /context (Beta) Manage context files for a profile -/compact Summarize conversation to free up context space +/compact Summarize the conversation to free up context space /help Show the help dialogue /quit Quit the application @@ -175,7 +173,7 @@ const HELP_TEXT: &str = color_print::cstr! {" /editor Open $EDITOR (defaults to vi) to compose a prompt /help Show this help dialogue /quit Quit the application -/compact Summarize conversation to free up context space +/compact Summarize the conversation to free up context space help Show help for the compact command [prompt] Optional custom prompt to guide summarization --summary Display the summary after compacting @@ -858,7 +856,7 @@ where execute!( self.output, style::SetForegroundColor(Color::Green), - style::Print("\nConversation history and summary cleared.\n\n"), + style::Print("\nConversation history cleared.\n\n"), style::SetForegroundColor(Color::Reset) )?; @@ -890,7 +888,7 @@ where } // Check if conversation history is long enough to compact - if self.conversation_state.history_len() <= 3 { + if self.conversation_state.history().len() <= 3 { execute!( self.output, style::SetForegroundColor(Color::Yellow), @@ -905,23 +903,12 @@ where }); } - // Get private history field using accessor - let saved_history = VecDeque::from(self.conversation_state.history_as_vec().clone()); - // Set up summarization state with history, custom prompt, and show_summary flag let mut summarization_state = SummarizationState::with_prompt(prompt.clone()); - summarization_state.original_history = Some(saved_history); + summarization_state.original_history = Some(self.conversation_state.history().clone()); summarization_state.show_summary = show_summary; // Store the show_summary flag self.summarization_state = Some(summarization_state); - // Tell user we're working on the summary - execute!( - self.output, - style::SetForegroundColor(Color::Green), - style::Print("\nCompacting conversation history...\n"), - style::SetForegroundColor(Color::Reset) - )?; - // Create a summary request based on user input or default let summary_request = match prompt { Some(custom_prompt) => { @@ -971,7 +958,7 @@ where // Use spinner while we wait if self.interactive { - queue!(self.output, cursor::Hide)?; + execute!(self.output, cursor::Hide, style::Print("\n"))?; self.spinner = Some(Spinner::new(Spinners::Dots, "Creating summary...".to_string())); } @@ -1821,7 +1808,13 @@ where buf.push('\n'); } - if tool_name_being_recvd.is_none() && !buf.is_empty() && self.interactive && self.spinner.is_some() { + // TODO: refactor summarization into a separate ChatState value + if tool_name_being_recvd.is_none() + && !buf.is_empty() + && self.interactive + && self.spinner.is_some() + && !is_summarization + { drop(self.spinner.take()); queue!( self.output, @@ -1905,8 +1898,18 @@ where // Handle summarization completion if we were in summarization mode if let Some(summarization_state) = self.summarization_state.take() { + if self.spinner.is_some() { + drop(self.spinner.take()); + queue!( + self.output, + terminal::Clear(terminal::ClearType::CurrentLine), + cursor::MoveToColumn(0), + cursor::Show + )?; + } + // Get the latest message content (the summary) - let summary = match self.conversation_state.history_as_vec().last() { + let summary = match self.conversation_state.history().back() { Some(ChatMessage::AssistantResponseMessage(message)) => message.content.clone(), _ => "Summary could not be generated.".to_string(), }; @@ -1928,12 +1931,10 @@ where // Add the message self.conversation_state.push_assistant_message(special_message); - // Display confirmation message - execute!( self.output, style::SetForegroundColor(Color::Green), - style::Print("\nāœ… Conversation has been successfully summarized and cleared!\n\n"), + style::Print("āœ” Conversation history has been compacted successfully!\n\n"), style::SetForegroundColor(Color::DarkGrey) )?; @@ -1969,10 +1970,9 @@ where style::Print("\n"), style::SetAttribute(Attribute::Bold), style::Print(" CONVERSATION SUMMARY"), - style::SetAttribute(Attribute::Reset), style::Print("\n"), style::Print(&border), - style::SetForegroundColor(Color::Reset), + style::SetAttribute(Attribute::Reset), style::Print("\n\n"), style::Print(&summary), style::Print("\n\n"), @@ -1980,7 +1980,7 @@ where style::Print("This summary is stored in memory and available to the assistant.\n"), style::Print("It contains all important details from previous interactions.\n"), style::Print(&border), - style::Print("\n"), + style::Print("\n\n"), style::SetForegroundColor(Color::Reset) )?; } diff --git a/crates/q_cli/src/cli/chat/prompt.rs b/crates/q_cli/src/cli/chat/prompt.rs index 2e25e8245b..2be0e1828b 100644 --- a/crates/q_cli/src/cli/chat/prompt.rs +++ b/crates/q_cli/src/cli/chat/prompt.rs @@ -64,6 +64,9 @@ const COMMANDS: &[&str] = &[ "/context rm --global", "/context clear", "/context clear --global", + "/compact", + "/compact help", + "/compact --summary", ]; pub fn generate_prompt(current_profile: Option<&str>, warning: bool) -> String { diff --git a/crates/q_cli/src/cli/chat/chat_state.rs b/crates/q_cli/src/cli/chat/summarization_state.rs similarity index 100% rename from crates/q_cli/src/cli/chat/chat_state.rs rename to crates/q_cli/src/cli/chat/summarization_state.rs