From fd7720480b7d40e962f3991aeb4db97cfa4857cb Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Tue, 2 Jun 2026 11:49:53 -0400 Subject: [PATCH] usage data types --- crates/goose-types/src/lib.rs | 2 + crates/goose-types/src/usage.rs | 143 +++++++++++++++++++++++++++ crates/goose/src/context_mgmt/mod.rs | 14 ++- crates/goose/src/providers/base.rs | 125 +---------------------- 4 files changed, 156 insertions(+), 128 deletions(-) create mode 100644 crates/goose-types/src/usage.rs diff --git a/crates/goose-types/src/lib.rs b/crates/goose-types/src/lib.rs index 34b520ef50..82d7d51983 100644 --- a/crates/goose-types/src/lib.rs +++ b/crates/goose-types/src/lib.rs @@ -11,9 +11,11 @@ use unicode_normalization::UnicodeNormalization; pub mod canonical; pub mod conversation; pub mod model; +pub mod usage; pub use conversation::Conversation; pub use model::{Config, ConfigError, ConfigParamError, EnvConfig, ModelConfig, ThinkingEffort}; +pub use usage::{ProviderUsage, Usage}; pub type ToolResult = Result; diff --git a/crates/goose-types/src/usage.rs b/crates/goose-types/src/usage.rs new file mode 100644 index 0000000000..6cfe171aa8 --- /dev/null +++ b/crates/goose-types/src/usage.rs @@ -0,0 +1,143 @@ +use serde::{Deserialize, Serialize}; +use std::ops::{Add, AddAssign}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProviderUsage { + pub model: String, + pub usage: Usage, +} + +impl ProviderUsage { + pub fn new(model: String, usage: Usage) -> Self { + Self { model, usage } + } + + pub fn combine_with(&self, other: &ProviderUsage) -> ProviderUsage { + ProviderUsage { + model: self.model.clone(), + usage: self.usage + other.usage, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)] +pub struct Usage { + pub input_tokens: Option, + pub output_tokens: Option, + pub total_tokens: Option, + pub cache_read_input_tokens: Option, + pub cache_write_input_tokens: Option, +} + +fn sum_optionals(a: Option, b: Option) -> Option +where + T: Add + Default, +{ + match (a, b) { + (Some(x), Some(y)) => Some(x + y), + (Some(x), None) => Some(x + T::default()), + (None, Some(y)) => Some(T::default() + y), + (None, None) => None, + } +} + +impl Add for Usage { + type Output = Self; + + fn add(self, other: Self) -> Self { + Self::new( + sum_optionals(self.input_tokens, other.input_tokens), + sum_optionals(self.output_tokens, other.output_tokens), + sum_optionals(self.total_tokens, other.total_tokens), + ) + .with_cache_tokens( + sum_optionals(self.cache_read_input_tokens, other.cache_read_input_tokens), + sum_optionals( + self.cache_write_input_tokens, + other.cache_write_input_tokens, + ), + ) + } +} + +impl AddAssign for Usage { + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } +} + +impl Usage { + pub fn new( + input_tokens: Option, + output_tokens: Option, + total_tokens: Option, + ) -> Self { + let calculated_total = if total_tokens.is_none() { + match (input_tokens, output_tokens) { + (Some(input), Some(output)) => Some(input + output), + (Some(input), None) => Some(input), + (None, Some(output)) => Some(output), + (None, None) => None, + } + } else { + total_tokens + }; + + Self { + input_tokens, + output_tokens, + total_tokens: calculated_total, + cache_read_input_tokens: None, + cache_write_input_tokens: None, + } + } + + pub fn with_cache_tokens( + mut self, + cache_read_input_tokens: Option, + cache_write_input_tokens: Option, + ) -> Self { + self.cache_read_input_tokens = cache_read_input_tokens; + self.cache_write_input_tokens = cache_write_input_tokens; + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn usage_new_calculates_total_when_missing() { + let usage = Usage::new(Some(10), Some(20), None); + assert_eq!(usage.total_tokens, Some(30)); + } + + #[test] + fn usage_add_combines_token_counts() { + let usage_a = + Usage::new(Some(100), Some(20), Some(120)).with_cache_tokens(Some(10), Some(5)); + let usage_b = Usage::new(Some(50), Some(8), Some(58)).with_cache_tokens(Some(4), Some(1)); + + let combined = usage_a + usage_b; + + assert_eq!(combined.input_tokens, Some(150)); + assert_eq!(combined.output_tokens, Some(28)); + assert_eq!(combined.total_tokens, Some(178)); + assert_eq!(combined.cache_read_input_tokens, Some(14)); + assert_eq!(combined.cache_write_input_tokens, Some(6)); + } + + #[test] + fn provider_usage_combines_with_same_model() { + let usage_a = ProviderUsage::new("model-a".to_string(), Usage::new(Some(1), Some(2), None)); + let usage_b = ProviderUsage::new("model-b".to_string(), Usage::new(Some(3), Some(4), None)); + + let combined = usage_a.combine_with(&usage_b); + + assert_eq!(combined.model, "model-a"); + assert_eq!(combined.usage.input_tokens, Some(4)); + assert_eq!(combined.usage.output_tokens, Some(6)); + assert_eq!(combined.usage.total_tokens, Some(10)); + } +} diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 4d62daf8ae..d26762b84f 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -6,6 +6,7 @@ use crate::prompt_template::render_template; use crate::providers::base::{stream_from_single_message, MessageStream}; use crate::providers::base::{Provider, ProviderUsage}; use crate::providers::errors::ProviderError; +use crate::providers::usage_estimator::ensure_usage_tokens; use crate::{config::Config, token_counter::create_token_counter}; use anyhow::Result; use indoc::indoc; @@ -319,10 +320,15 @@ async fn do_compact( Ok((mut response, mut provider_usage)) => { response.role = Role::User; - provider_usage - .ensure_tokens(&system_prompt, &summarization_request, &response, &[]) - .await - .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?; + ensure_usage_tokens( + &mut provider_usage, + &system_prompt, + &summarization_request, + &response, + &[], + ) + .await + .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?; return Ok((response, provider_usage)); } diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 05cf9d60d5..6f86f62311 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -20,12 +20,12 @@ use crate::conversation::Conversation; use crate::model::ModelConfig; use crate::permission::PermissionConfirmation; use crate::utils::safe_truncate; +pub use goose_types::usage::{ProviderUsage, Usage}; use rmcp::model::Tool; use utoipa::ToSchema; use once_cell::sync::Lazy; use regex::Regex; -use std::ops::{Add, AddAssign}; use std::path::PathBuf; use std::pin::Pin; use std::sync::LazyLock; @@ -678,129 +678,6 @@ impl ConfigKey { } } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ProviderUsage { - pub model: String, - pub usage: Usage, -} - -impl ProviderUsage { - pub fn new(model: String, usage: Usage) -> Self { - Self { model, usage } - } - - /// Ensures this ProviderUsage has token counts, estimating them if necessary - pub async fn ensure_tokens( - &mut self, - system_prompt: &str, - request_messages: &[Message], - response: &Message, - tools: &[Tool], - ) -> Result<(), ProviderError> { - crate::providers::usage_estimator::ensure_usage_tokens( - self, - system_prompt, - request_messages, - response, - tools, - ) - .await - .map_err(|e| ProviderError::ExecutionError(format!("Failed to ensure usage tokens: {}", e))) - } - - /// Combine this ProviderUsage with another, adding their token counts - /// Uses the model from this ProviderUsage - pub fn combine_with(&self, other: &ProviderUsage) -> ProviderUsage { - ProviderUsage { - model: self.model.clone(), - usage: self.usage + other.usage, - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)] -pub struct Usage { - pub input_tokens: Option, - pub output_tokens: Option, - pub total_tokens: Option, - pub cache_read_input_tokens: Option, - pub cache_write_input_tokens: Option, -} - -fn sum_optionals(a: Option, b: Option) -> Option -where - T: Add + Default, -{ - match (a, b) { - (Some(x), Some(y)) => Some(x + y), - (Some(x), None) => Some(x + T::default()), - (None, Some(y)) => Some(T::default() + y), - (None, None) => None, - } -} - -impl Add for Usage { - type Output = Self; - - fn add(self, other: Self) -> Self { - Self::new( - sum_optionals(self.input_tokens, other.input_tokens), - sum_optionals(self.output_tokens, other.output_tokens), - sum_optionals(self.total_tokens, other.total_tokens), - ) - .with_cache_tokens( - sum_optionals(self.cache_read_input_tokens, other.cache_read_input_tokens), - sum_optionals( - self.cache_write_input_tokens, - other.cache_write_input_tokens, - ), - ) - } -} - -impl AddAssign for Usage { - fn add_assign(&mut self, rhs: Self) { - *self = *self + rhs; - } -} - -impl Usage { - pub fn new( - input_tokens: Option, - output_tokens: Option, - total_tokens: Option, - ) -> Self { - let calculated_total = if total_tokens.is_none() { - match (input_tokens, output_tokens) { - (Some(input), Some(output)) => Some(input + output), - (Some(input), None) => Some(input), - (None, Some(output)) => Some(output), - (None, None) => None, - } - } else { - total_tokens - }; - - Self { - input_tokens, - output_tokens, - total_tokens: calculated_total, - cache_read_input_tokens: None, - cache_write_input_tokens: None, - } - } - - pub fn with_cache_tokens( - mut self, - cache_read_input_tokens: Option, - cache_write_input_tokens: Option, - ) -> Self { - self.cache_read_input_tokens = cache_read_input_tokens; - self.cache_write_input_tokens = cache_write_input_tokens; - self - } -} - pub(crate) fn current_working_dir() -> PathBuf { std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")) }