Salman Mohammed
88 lines
4.3 KiB
Rust
88 lines
4.3 KiB
Rust
use anyhow::Ok;
|
|
|
|
use crate::message::Message;
|
|
use crate::token_counter::create_async_token_counter;
|
|
|
|
use crate::context_mgmt::summarize::summarize_messages_async;
|
|
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<(Vec<Message>, 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.
|
|
pub async fn summarize_context(
|
|
&self,
|
|
messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded
|
|
) -> Result<(Vec<Message>, 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.clone());
|
|
|
|
let (mut new_messages, mut new_token_counts) =
|
|
summarize_messages_async(provider, messages, &token_counter, target_context_limit)
|
|
.await?;
|
|
|
|
// 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.",
|
|
);
|
|
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);
|
|
}
|
|
}
|
|
|
|
Ok((new_messages, new_token_counts))
|
|
}
|
|
}
|