tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: implement async token counter with network resilience and performance optimizations (#3111)

Browse Source 
Co-authored-by: jack <>
Co-authored-by: Salman Mohammed <smohammed@squareup.com>
This commit is contained in:
jack
2025-06-30 14:45:17 +02:00
committed by GitHub
parent fdafbca92e
commit 495cdfb33c
8 changed files with 981 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/context.rs
+14 -7
View File
@@ -1,11 +1,11 @@
use anyhow::Ok;
use crate::message::Message;
use crate::token_counter::TokenCounter;
use crate::token_counter::create_async_token_counter;
use crate::context_mgmt::summarize::summarize_messages;
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};
use crate::context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async};
use super::super::agents::Agent;
@@ -16,9 +16,12 @@ impl Agent {
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 = TokenCounter::new(provider.get_model_config().tokenizer_name());
let token_counter =
create_async_token_counter(provider.get_model_config().tokenizer_name())
.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(&token_counter, messages);
let token_counts = get_messages_token_counts_async(&token_counter, messages);
let (mut new_messages, mut new_token_counts) = truncate_messages(
messages,
@@ -51,11 +54,15 @@ impl Agent {
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 = TokenCounter::new(provider.get_model_config().tokenizer_name());
let token_counter =
create_async_token_counter(provider.get_model_config().tokenizer_name())
.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(provider, messages, &token_counter, target_context_limit).await?;
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.