tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 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/context_mgmt/auto_compact.rs
+23 -113
View File
@@ -1,12 +1,7 @@
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::{
agents::Agent,
config::Config,
context_mgmt::{
common::{SYSTEM_PROMPT_TOKEN_OVERHEAD, TOOLS_TOKEN_OVERHEAD},
get_messages_token_counts_async,
},
agents::Agent, config::Config, context_mgmt::get_messages_token_counts_async,
token_counter::create_async_token_counter,
};
use anyhow::Result;
@@ -19,10 +14,9 @@ pub struct AutoCompactResult {
pub compacted: bool,
/// The messages after potential compaction
pub messages: Conversation,
/// Token count before compaction (if compaction occurred)
pub tokens_before: Option<usize>,
/// Token count after compaction (if compaction occurred)
pub tokens_after: Option<usize>,
/// Provider usage from summarization (if compaction occurred)
/// This contains the actual token counts after compaction
pub summarization_usage: Option<crate::providers::base::ProviderUsage>,
}
/// Result of checking if compaction is needed
@@ -126,47 +120,6 @@ pub async fn check_compaction_needed(
})
}
/// Perform compaction on messages
///
/// This function performs the actual compaction using the agent's summarization
/// capabilities. It assumes compaction is needed and should be called after
/// `check_compaction_needed` confirms it's necessary.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history to compact
///
/// # Returns
/// * Tuple of (compacted_messages, tokens_before, tokens_after)
pub async fn perform_compaction(
agent: &Agent,
messages: &[Message],
) -> Result<(Conversation, usize, usize)> {
// Get token counter to measure before/after
let token_counter = create_async_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
// Calculate tokens before compaction
let token_counts_before = get_messages_token_counts_async(&token_counter, messages);
let tokens_before: usize = token_counts_before.iter().sum();
info!("Performing compaction on {} tokens", tokens_before);
// Perform compaction
let (compacted_messages, compacted_token_counts) = agent.summarize_context(messages).await?;
let tokens_after: usize = compacted_token_counts.iter().sum();
info!(
"Compaction complete: {} tokens -> {} tokens ({:.1}% reduction)",
tokens_before,
tokens_after,
(1.0 - (tokens_after as f64 / tokens_before as f64)) * 100.0
);
Ok((compacted_messages, tokens_before, tokens_after))
}
/// Check if messages need compaction and compact them if necessary
///
/// This is a convenience wrapper function that combines checking and compaction.
@@ -201,8 +154,7 @@ pub async fn check_and_compact_messages(
return Ok(AutoCompactResult {
compacted: false,
messages: Conversation::new_unvalidated(messages.to_vec()),
tokens_before: None,
tokens_after: None,
summarization_usage: None,
});
}
@@ -225,8 +177,8 @@ pub async fn check_and_compact_messages(
};
// Perform the compaction on messages excluding the preserved user message
let (mut compacted_messages, tokens_before, tokens_after) =
perform_compaction(agent, messages_to_compact).await?;
let (mut compacted_messages, _, summarization_usage) =
agent.summarize_context(messages_to_compact).await?;
// Add back the preserved user message if it exists
if let Some(user_message) = preserved_user_message {
@@ -236,8 +188,7 @@ pub async fn check_and_compact_messages(
Ok(AutoCompactResult {
compacted: true,
messages: compacted_messages,
tokens_before: Some(tokens_before + SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD),
tokens_after: Some(tokens_after + SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD),
summarization_usage,
})
}
@@ -326,7 +277,7 @@ mod tests {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(100_000.into()),
.with_context_limit(Some(100_000)),
});
let agent = Agent::new();
@@ -352,7 +303,7 @@ mod tests {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(100_000.into()),
.with_context_limit(Some(100_000)),
});
let agent = Agent::new();
@@ -375,40 +326,12 @@ mod tests {
assert!(!result.needs_compaction);
}
#[tokio::test]
async fn test_perform_compaction() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(50_000.into()),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create some messages to compact
let messages = vec![
create_test_message("First message"),
create_test_message("Second message"),
create_test_message("Third message"),
];
let (compacted_messages, tokens_before, tokens_after) =
perform_compaction(&agent, &messages).await.unwrap();
assert!(tokens_before > 0);
assert!(tokens_after > 0);
// Note: The mock provider returns a fixed summary, which might not always be smaller
// In real usage, compaction should reduce tokens, but for testing we just verify it works
assert!(!compacted_messages.is_empty());
}
#[tokio::test]
async fn test_auto_compact_disabled() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(10_000.into()),
.with_context_limit(Some(10_000)),
});
let agent = Agent::new();
@@ -423,8 +346,7 @@ mod tests {
assert!(!result.compacted);
assert_eq!(result.messages.len(), messages.len());
assert!(result.tokens_before.is_none());
assert!(result.tokens_after.is_none());
assert!(result.summarization_usage.is_none());
// Test with threshold 1.0 (disabled)
let result = check_and_compact_messages(&agent, &messages, Some(1.0), None)
@@ -439,7 +361,7 @@ mod tests {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(100_000.into()), // Increased to ensure overhead doesn't dominate
.with_context_limit(Some(100_000)), // Increased to ensure overhead doesn't dominate
});
let agent = Agent::new();
@@ -499,14 +421,15 @@ mod tests {
}
assert!(result.compacted);
assert!(result.tokens_before.is_some());
assert!(result.tokens_after.is_some());
assert!(result.summarization_usage.is_some());
// Should have fewer tokens after compaction
if let (Some(before), Some(after)) = (result.tokens_before, result.tokens_after) {
// Verify that summarization usage contains token counts
if let Some(usage) = &result.summarization_usage {
assert!(usage.usage.total_tokens.is_some());
let after = usage.usage.total_tokens.unwrap_or(0) as usize;
assert!(
after < before,
"Token count should decrease after compaction"
after > 0,
"Token count after compaction should be greater than 0"
);
}
@@ -519,7 +442,7 @@ mod tests {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(30_000.into()), // Smaller context limit to make threshold easier to hit
.with_context_limit(Some(30_000)), // Smaller context limit to make threshold easier to hit
});
let agent = Agent::new();
@@ -553,19 +476,7 @@ mod tests {
// Debug info if not compacted
if !result.compacted {
let provider = agent.provider().await.unwrap();
let token_counter = create_async_token_counter().await.unwrap();
let token_counts = get_messages_token_counts_async(&token_counter, &messages);
let total_tokens: usize = token_counts.iter().sum();
let context_limit = provider.get_model_config().context_limit();
let usage_ratio = total_tokens as f64 / context_limit as f64;
eprintln!(
"Config test not compacted - tokens: {} / {} ({:.1}%)",
total_tokens,
context_limit,
usage_ratio * 100.0
);
eprintln!("Test failed - compaction not triggered");
}
// With such a low threshold (10%), it should compact
@@ -701,8 +612,7 @@ mod tests {
// Should have triggered compaction
assert!(result.compacted);
assert!(result.tokens_before.is_some());
assert!(result.tokens_after.is_some());
assert!(result.summarization_usage.is_some());
// Verify the compacted messages are returned
assert!(!result.messages.is_empty());
crates/goose/src/context_mgmt/summarize.rs
+44 -553
View File
@@ -1,70 +1,26 @@
use super::common::get_messages_token_counts_async;
use crate::context_mgmt::get_messages_token_counts;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::prompt_template::render_global_file;
use crate::providers::base::Provider;
use crate::token_counter::{AsyncTokenCounter, TokenCounter};
use anyhow::Result;
use rmcp::model::Role;
use serde::Serialize;
use std::sync::Arc;
// Constants for the summarization prompt and a follow-up user message.
const SUMMARY_PROMPT: &str = "You are good at summarizing conversations";
#[derive(Serialize)]
struct SummarizeContext {
messages: String,
}
/// Summarize the combined messages from the accumulated summary and the current chunk.
///
/// This method builds the summarization request, sends it to the provider, and returns the summarized response.
async fn summarize_combined_messages(
provider: &Arc<dyn Provider>,
accumulated_summary: &[Message],
current_chunk: &[Message],
) -> Result<Conversation, anyhow::Error> {
// Combine the accumulated summary and current chunk into a single batch.
let combined_messages = Conversation::new_unvalidated(
accumulated_summary
.iter()
.cloned()
.chain(current_chunk.iter().cloned())
.collect::<Vec<_>>(),
);
use crate::providers::base::ProviderUsage;
// Format the batch as a summarization request.
let request_text = format!(
"Please summarize the following conversation history, preserving the key points. This summarization will be used for the later conversations.\n\n```\n{:?}\n```",
combined_messages
);
let summarization_request = vec![Message::user().with_text(&request_text)];
// Send the request to the provider and fetch the response.
let mut response = provider
.complete(SUMMARY_PROMPT, &summarization_request, &[])
.await?
.0;
// Set role to user as it will be used in following conversation as user content.
response.role = Role::User;
// Return the summary as the new accumulated summary.
Ok(Conversation::new_unvalidated(vec![response]))
}
// Summarization steps:
// Using a single tailored prompt, summarize the entire conversation history.
pub async fn summarize_messages_oneshot(
/// Summarization function that uses the detailed prompt from the markdown template
pub async fn summarize_messages(
provider: Arc<dyn Provider>,
messages: &[Message],
token_counter: &TokenCounter,
_context_limit: usize,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
) -> Result<Option<(Message, ProviderUsage)>, anyhow::Error> {
if messages.is_empty() {
// If no messages to summarize, return empty
return Ok((Conversation::empty(), vec![]));
return Ok(None);
}
// Format all messages as a single string for the summarization prompt
@@ -86,176 +42,21 @@ pub async fn summarize_messages_oneshot(
.with_text("Please summarize the conversation history provided in the system prompt.");
let summarization_request = vec![user_message];
// Send the request to the provider and fetch the response.
let mut response = provider
// Send the request to the provider and fetch the response
let (mut response, mut provider_usage) = provider
.complete(&system_prompt, &summarization_request, &[])
.await?
.0;
.await?;
// Set role to user as it will be used in following conversation as user content.
// Set role to user as it will be used in following conversation as user content
response.role = Role::User;
// Return just the summary without any tool response preservation
let final_summary = Conversation::new_unvalidated([response].into_iter());
let counts = get_messages_token_counts(token_counter, final_summary.messages());
Ok((final_summary, counts))
}
// Summarization steps:
// 1. Break down large text into smaller chunks (roughly 30% of the model’s context window).
// 2. For each chunk:
// a. Combine it with the previous summary (or leave blank for the first iteration).
// b. Summarize the combined text, focusing on extracting only the information we need.
// 3. Generate a final summary using a tailored prompt.
pub async fn summarize_messages_chunked(
provider: Arc<dyn Provider>,
messages: &[Message],
token_counter: &TokenCounter,
context_limit: usize,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
let chunk_size = context_limit / 3; // 33% of the context window.
let summary_prompt_tokens = token_counter.count_tokens(SUMMARY_PROMPT);
let mut accumulated_summary = Conversation::empty();
// Get token counts for each message.
let token_counts = get_messages_token_counts(token_counter, messages);
// Tokenize and break messages into chunks.
let mut current_chunk: Vec<Message> = Vec::new();
let mut current_chunk_tokens = 0;
for (message, message_tokens) in messages.iter().zip(token_counts.iter()) {
if current_chunk_tokens + message_tokens > chunk_size - summary_prompt_tokens {
// Summarize the current chunk with the accumulated summary.
accumulated_summary = summarize_combined_messages(
&provider,
accumulated_summary.messages(),
&current_chunk,
)
.await?;
// Reset for the next chunk.
current_chunk.clear();
current_chunk_tokens = 0;
}
// Add message to the current chunk.
current_chunk.push(message.clone());
current_chunk_tokens += message_tokens;
}
// Summarize the final chunk if it exists.
if !current_chunk.is_empty() {
accumulated_summary =
summarize_combined_messages(&provider, accumulated_summary.messages(), &current_chunk)
.await?;
}
// Return just the summary without any tool response preservation
Ok((
accumulated_summary.clone(),
get_messages_token_counts(token_counter, accumulated_summary.messages()),
))
}
/// Main summarization function that chooses the best algorithm based on context size.
///
/// This function will:
/// 1. First try the one-shot summarization if there's enough context window available
/// 2. Fall back to the chunked approach if the one-shot fails or if context is too limited
/// 3. Choose the algorithm based on absolute token requirements rather than percentages
pub async fn summarize_messages(
provider: Arc<dyn Provider>,
messages: &[Message],
token_counter: &TokenCounter,
context_limit: usize,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
// Calculate total tokens in messages
let total_tokens: usize = get_messages_token_counts(token_counter, messages)
.iter()
.sum();
// Calculate absolute token requirements (future-proof for large context models)
let system_prompt_overhead = 1000; // Conservative estimate for the summarization prompt
let response_overhead = 4000; // Generous buffer for response generation
let safety_buffer = 1000; // Small safety margin for tokenization variations
let total_required = total_tokens + system_prompt_overhead + response_overhead + safety_buffer;
// Use one-shot if we have enough absolute space (no percentage-based limits)
if total_required <= context_limit {
match summarize_messages_oneshot(
Arc::clone(&provider),
messages,
token_counter,
context_limit,
)
// Ensure we have token counts, estimating if necessary
provider_usage
.ensure_tokens(&system_prompt, &summarization_request, &response, &[])
.await
{
Ok(result) => return Ok(result),
Err(e) => {
// Log the error but continue to fallback
tracing::warn!(
"One-shot summarization failed, falling back to chunked approach: {}",
e
);
}
}
}
.map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
// Fall back to the chunked approach
summarize_messages_chunked(provider, messages, token_counter, context_limit).await
}
/// Async version using AsyncTokenCounter for better performance
pub async fn summarize_messages_async(
provider: Arc<dyn Provider>,
messages: &[Message],
token_counter: &AsyncTokenCounter,
context_limit: usize,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
let chunk_size = context_limit / 3; // 33% of the context window.
let summary_prompt_tokens = token_counter.count_tokens(SUMMARY_PROMPT);
let mut accumulated_summary = Conversation::empty();
// Get token counts for each message.
let token_counts = get_messages_token_counts_async(token_counter, messages);
// Tokenize and break messages into chunks.
let mut current_chunk = Vec::new();
let mut current_chunk_tokens = 0;
for (message, message_tokens) in messages.iter().zip(token_counts.iter()) {
if current_chunk_tokens + message_tokens > chunk_size - summary_prompt_tokens {
// Summarize the current chunk with the accumulated summary.
accumulated_summary = summarize_combined_messages(
&provider,
accumulated_summary.messages(),
&current_chunk,
)
.await?;
// Reset for the next chunk.
current_chunk.clear();
current_chunk_tokens = 0;
}
// Add message to the current chunk.
current_chunk.push(message.clone());
current_chunk_tokens += message_tokens;
}
// Summarize the final chunk if it exists.
if !current_chunk.is_empty() {
accumulated_summary =
summarize_combined_messages(&provider, accumulated_summary.messages(), &current_chunk)
.await?;
}
let count = get_messages_token_counts_async(token_counter, accumulated_summary.messages());
// Return just the summary without any tool response preservation
Ok((accumulated_summary.clone(), count))
Ok(Some((response, provider_usage)))
}
#[cfg(test)]
@@ -303,13 +104,20 @@ mod tests {
.no_annotation(),
)],
),
ProviderUsage::new("mock".to_string(), Usage::default()),
ProviderUsage::new(
"mock".to_string(),
Usage {
input_tokens: Some(100),
output_tokens: Some(50),
total_tokens: Some(150),
},
),
))
}
}
fn create_mock_provider() -> Result<Arc<dyn Provider>> {
let mock_model_config = ModelConfig::new("test-model")?.with_context_limit(200_000.into());
let mock_model_config = ModelConfig::new("test-model")?.with_context_limit(Some(200_000));
Ok(Arc::new(MockProvider {
model_config: mock_model_config,
@@ -329,366 +137,49 @@ mod tests {
}
#[tokio::test]
async fn test_summarize_messages_single_chunk() {
async fn test_summarize_messages_basic() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 10_000; // Higher limit to avoid underflow
let messages = create_test_messages();
let result = summarize_messages(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
let result = summarize_messages(Arc::clone(&provider), &messages).await;
assert!(result.is_ok(), "The function should return Ok.");
let (summarized_messages, token_counts) = result.unwrap();
let summary_result = result.unwrap();
assert!(
summary_result.is_some(),
"The summary should contain a result."
);
let (summarized_message, provider_usage) = summary_result.unwrap();
assert_eq!(
summarized_messages.len(),
1,
"The summary should contain one message."
);
assert_eq!(
summarized_messages.first().unwrap().role,
summarized_message.role,
Role::User,
"The summarized message should be from the user."
);
assert_eq!(
token_counts.len(),
1,
"Token counts should match the number of summarized messages."
assert!(
provider_usage.usage.input_tokens.unwrap_or(0) > 0,
"Should have input token count"
);
}
#[tokio::test]
async fn test_summarize_messages_multiple_chunks() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 10_000; // Higher limit to avoid underflow
let messages = create_test_messages();
let result = summarize_messages(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
assert!(result.is_ok(), "The function should return Ok.");
let (summarized_messages, token_counts) = result.unwrap();
assert_eq!(
summarized_messages.len(),
1,
"There should be one final summarized message."
);
assert_eq!(
summarized_messages.first().unwrap().role,
Role::User,
"The summarized message should be from the user."
);
assert_eq!(
token_counts.len(),
1,
"Token counts should match the number of summarized messages."
assert!(
provider_usage.usage.output_tokens.unwrap_or(0) > 0,
"Should have output token count"
);
}
#[tokio::test]
async fn test_summarize_messages_empty_input() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 10_000; // Higher limit to avoid underflow
let messages: Vec<Message> = Vec::new();
let result = summarize_messages(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
let result = summarize_messages(Arc::clone(&provider), &messages).await;
assert!(result.is_ok(), "The function should return Ok.");
let (summarized_messages, token_counts) = result.unwrap();
assert_eq!(
summarized_messages.len(),
0,
"The summary should be empty for an empty input."
);
assert!(
token_counts.is_empty(),
"Token counts should be empty for an empty input."
);
}
#[tokio::test]
async fn test_summarize_messages_uses_oneshot_for_small_context() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 100_000; // Large context limit
let messages = create_test_messages(); // Small message set
let result = summarize_messages(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
assert!(result.is_ok(), "The function should return Ok.");
let (summarized_messages, _) = result.unwrap();
// Should use one-shot and return a single summarized message
assert_eq!(
summarized_messages.len(),
1,
"Should use one-shot summarization for small context."
);
}
#[tokio::test]
async fn test_summarize_messages_uses_chunked_for_large_context() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 10_000; // Higher limit to avoid underflow
let messages = create_test_messages();
let result = summarize_messages(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
assert!(result.is_ok(), "The function should return Ok.");
let (summarized_messages, _) = result.unwrap();
// Should fall back to chunked approach
assert_eq!(
summarized_messages.len(),
1,
"Should use chunked summarization for large context."
);
}
// Mock provider that fails on one-shot but succeeds on chunked
#[derive(Clone)]
struct FailingOneshotProvider {
model_config: ModelConfig,
call_count: Arc<std::sync::Mutex<usize>>,
}
#[async_trait::async_trait]
impl Provider for FailingOneshotProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete(
&self,
system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let mut count = self.call_count.lock().unwrap();
*count += 1;
// Fail if this looks like a one-shot request
if system.contains("reasoning in `<analysis>` tags") {
return Err(ProviderError::RateLimitExceeded(
"Simulated one-shot failure".to_string(),
));
}
// Succeed for chunked requests (uses the old SUMMARY_PROMPT)
Ok((
Message::new(
Role::Assistant,
Utc::now().timestamp(),
vec![MessageContent::Text(
RawTextContent {
text: "Chunked summary".to_string(),
}
.no_annotation(),
)],
),
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
}
#[tokio::test]
async fn test_summarize_messages_fallback_on_oneshot_failure() {
let call_count = Arc::new(std::sync::Mutex::new(0));
let provider = Arc::new(FailingOneshotProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(200_000.into()),
call_count: Arc::clone(&call_count),
});
let token_counter = TokenCounter::new();
let context_limit = 100_000; // Large enough to try one-shot first
let messages = create_test_messages();
let result = summarize_messages(provider, &messages, &token_counter, context_limit).await;
let summary_result = result.unwrap();
assert!(
result.is_ok(),
"The function should return Ok after fallback."
summary_result.is_none(),
"The summary should be None for empty input."
);
let (summarized_messages, _) = result.unwrap();
// Should have fallen back to chunked approach
assert_eq!(
summarized_messages.len(),
1,
"Should successfully fall back to chunked approach."
);
// Verify the content comes from the chunked approach
if let MessageContent::Text(text_content) = &summarized_messages.first().unwrap().content[0]
{
assert_eq!(text_content.text, "Chunked summary");
} else {
panic!("Expected text content");
}
// Should have made multiple calls (one-shot attempt + chunked calls)
let final_count = *call_count.lock().unwrap();
assert!(
final_count > 1,
"Should have made multiple provider calls during fallback"
);
}
#[tokio::test]
async fn test_summarize_messages_oneshot_direct_call() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 100_000;
let messages = create_test_messages();
let result = summarize_messages_oneshot(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
assert!(
result.is_ok(),
"One-shot summarization should work directly."
);
let (summarized_messages, token_counts) = result.unwrap();
assert_eq!(
summarized_messages.len(),
1,
"One-shot should return a single summary message."
);
assert_eq!(
summarized_messages.first().unwrap().role,
Role::User,
"Summary should be from user role for context."
);
assert_eq!(
token_counts.len(),
1,
"Should have token count for the summary."
);
}
#[tokio::test]
async fn test_summarize_messages_chunked_direct_call() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
let context_limit = 10_000; // Higher limit to avoid underflow
let messages = create_test_messages();
let result = summarize_messages_chunked(
Arc::clone(&provider),
&messages,
&token_counter,
context_limit,
)
.await;
assert!(
result.is_ok(),
"Chunked summarization should work directly."
);
let (summarized_messages, token_counts) = result.unwrap();
assert_eq!(
summarized_messages.len(),
1,
"Chunked should return a single final summary."
);
assert_eq!(
summarized_messages.first().unwrap().role,
Role::User,
"Summary should be from user role for context."
);
assert_eq!(
token_counts.len(),
1,
"Should have token count for the summary."
);
}
#[tokio::test]
async fn test_absolute_token_threshold_calculation() {
let provider = create_mock_provider().expect("failed to create mock provider");
let token_counter = TokenCounter::new();
// Test with a context limit where absolute token calculation matters
let context_limit = 10_000;
let system_prompt_overhead = 1000;
let response_overhead = 4000;
let safety_buffer = 1000;
let max_message_tokens =
context_limit - system_prompt_overhead - response_overhead - safety_buffer; // 4000 tokens
// Create messages that are just under the absolute threshold
let mut large_messages = Vec::new();
let base_message = set_up_text_message("x".repeat(50).as_str(), Role::User);
// Add enough messages to approach but not exceed the absolute threshold
let message_tokens = token_counter.count_tokens(&format!("{:?}", base_message));
let num_messages = (max_message_tokens / message_tokens).saturating_sub(1);
for i in 0..num_messages {
large_messages.push(set_up_text_message(&format!("Message {}", i), Role::User));
}
let result = summarize_messages(
Arc::clone(&provider),
&large_messages,
&token_counter,
context_limit,
)
.await;
assert!(
result.is_ok(),
"Should handle absolute threshold calculation correctly."
);
let (summarized_messages, _) = result.unwrap();
assert_eq!(summarized_messages.len(), 1, "Should produce a summary.");
}
}