use bamboo_agent_core::Message;
use bamboo_compression::{TiktokenTokenCounter, TokenCounter};
use bamboo_engine::metrics::types::TokenUsage as MetricsTokenUsage;
pub(super) fn estimate_prompt_tokens(messages: &[Message]) -> u64 {
let counter = TiktokenTokenCounter::default();
u64::from(counter.count_messages(messages))
}
pub(super) fn estimate_completion_tokens(output_text: &str) -> u64 {
let counter = TiktokenTokenCounter::default();
u64::from(counter.count_text(output_text))
}
pub(super) fn build_estimated_usage(
prompt_tokens: u64,
completion_tokens: u64,
) -> MetricsTokenUsage {
MetricsTokenUsage {
prompt_tokens,
completion_tokens,
total_tokens: prompt_tokens.saturating_add(completion_tokens),
}
}