tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Token counting reliability + summarization integration. (#3721)

Browse Source
This commit is contained in:
David Katz
2025-08-14 11:23:37 -04:00
committed by GitHub
parent f2e335cb81
commit 80826c2b23
15 changed files with 359 additions and 891 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/web.rs
+1 -1
View File
@@ -620,7 +620,7 @@ async fn process_message_streaming(
// For now, auto-summarize in web mode
// TODO: Implement proper UI for context handling
let (summarized_messages, _) =
let (summarized_messages, _, _) =
agent.summarize_context(messages.messages()).await?;
{
let mut session_msgs = session_messages.lock().await;
crates/goose-cli/src/session/mod.rs
+43 -3
View File
@@ -162,7 +162,7 @@ impl Session {
message_suffix: &str,
) -> Result<()> {
// Summarize messages to fit within context length
let (summarized_messages, _) = agent.summarize_context(messages.messages()).await?;
let (summarized_messages, _, _) = agent.summarize_context(messages.messages()).await?;
let msg = format!("Context maxed out\n{}\n{}", "-".repeat(50), message_suffix);
output::render_text(&msg, Some(Color::Yellow), true);
*messages = summarized_messages;
@@ -719,7 +719,7 @@ impl Session {
let provider = self.agent.provider().await?;
// Call the summarize_context method which uses the summarize_messages function
let (summarized_messages, _) = self
let (summarized_messages, _token_counts, summarization_usage) = self
.agent
.summarize_context(self.messages.messages())
.await?;
@@ -727,7 +727,7 @@ impl Session {
// Update the session messages with the summarized ones
self.messages = summarized_messages;
// Persist the summarized messages
// Persist the summarized messages and update session metadata with new token counts
if let Some(session_file) = &self.session_file {
let working_dir = std::env::current_dir().ok();
session::persist_messages_with_schedule_id(
@@ -738,6 +738,46 @@ impl Session {
working_dir,
)
.await?;
// Update session metadata with the new token counts from summarization
if let Some(usage) = summarization_usage {
let session_file_path = session::storage::get_path(
session::storage::Identifier::Path(session_file.to_path_buf()),
)?;
let mut metadata =
session::storage::read_metadata(&session_file_path)?;
// Update token counts with the summarization usage
// Use output tokens as total since that's what's actually in the context going forward
let summary_tokens = usage.usage.output_tokens.unwrap_or(0);
metadata.total_tokens = Some(summary_tokens);
metadata.input_tokens = None; // Clear input tokens since we now have a summary
metadata.output_tokens = Some(summary_tokens);
metadata.message_count = self.messages.len();
// Update accumulated tokens (add the summarization cost)
let accumulate = |a: Option<i32>, b: Option<i32>| -> Option<i32> {
match (a, b) {
(Some(x), Some(y)) => Some(x + y),
_ => a.or(b),
}
};
metadata.accumulated_total_tokens = accumulate(
metadata.accumulated_total_tokens,
usage.usage.total_tokens,
);
metadata.accumulated_input_tokens = accumulate(
metadata.accumulated_input_tokens,
usage.usage.input_tokens,
);
metadata.accumulated_output_tokens = accumulate(
metadata.accumulated_output_tokens,
usage.usage.output_tokens,
);
session::storage::update_metadata(&session_file_path, &metadata)
.await?;
}
}
output::hide_thinking();