Wendy Tang
469 lines
16 KiB
Rust
469 lines
16 KiB
Rust
use crate::message::Message;
|
|
use anyhow::{anyhow, Result};
|
|
use mcp_core::Role;
|
|
use std::collections::HashSet;
|
|
use tracing::debug;
|
|
|
|
/// Trait representing a truncation strategy
|
|
pub trait TruncationStrategy {
|
|
/// Determines the indices of messages to remove to fit within the context limit.
|
|
///
|
|
/// - `messages`: The list of messages in the conversation.
|
|
/// - `token_counts`: A parallel array containing the token count for each message.
|
|
/// - `context_limit`: The maximum allowed context length in tokens.
|
|
///
|
|
/// Returns a vector of indices to remove.
|
|
fn determine_indices_to_remove(
|
|
&self,
|
|
messages: &[Message],
|
|
token_counts: &[usize],
|
|
context_limit: usize,
|
|
) -> Result<HashSet<usize>>;
|
|
}
|
|
|
|
/// Strategy to truncate messages by removing the oldest first
|
|
pub struct OldestFirstTruncation;
|
|
/// Strategy to truncate messages explicitly
|
|
pub struct ExplicitTruncation;
|
|
|
|
impl TruncationStrategy for OldestFirstTruncation {
|
|
fn determine_indices_to_remove(
|
|
&self,
|
|
messages: &[Message],
|
|
token_counts: &[usize],
|
|
context_limit: usize,
|
|
) -> Result<HashSet<usize>> {
|
|
let mut indices_to_remove = HashSet::new();
|
|
let mut total_tokens: usize = token_counts.iter().sum();
|
|
let mut tool_ids_to_remove = HashSet::new();
|
|
|
|
for (i, message) in messages.iter().enumerate() {
|
|
if total_tokens <= context_limit {
|
|
break;
|
|
}
|
|
|
|
// Remove the message
|
|
indices_to_remove.insert(i);
|
|
total_tokens -= token_counts[i];
|
|
debug!(
|
|
"OldestFirst: Removing message at index {}. Tokens removed: {}",
|
|
i, token_counts[i]
|
|
);
|
|
|
|
// If it's a ToolRequest or ToolResponse, mark its pair for removal
|
|
if message.is_tool_call() || message.is_tool_response() {
|
|
message.get_tool_ids().iter().for_each(|id| {
|
|
tool_ids_to_remove.insert((i, id.to_string()));
|
|
});
|
|
}
|
|
}
|
|
|
|
// Now, find and remove paired ToolResponses or ToolRequests
|
|
for (i, message) in messages.iter().enumerate() {
|
|
let message_tool_ids = message.get_tool_ids();
|
|
// Find the other part of the pair - same tool_id but different message index
|
|
for (message_idx, tool_id) in &tool_ids_to_remove {
|
|
if message_idx != &i && message_tool_ids.contains(tool_id.as_str()) {
|
|
indices_to_remove.insert(i);
|
|
// No need to check other tool_ids for this message since it's already marked
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(indices_to_remove)
|
|
}
|
|
}
|
|
|
|
/// Truncates the messages to fit within the model's context window.
|
|
/// Mutates the input messages and token counts in place.
|
|
/// Returns an error if it's impossible to truncate the messages within the context limit.
|
|
/// - messages: The vector of messages in the conversation.
|
|
/// - token_counts: A parallel vector containing the token count for each message.
|
|
/// - context_limit: The maximum allowed context length in tokens.
|
|
/// - strategy: The truncation strategy to use. Only option is OldestFirstTruncation.
|
|
pub fn truncate_messages(
|
|
messages: &mut Vec<Message>,
|
|
token_counts: &mut Vec<usize>,
|
|
context_limit: usize,
|
|
strategy: &dyn TruncationStrategy,
|
|
) -> Result<(), anyhow::Error> {
|
|
if messages.len() != token_counts.len() {
|
|
return Err(anyhow!(
|
|
"The vector for messages and token_counts must have same length"
|
|
));
|
|
}
|
|
|
|
// Step 1: Calculate total tokens
|
|
let mut total_tokens: usize = token_counts.iter().sum();
|
|
debug!("Total tokens before truncation: {}", total_tokens);
|
|
|
|
// Check if any individual message is larger than the context limit
|
|
let min_user_msg_tokens = messages
|
|
.iter()
|
|
.zip(token_counts.iter())
|
|
.filter(|(msg, _)| msg.role == Role::User && msg.has_only_text_content())
|
|
.map(|(_, &tokens)| tokens)
|
|
.min();
|
|
|
|
// If there are no valid user messages, or the smallest one is too big for the context
|
|
if min_user_msg_tokens.is_none() || min_user_msg_tokens.unwrap() > context_limit {
|
|
return Err(anyhow!(
|
|
"Not possible to truncate messages within context limit"
|
|
));
|
|
}
|
|
|
|
if total_tokens <= context_limit {
|
|
return Ok(()); // No truncation needed
|
|
}
|
|
|
|
// Step 2: Determine indices to remove based on strategy
|
|
let indices_to_remove =
|
|
strategy.determine_indices_to_remove(messages, token_counts, context_limit)?;
|
|
|
|
// Step 3: Remove the marked messages
|
|
// Vectorize the set and sort in reverse order to avoid shifting indices when removing
|
|
let mut indices_to_remove = indices_to_remove.iter().cloned().collect::<Vec<usize>>();
|
|
indices_to_remove.sort_unstable_by(|a, b| b.cmp(a));
|
|
|
|
for &index in &indices_to_remove {
|
|
if index < messages.len() {
|
|
let _ = messages.remove(index);
|
|
let removed_tokens = token_counts.remove(index);
|
|
total_tokens -= removed_tokens;
|
|
}
|
|
}
|
|
|
|
// Step 4: Ensure the last message is a user message with TextContent only
|
|
while let Some(last_msg) = messages.last() {
|
|
if last_msg.role != Role::User || !last_msg.has_only_text_content() {
|
|
let _ = messages.pop().ok_or(anyhow!("Failed to pop message"))?;
|
|
let removed_tokens = token_counts
|
|
.pop()
|
|
.ok_or(anyhow!("Failed to pop token count"))?;
|
|
total_tokens -= removed_tokens;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Step 5: Check first msg is a User message with TextContent only
|
|
while let Some(first_msg) = messages.first() {
|
|
if first_msg.role != Role::User || !first_msg.has_only_text_content() {
|
|
let _ = messages.remove(0);
|
|
let removed_tokens = token_counts.remove(0);
|
|
total_tokens -= removed_tokens;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
debug!("Total tokens after truncation: {}", total_tokens);
|
|
|
|
// Ensure we have at least one message remaining and it's within context limit
|
|
if messages.is_empty() {
|
|
return Err(anyhow!(
|
|
"Unable to preserve any messages within context limit"
|
|
));
|
|
}
|
|
|
|
if total_tokens > context_limit {
|
|
return Err(anyhow!(
|
|
"Unable to truncate messages within context window."
|
|
));
|
|
}
|
|
|
|
debug!("Truncation complete. Total tokens: {}", total_tokens);
|
|
Ok(())
|
|
}
|
|
|
|
// truncate.rs
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::message::Message;
|
|
use anyhow::Result;
|
|
use mcp_core::content::Content;
|
|
use mcp_core::tool::ToolCall;
|
|
use serde_json::json;
|
|
|
|
// Helper function to create a user text message with a specified token count
|
|
fn user_text(index: usize, tokens: usize) -> (Message, usize) {
|
|
let content = format!("User message {}", index);
|
|
(Message::user().with_text(content), tokens)
|
|
}
|
|
|
|
// Helper function to create an assistant text message with a specified token count
|
|
fn assistant_text(index: usize, tokens: usize) -> (Message, usize) {
|
|
let content = format!("Assistant message {}", index);
|
|
(Message::assistant().with_text(content), tokens)
|
|
}
|
|
|
|
// Helper function to create a tool request message with a specified token count
|
|
fn assistant_tool_request(id: &str, tool_call: ToolCall, tokens: usize) -> (Message, usize) {
|
|
(
|
|
Message::assistant().with_tool_request(id, Ok(tool_call)),
|
|
tokens,
|
|
)
|
|
}
|
|
|
|
// Helper function to create a tool response message with a specified token count
|
|
fn user_tool_response(id: &str, result: Vec<Content>, tokens: usize) -> (Message, usize) {
|
|
(Message::user().with_tool_response(id, Ok(result)), tokens)
|
|
}
|
|
|
|
// Helper function to create messages with alternating user and assistant
|
|
// text messages of a fixed token count
|
|
fn create_messages_with_counts(
|
|
num_pairs: usize,
|
|
tokens: usize,
|
|
remove_last: bool,
|
|
) -> (Vec<Message>, Vec<usize>) {
|
|
let mut messages: Vec<Message> = (0..num_pairs)
|
|
.flat_map(|i| {
|
|
vec![
|
|
user_text(i * 2, tokens).0,
|
|
assistant_text((i * 2) + 1, tokens).0,
|
|
]
|
|
})
|
|
.collect();
|
|
|
|
if remove_last {
|
|
messages.pop();
|
|
}
|
|
|
|
let token_counts = vec![tokens; messages.len()];
|
|
|
|
(messages, token_counts)
|
|
}
|
|
|
|
#[test]
|
|
fn test_oldest_first_no_truncation() -> Result<()> {
|
|
let (messages, token_counts) = create_messages_with_counts(1, 10, false);
|
|
let context_limit = 25;
|
|
|
|
let mut messages_clone = messages.clone();
|
|
let mut token_counts_clone = token_counts.clone();
|
|
truncate_messages(
|
|
&mut messages_clone,
|
|
&mut token_counts_clone,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
assert_eq!(messages_clone, messages);
|
|
assert_eq!(token_counts_clone, token_counts);
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_complex_conversation_with_tools() -> Result<()> {
|
|
// Simulating a real conversation with multiple tool interactions
|
|
let tool_call1 = ToolCall::new("file_read", json!({"path": "/tmp/test.txt"}));
|
|
let tool_call2 = ToolCall::new("database_query", json!({"query": "SELECT * FROM users"}));
|
|
|
|
let messages = vec![
|
|
user_text(1, 15).0, // Initial user query
|
|
assistant_tool_request("tool1", tool_call1.clone(), 20).0,
|
|
user_tool_response(
|
|
"tool1",
|
|
vec![Content::text("File contents".to_string())],
|
|
10,
|
|
)
|
|
.0,
|
|
assistant_text(2, 25).0, // Assistant processes file contents
|
|
user_text(3, 10).0, // User follow-up
|
|
assistant_tool_request("tool2", tool_call2.clone(), 30).0,
|
|
user_tool_response(
|
|
"tool2",
|
|
vec![Content::text("Query results".to_string())],
|
|
20,
|
|
)
|
|
.0,
|
|
assistant_text(4, 35).0, // Assistant analyzes query results
|
|
user_text(5, 5).0, // Final user confirmation
|
|
];
|
|
|
|
let token_counts = vec![15, 20, 10, 25, 10, 30, 20, 35, 5];
|
|
let context_limit = 100; // Force truncation while preserving some tool interactions
|
|
|
|
let mut messages_clone = messages.clone();
|
|
let mut token_counts_clone = token_counts.clone();
|
|
truncate_messages(
|
|
&mut messages_clone,
|
|
&mut token_counts_clone,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
// Verify that tool pairs are kept together and the conversation remains coherent
|
|
assert!(messages_clone.len() >= 3); // At least one complete interaction should remain
|
|
assert!(messages_clone.last().unwrap().role == Role::User); // Last message should be from user
|
|
|
|
// Verify tool pairs are either both present or both removed
|
|
let tool_ids: HashSet<_> = messages_clone
|
|
.iter()
|
|
.flat_map(|m| m.get_tool_ids())
|
|
.collect();
|
|
|
|
// Each tool ID should appear 0 or 2 times (request + response)
|
|
for id in tool_ids {
|
|
let count = messages_clone
|
|
.iter()
|
|
.flat_map(|m| m.get_tool_ids().into_iter())
|
|
.filter(|&tool_id| tool_id == id)
|
|
.count();
|
|
assert!(count == 0 || count == 2, "Tool pair was split: {}", id);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_edge_case_context_window() -> Result<()> {
|
|
// Test case where we're exactly at the context limit
|
|
let (mut messages, mut token_counts) = create_messages_with_counts(2, 25, false);
|
|
let context_limit = 100; // Exactly matches total tokens
|
|
|
|
truncate_messages(
|
|
&mut messages,
|
|
&mut token_counts,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
assert_eq!(messages.len(), 4); // No truncation needed
|
|
assert_eq!(token_counts.iter().sum::<usize>(), 100);
|
|
|
|
// Now add one more token to force truncation
|
|
messages.push(user_text(5, 1).0);
|
|
token_counts.push(1);
|
|
|
|
truncate_messages(
|
|
&mut messages,
|
|
&mut token_counts,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
assert!(token_counts.iter().sum::<usize>() <= context_limit);
|
|
assert!(messages.last().unwrap().role == Role::User);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_tool_chain() -> Result<()> {
|
|
// Simulate a chain of dependent tool calls
|
|
let tool_calls = vec![
|
|
ToolCall::new("git_status", json!({})),
|
|
ToolCall::new("git_diff", json!({"file": "main.rs"})),
|
|
ToolCall::new("git_commit", json!({"message": "Update"})),
|
|
];
|
|
|
|
let mut messages = Vec::new();
|
|
let mut token_counts = Vec::new();
|
|
|
|
// Build a chain of related tool calls
|
|
// 30 tokens each round
|
|
for (i, tool_call) in tool_calls.into_iter().enumerate() {
|
|
let id = format!("git_{}", i);
|
|
messages.push(user_text(i, 10).0);
|
|
token_counts.push(10);
|
|
|
|
messages.push(assistant_tool_request(&id, tool_call, 15).0);
|
|
token_counts.push(20);
|
|
}
|
|
|
|
let context_limit = 50; // Force partial truncation
|
|
let mut messages_clone = messages.clone();
|
|
let mut token_counts_clone = token_counts.clone();
|
|
|
|
truncate_messages(
|
|
&mut messages_clone,
|
|
&mut token_counts_clone,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
// Verify that remaining tool chains are complete
|
|
let remaining_tool_ids: HashSet<_> = messages_clone
|
|
.iter()
|
|
.flat_map(|m| m.get_tool_ids())
|
|
.collect();
|
|
|
|
for _id in remaining_tool_ids {
|
|
// Count request/response pairs
|
|
let requests = messages_clone
|
|
.iter()
|
|
.flat_map(|m| m.get_tool_request_ids().into_iter())
|
|
.count();
|
|
|
|
let responses = messages_clone
|
|
.iter()
|
|
.flat_map(|m| m.get_tool_response_ids().into_iter())
|
|
.count();
|
|
|
|
assert_eq!(requests, 1, "Each remaining tool should have one request");
|
|
assert_eq!(responses, 1, "Each remaining tool should have one response");
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_truncation_with_image_content() -> Result<()> {
|
|
// Create a conversation with image content mixed in
|
|
let mut messages = vec![
|
|
Message::user().with_image("base64_data", "image/png"), // 50 tokens
|
|
Message::assistant().with_text("I see the image"), // 10 tokens
|
|
Message::user().with_text("Can you describe it?"), // 10 tokens
|
|
Message::assistant().with_text("It shows..."), // 20 tokens
|
|
Message::user().with_text("Thanks!"), // 5 tokens
|
|
];
|
|
let mut token_counts = vec![50, 10, 10, 20, 5];
|
|
let context_limit = 45; // Force truncation
|
|
|
|
truncate_messages(
|
|
&mut messages,
|
|
&mut token_counts,
|
|
context_limit,
|
|
&OldestFirstTruncation,
|
|
)?;
|
|
|
|
// Verify the conversation still makes sense
|
|
assert!(messages.len() >= 1);
|
|
assert!(messages.last().unwrap().role == Role::User);
|
|
assert!(token_counts.iter().sum::<usize>() <= context_limit);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_error_cases() -> Result<()> {
|
|
// Test impossibly small context window
|
|
let (mut messages, mut token_counts) = create_messages_with_counts(1, 10, false);
|
|
let result = truncate_messages(
|
|
&mut messages,
|
|
&mut token_counts,
|
|
5, // Impossibly small context
|
|
&OldestFirstTruncation,
|
|
);
|
|
assert!(result.is_err());
|
|
|
|
// Test unmatched token counts
|
|
let mut messages = vec![user_text(1, 10).0];
|
|
let mut token_counts = vec![10, 10]; // Mismatched length
|
|
let result = truncate_messages(
|
|
&mut messages,
|
|
&mut token_counts,
|
|
100,
|
|
&OldestFirstTruncation,
|
|
);
|
|
assert!(result.is_err());
|
|
|
|
Ok(())
|
|
}
|
|
}
|