tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 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/src/agents/context.rs
+37 -24
View File
@@ -4,7 +4,7 @@ use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::token_counter::create_async_token_counter;
use crate::context_mgmt::summarize::summarize_messages_async;
use crate::context_mgmt::summarize::summarize_messages;
use crate::context_mgmt::truncate::{truncate_messages, OldestFirstTruncation};
use crate::context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async};
@@ -49,40 +49,53 @@ impl Agent {
}
/// Public API to summarize the conversation so that its token count is within the allowed context limit.
/// Returns the summarized messages, token counts, and the ProviderUsage from summarization
pub async fn summarize_context(
&self,
messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
) -> Result<
(
Conversation,
Vec<usize>,
Option<crate::providers::base::ProviderUsage>,
),
anyhow::Error,
> {
let provider = self.provider().await?;
let token_counter = create_async_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
let target_context_limit = estimate_target_context_limit(provider.clone());
let summary_result = summarize_messages(provider.clone(), messages).await?;
let (mut new_messages, mut new_token_counts) =
summarize_messages_async(provider, messages, &token_counter, target_context_limit)
.await?;
let (mut new_messages, mut new_token_counts, summarization_usage) = match summary_result {
Some((summary_message, provider_usage)) => {
// For token counting purposes, we use the output tokens (the actual summary content)
// since that's what will be in the context going forward
let total_tokens = provider_usage.usage.output_tokens.unwrap_or(0) as usize;
(
vec![summary_message],
vec![total_tokens],
Some(provider_usage),
)
}
None => {
// No summary was generated (empty input)
tracing::warn!("Summarization failed. Returning empty messages.");
return Ok((Conversation::empty(), vec![], None));
}
};
// If the summarized messages only contains one message, it means no tool request and response message in the summarized messages,
// Add an assistant message to the summarized messages to ensure the assistant's response is included in the context.
if new_messages.len() == 1 {
let assistant_message = Message::assistant().with_text(
"I had run into a context length exceeded error so I summarized our conversation.",
"I ran into a context length exceeded error so I summarized our conversation.",
);
let assistant_tokens =
token_counter.count_chat_tokens("", &[assistant_message.clone()], &[]);
let current_total: usize = new_token_counts.iter().sum();
if current_total + assistant_tokens <= target_context_limit {
new_messages.push(assistant_message);
new_token_counts.push(assistant_tokens);
} else {
// If we can't fit the assistant message, at least log what happened
tracing::warn!("Cannot add summarization notice message due to context limits. Current: {}, Assistant: {}, Limit: {}",
current_total, assistant_tokens, target_context_limit);
}
let assistant_message_tokens: usize = 14;
new_messages.push(assistant_message);
new_token_counts.push(assistant_message_tokens);
}
Ok((new_messages, new_token_counts))
Ok((
Conversation::new_unvalidated(new_messages),
new_token_counts,
summarization_usage,
))
}
}