diff --git a/Cargo.lock b/Cargo.lock index a962076cdd..d47c135f88 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2668,20 +2668,6 @@ dependencies = [ "parking_lot_core", ] -[[package]] -name = "dashmap" -version = "6.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" -dependencies = [ - "cfg-if", - "crossbeam-utils", - "hashbrown 0.14.5", - "lock_api", - "once_cell", - "parking_lot_core", -] - [[package]] name = "data-encoding" version = "2.11.0" @@ -4449,7 +4435,6 @@ version = "1.37.0" dependencies = [ "agent-client-protocol", "agent-client-protocol-schema", - "ahash", "anyhow", "arboard", "async-stream", @@ -4469,7 +4454,6 @@ dependencies = [ "chrono", "clap", "ctor", - "dashmap 6.2.1", "dirs 5.0.1", "dotenvy", "dtor", @@ -6313,7 +6297,7 @@ dependencies = [ "async-trait", "boxed_error", "capacity_builder", - "dashmap 5.5.3", + "dashmap", "deno_error", "deno_maybe_sync", "deno_media_type", diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index fccd7dbbc6..c59286e09e 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -152,8 +152,6 @@ blake3 = { version = "1", default-features = false, features = ["std"] } fs2 = { workspace = true } tokio-stream = { workspace = true, features = ["io-util"] } tempfile = { workspace = true } -dashmap = { version = "6", default-features = false } -ahash = { version = "0.8.11", default-features = false, features = ["std"] } tokio-util = { workspace = true, features = ["compat"] } agent-client-protocol-schema = { workspace = true } agent-client-protocol = { workspace = true, features = ["unstable"] } diff --git a/crates/goose/src/token_counter.rs b/crates/goose/src/token_counter.rs index cb41621427..6b1f8bb5ff 100644 --- a/crates/goose/src/token_counter.rs +++ b/crates/goose/src/token_counter.rs @@ -1,8 +1,7 @@ -use ahash::AHasher; -use dashmap::DashMap; +use lru::LruCache; use rmcp::model::Tool; -use std::hash::{Hash, Hasher}; -use std::sync::Arc; +use std::num::NonZeroUsize; +use std::sync::{Arc, Mutex}; use tiktoken_rs::CoreBPE; use tokio::sync::OnceCell; @@ -10,7 +9,7 @@ use crate::conversation::message::Message; static TOKENIZER: OnceCell> = OnceCell::const_new(); -const MAX_TOKEN_CACHE_SIZE: usize = 10_000; +const MAX_TOKEN_CACHE_SIZE: usize = 1_024; // token use for various bits of a tool calls: const FUNC_INIT: usize = 7; @@ -22,38 +21,54 @@ const FUNC_END: usize = 12; pub struct TokenCounter { tokenizer: Arc, - token_cache: Arc>, + token_cache: Mutex>, +} + +#[derive(Clone, Copy, Eq, Hash, PartialEq)] +struct TokenCacheKey { + len: usize, + hash: [u8; 32], +} + +impl TokenCacheKey { + fn from_text(text: &str) -> Self { + Self { + len: text.len(), + hash: *blake3::hash(text.as_bytes()).as_bytes(), + } + } } impl TokenCounter { pub async fn new() -> Result { let tokenizer = get_tokenizer().await?; + let cache_capacity = + NonZeroUsize::new(MAX_TOKEN_CACHE_SIZE).expect("token cache capacity must be non-zero"); Ok(Self { tokenizer, - token_cache: Arc::new(DashMap::new()), + token_cache: Mutex::new(LruCache::new(cache_capacity)), }) } pub fn count_tokens(&self, text: &str) -> usize { - let mut hasher = AHasher::default(); - text.hash(&mut hasher); - let hash = hasher.finish(); - - if let Some(count) = self.token_cache.get(&hash) { - return *count; + let cache_key = TokenCacheKey::from_text(text); + if let Some(count) = self + .token_cache + .lock() + .expect("token cache mutex poisoned") + .get(&cache_key) + .copied() + { + return count; } let tokens = self.tokenizer.encode_with_special_tokens(text); let count = tokens.len(); - if self.token_cache.len() >= MAX_TOKEN_CACHE_SIZE { - if let Some(entry) = self.token_cache.iter().next() { - let old_hash = *entry.key(); - self.token_cache.remove(&old_hash); - } - } - - self.token_cache.insert(hash, count); + self.token_cache + .lock() + .expect("token cache mutex poisoned") + .put(cache_key, count); count } @@ -173,11 +188,17 @@ impl TokenCounter { } pub fn clear_cache(&self) { - self.token_cache.clear(); + self.token_cache + .lock() + .expect("token cache mutex poisoned") + .clear(); } pub fn cache_size(&self) -> usize { - self.token_cache.len() + self.token_cache + .lock() + .expect("token cache mutex poisoned") + .len() } } @@ -260,13 +281,13 @@ mod tests { let counter = create_token_counter().await.unwrap(); let mut cached_texts = Vec::new(); - for i in 0..50 { + for i in 0..=MAX_TOKEN_CACHE_SIZE { let text = format!("Test string number {}", i); counter.count_tokens(&text); cached_texts.push(text); } - assert!(counter.cache_size() <= MAX_TOKEN_CACHE_SIZE); + assert_eq!(counter.cache_size(), MAX_TOKEN_CACHE_SIZE); let recent_text = &cached_texts[cached_texts.len() - 1]; let start_size = counter.cache_size();