use crate::types::*;
use std::sync::Arc;
pub trait TokenCounter: Send + Sync {
fn estimate_text(&self, text: &str) -> usize;
fn estimate_content(&self, content: &[Content]) -> usize {
content
.iter()
.map(|c| match c {
Content::Text { text } => self.estimate_text(text),
Content::Image { data, .. } => {
let raw_bytes = data.len() * 3 / 4;
(raw_bytes / 750).clamp(85, 16_000)
}
Content::Thinking { thinking, .. } => self.estimate_text(thinking),
Content::ToolCall {
name, arguments, ..
} => self.estimate_text(name) + self.estimate_text(&arguments.to_string()) + 8,
})
.sum()
}
fn estimate_message(&self, msg: &AgentMessage) -> usize {
match msg {
AgentMessage::Llm(lm) => match &lm.message {
Message::User { content, .. } => self.estimate_content(content) + 4,
Message::Assistant { content, .. } => self.estimate_content(content) + 4,
Message::ToolResult {
content, tool_name, ..
} => self.estimate_content(content) + self.estimate_text(tool_name) + 8,
},
AgentMessage::Extension(ext) => self.estimate_text(&ext.data.to_string()) + 4,
}
}
fn estimate_messages(&self, msgs: &[AgentMessage]) -> usize {
msgs.iter().map(|m| self.estimate_message(m)).sum()
}
}
pub struct HeuristicTokenCounter;
impl TokenCounter for HeuristicTokenCounter {
fn estimate_text(&self, text: &str) -> usize {
text.len().div_ceil(4)
}
}
pub fn estimate_tokens(text: &str) -> usize {
HeuristicTokenCounter.estimate_text(text)
}
pub fn message_tokens(msg: &AgentMessage) -> usize {
HeuristicTokenCounter.estimate_message(msg)
}
pub fn content_tokens(content: &[Content]) -> usize {
HeuristicTokenCounter.estimate_content(content)
}
pub fn total_tokens(messages: &[AgentMessage]) -> usize {
HeuristicTokenCounter.estimate_messages(messages)
}
pub fn resolve_counter(counter: Option<&Arc<dyn TokenCounter>>) -> &dyn TokenCounter {
counter
.map(|c| c.as_ref() as &dyn TokenCounter)
.unwrap_or(&HeuristicTokenCounter)
}