David Katz
102 lines
4.5 KiB
Rust
102 lines
4.5 KiB
Rust
use anyhow::Ok;
|
|
|
|
use crate::conversation::message::Message;
|
|
use crate::conversation::Conversation;
|
|
use crate::token_counter::create_async_token_counter;
|
|
|
|
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};
|
|
|
|
use super::super::agents::Agent;
|
|
|
|
impl Agent {
|
|
/// Public API to truncate oldest messages so that the conversation's token count is within the allowed context limit.
|
|
pub async fn truncate_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> {
|
|
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);
|
|
let token_counts = get_messages_token_counts_async(&token_counter, messages);
|
|
|
|
let (mut new_messages, mut new_token_counts) = truncate_messages(
|
|
messages,
|
|
&token_counts,
|
|
target_context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
// Only add an assistant message if we have room for it and it won't cause another overflow
|
|
let assistant_message = Message::assistant().with_text("I had run into a context length exceeded error so I truncated some of the oldest messages in 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 truncation notice message due to context limits. Current: {}, Assistant: {}, Limit: {}",
|
|
current_total, assistant_tokens, target_context_limit);
|
|
}
|
|
|
|
Ok((new_messages, new_token_counts))
|
|
}
|
|
|
|
/// 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>,
|
|
Option<crate::providers::base::ProviderUsage>,
|
|
),
|
|
anyhow::Error,
|
|
> {
|
|
let provider = self.provider().await?;
|
|
let summary_result = summarize_messages(provider.clone(), messages).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));
|
|
}
|
|
};
|
|
|
|
// 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 ran into a context length exceeded error so I summarized our conversation.",
|
|
);
|
|
let assistant_message_tokens: usize = 14;
|
|
new_messages.push(assistant_message);
|
|
new_token_counts.push(assistant_message_tokens);
|
|
}
|
|
|
|
Ok((
|
|
Conversation::new_unvalidated(new_messages),
|
|
new_token_counts,
|
|
summarization_usage,
|
|
))
|
|
}
|
|
}
|