diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 6737d069..c30bb6a1 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::time::Duration; -use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; use super::errors::ProviderError; use super::formats::openai::{create_request, get_usage, response_to_message}; use super::oauth; @@ -243,7 +243,14 @@ impl Provider for DatabricksProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = get_usage(&response)?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::warn!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); diff --git a/crates/goose/src/providers/errors.rs b/crates/goose/src/providers/errors.rs index 3287866a..7bd50f92 100644 --- a/crates/goose/src/providers/errors.rs +++ b/crates/goose/src/providers/errors.rs @@ -19,6 +19,9 @@ pub enum ProviderError { #[error("Execution error: {0}")] ExecutionError(String), + + #[error("Usage data error: {0}")] + UsageError(String), } impl From for ProviderError { diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index 070ad636..06f61ff5 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -1,6 +1,7 @@ use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; use crate::providers::base::Usage; +use crate::providers::errors::ProviderError; use anyhow::{anyhow, Result}; use mcp_core::content::Content; use mcp_core::role::Role; @@ -201,6 +202,10 @@ pub fn get_usage(data: &Value) -> Result { Ok(Usage::new(input_tokens, output_tokens, total_tokens)) } else { + tracing::warn!( + "Failed to get usage data: {}", + ProviderError::UsageError("No usage data found in response".to_string()) + ); // If no usage data, return None for all values Ok(Usage::new(None, None, None)) } diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index 65e94952..02b64ddd 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -1,6 +1,7 @@ use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; use crate::providers::base::Usage; +use crate::providers::errors::ProviderError; use crate::providers::utils::{is_valid_function_name, sanitize_function_name}; use anyhow::Result; use mcp_core::content::Content; @@ -254,6 +255,10 @@ pub fn get_usage(data: &Value) -> Result { .map(|v| v as i32); Ok(Usage::new(input_tokens, output_tokens, total_tokens)) } else { + tracing::warn!( + "Failed to get usage data: {}", + ProviderError::UsageError("No usage data found in response".to_string()) + ); // If no usage data, return None for all values Ok(Usage::new(None, None, None)) } diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 26b33f45..2be6796d 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -1,6 +1,7 @@ use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; use crate::providers::base::Usage; +use crate::providers::errors::ProviderError; use crate::providers::utils::{ convert_image, is_valid_function_name, sanitize_function_name, ImageFormat, }; @@ -221,10 +222,10 @@ pub fn response_to_message(response: Value) -> anyhow::Result { }) } -pub fn get_usage(data: &Value) -> anyhow::Result { +pub fn get_usage(data: &Value) -> Result { let usage = data .get("usage") - .ok_or_else(|| anyhow!("No usage data in response"))?; + .ok_or_else(|| ProviderError::UsageError("No usage data in response".to_string()))?; let input_tokens = usage .get("prompt_tokens") diff --git a/crates/goose/src/providers/groq.rs b/crates/goose/src/providers/groq.rs index 2fcb5672..c230952e 100644 --- a/crates/goose/src/providers/groq.rs +++ b/crates/goose/src/providers/groq.rs @@ -1,7 +1,7 @@ use super::errors::ProviderError; use crate::message::Message; use crate::model::ModelConfig; -use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use crate::providers::utils::get_model; use anyhow::Result; @@ -137,7 +137,14 @@ impl Provider for GroqProvider { let response = self.post(payload.clone()).await?; let message = response_to_message(response.clone())?; - let usage = get_usage(&response)?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::warn!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 6873e2fe..db3f1f59 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -1,4 +1,4 @@ -use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; use super::errors::ProviderError; use super::utils::{get_model, handle_response_openai_compat}; use crate::message::Message; @@ -104,7 +104,14 @@ impl Provider for OllamaProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = get_usage(&response)?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::warn!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) diff --git a/crates/goose/src/providers/openai.rs b/crates/goose/src/providers/openai.rs index b35ae7c2..26f28e56 100644 --- a/crates/goose/src/providers/openai.rs +++ b/crates/goose/src/providers/openai.rs @@ -4,7 +4,7 @@ use reqwest::Client; use serde_json::Value; use std::time::Duration; -use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; use super::errors::ProviderError; use super::formats::openai::{create_request, get_usage, response_to_message}; use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat}; @@ -115,7 +115,14 @@ impl Provider for OpenAiProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = get_usage(&response)?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::warn!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; let model = get_model(&response); emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index f8fb3f2c..1750776b 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -4,7 +4,7 @@ use reqwest::Client; use serde_json::{json, Value}; use std::time::Duration; -use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; use super::errors::ProviderError; use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat}; use crate::message::Message; @@ -203,7 +203,14 @@ impl Provider for OpenRouterProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = get_usage(&response)?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::warn!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; let model = get_model(&response); emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage)))