diff --git a/crates/rig-core/src/completion/request.rs b/crates/rig-core/src/completion/request.rs index 4e84fa1bff..1a32016002 100644 --- a/crates/rig-core/src/completion/request.rs +++ b/crates/rig-core/src/completion/request.rs @@ -254,7 +254,7 @@ where /// Struct representing the token usage for a completion request. /// If tokens used are `0`, then the provider failed to supply token usage metrics. -#[derive(Debug, PartialEq, Eq, Clone, Copy, Serialize, Deserialize)] +#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)] pub struct Usage { /// The number of input ("prompt") tokens used in a given request. pub input_tokens: u64, @@ -272,6 +272,9 @@ pub struct Usage { /// The number of tokens spent on internal reasoning / "thoughts" by reasoning-capable /// models (e.g. Gemini thinking, Anthropic extended thinking, OpenAI o-series). pub reasoning_tokens: u64, + /// Provider-reported request cost (USD), if the provider supplies it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cost: Option, } impl Usage { @@ -285,6 +288,7 @@ impl Usage { cache_creation_input_tokens: 0, tool_use_prompt_tokens: 0, reasoning_tokens: 0, + cost: None, } } @@ -303,6 +307,8 @@ impl Default for Usage { } } +impl Eq for Usage {} + impl Add for Usage { type Output = Self; @@ -316,6 +322,7 @@ impl Add for Usage { + other.cache_creation_input_tokens, tool_use_prompt_tokens: self.tool_use_prompt_tokens + other.tool_use_prompt_tokens, reasoning_tokens: self.reasoning_tokens + other.reasoning_tokens, + cost: self.cost.or(other.cost), } } } @@ -329,6 +336,9 @@ impl AddAssign for Usage { self.cache_creation_input_tokens += other.cache_creation_input_tokens; self.tool_use_prompt_tokens += other.tool_use_prompt_tokens; self.reasoning_tokens += other.reasoning_tokens; + if self.cost.is_none() { + self.cost = other.cost; + } } } diff --git a/crates/rig-core/src/providers/anthropic/completion.rs b/crates/rig-core/src/providers/anthropic/completion.rs index 2102fb01bf..1248cfa935 100644 --- a/crates/rig-core/src/providers/anthropic/completion.rs +++ b/crates/rig-core/src/providers/anthropic/completion.rs @@ -251,6 +251,7 @@ impl TryFrom for completion::CompletionResponse for completion::CompletionResponse for completion::CompletionResponse Self::FinalResponse { + fn build_final_response(&self, usage: Self::Usage, cost: Option) -> Self::FinalResponse { CopilotStreamingResponse::Chat(openai::completion::streaming::StreamingCompletionResponse { usage, + cost, }) } diff --git a/crates/rig-core/src/providers/deepseek.rs b/crates/rig-core/src/providers/deepseek.rs index 76597d596b..dd6acd4ff3 100644 --- a/crates/rig-core/src/providers/deepseek.rs +++ b/crates/rig-core/src/providers/deepseek.rs @@ -265,6 +265,7 @@ impl GetTokenUsage for Usage { .and_then(|details| details.reasoning_tokens) .map(u64::from) .unwrap_or(0), + cost: None, } } } diff --git a/crates/rig-core/src/providers/internal/mod.rs b/crates/rig-core/src/providers/internal/mod.rs index 754f746766..29812c00e1 100644 --- a/crates/rig-core/src/providers/internal/mod.rs +++ b/crates/rig-core/src/providers/internal/mod.rs @@ -17,5 +17,6 @@ pub(crate) fn completion_usage( cache_creation_input_tokens: 0, tool_use_prompt_tokens: 0, reasoning_tokens: 0, + cost: None, } } diff --git a/crates/rig-core/src/providers/internal/openai_chat_completions_compatible.rs b/crates/rig-core/src/providers/internal/openai_chat_completions_compatible.rs index 5a7fb23ee8..8f55af471b 100644 --- a/crates/rig-core/src/providers/internal/openai_chat_completions_compatible.rs +++ b/crates/rig-core/src/providers/internal/openai_chat_completions_compatible.rs @@ -104,6 +104,7 @@ pub(crate) struct CompatibleChunk { pub(crate) response_model: Option, pub(crate) choice: Option>, pub(crate) usage: Option, + pub(crate) cost: Option, } pub(crate) type NormalizedCompatibleChunk = @@ -128,6 +129,7 @@ pub(crate) fn normalize_first_choice_chunk( response_id: Option, response_model: Option, usage: Option, + cost: Option, choices: &[Choice], map_choice: F, ) -> CompatibleChunk @@ -142,6 +144,7 @@ where response_model, choice, usage, + cost, } } @@ -162,7 +165,7 @@ pub(crate) trait CompatibleStreamProfile: WasmCompatSend { fn normalize_chunk(&self, data: &str) -> NormalizedCompatibleChunk; - fn build_final_response(&self, usage: Self::Usage) -> Self::FinalResponse; + fn build_final_response(&self, usage: Self::Usage, cost: Option) -> Self::FinalResponse; fn uses_distinct_tool_call_eviction(&self) -> bool { false @@ -231,6 +234,7 @@ where let stream = stream! { let mut tool_calls: HashMap = HashMap::new(); let mut final_usage = None; + let mut final_cost = None; let mut terminated_with_error = false; while let Some(event_result) = event_source.next().await { @@ -270,6 +274,10 @@ where final_usage = Some(usage); } + if let Some(cost) = chunk.cost { + final_cost = Some(cost); + } + let Some(choice) = chunk.choice else { continue; }; @@ -387,7 +395,7 @@ where let final_usage = final_usage.unwrap_or_default(); record_usage(&span, &final_usage); yield Ok(RawStreamingChoice::FinalResponse( - profile.build_final_response(final_usage), + profile.build_final_response(final_usage, final_cost), )); } .instrument(instrument_span); diff --git a/crates/rig-core/src/providers/mira.rs b/crates/rig-core/src/providers/mira.rs index 632bc6e143..77210d1471 100644 --- a/crates/rig-core/src/providers/mira.rs +++ b/crates/rig-core/src/providers/mira.rs @@ -344,6 +344,7 @@ impl TryFrom for completion::CompletionResponse for completion::CompletionResponse for completion::CompletionResponse { model: Option, choices: Vec, usage: Option, + /// Some gateways report cost as a string, others as a number. + #[serde(default, deserialize_with = "deserialize_cost")] + cost: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum CostValue { + Number(f64), + String(String), +} + +fn deserialize_cost<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + Ok(match Option::::deserialize(deserializer)? { + Some(CostValue::Number(n)) => Some(n), + Some(CostValue::String(s)) => s.trim().parse().ok(), + None => None, + }) } /// Final streaming response. `U` is the provider's streaming usage payload @@ -118,6 +139,9 @@ struct StreamingCompletionChunk { #[derive(Clone, Serialize, Deserialize)] pub struct StreamingCompletionResponse { pub usage: U, + /// Provider-reported request cost (USD), when the provider supplies it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cost: Option, } impl GetTokenUsage for StreamingCompletionResponse @@ -125,7 +149,9 @@ where U: GetTokenUsage, { fn token_usage(&self) -> crate::completion::Usage { - self.usage.token_usage() + let mut usage = self.usage.token_usage(); + usage.cost = self.cost; + usage } } @@ -271,6 +297,7 @@ where data.id, data.model, data.usage, + data.cost, &data.choices, |choice| CompatibleChoiceData { // `function_call` is the deprecated pre-tools finish reason @@ -297,8 +324,8 @@ where )) } - fn build_final_response(&self, usage: Self::Usage) -> Self::FinalResponse { - StreamingCompletionResponse { usage } + fn build_final_response(&self, usage: Self::Usage, cost: Option) -> Self::FinalResponse { + StreamingCompletionResponse { usage, cost } } fn decorate_tool_call( @@ -579,6 +606,61 @@ mod tests { ); } + #[tokio::test] + async fn test_streaming_captures_provider_reported_cost() { + use crate::test_utils::MockStreamingClient; + use futures::StreamExt; + + // OpenCode Go emits a final usage chunk followed by a dedicated + // `inference-cost` chunk carrying `cost` as a string. + let client = MockStreamingClient { + sse_bytes: sse_bytes_from_data_lines([ + "{\"choices\":[{\"delta\":{\"content\":\"Hello\",\"tool_calls\":[]}}],\"usage\":null}", + "{\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}", + "{\"choices\":[],\"cost\":\"0.00006972\"}", + "[DONE]", + ]), + }; + + let req = http::Request::builder() + .method("POST") + .uri("http://localhost/v1/chat/completions") + .body(Vec::new()) + .unwrap(); + + let mut stream = send_compatible_streaming_request(client, req) + .await + .unwrap(); + + let mut final_response = None; + while let Some(chunk) = stream.next().await { + if let streaming::StreamedAssistantContent::Final(res) = chunk.unwrap() { + final_response = Some(res); + break; + } + } + + let response = final_response.expect("expected a final streaming response"); + assert_eq!(response.cost, Some(0.00006972)); + assert_eq!(response.token_usage().cost, Some(0.00006972)); + assert_eq!(response.token_usage().input_tokens, 10); + } + + #[test] + fn test_streaming_chunk_parses_cost_string_or_number() { + let chunk: StreamingCompletionChunk = + serde_json::from_str(r#"{"choices":[],"cost":"0.00006972"}"#).unwrap(); + assert_eq!(chunk.cost, Some(0.00006972)); + + let chunk: StreamingCompletionChunk = + serde_json::from_str(r#"{"choices":[],"cost":0.001}"#).unwrap(); + assert_eq!(chunk.cost, Some(0.001)); + + let chunk: StreamingCompletionChunk = + serde_json::from_str(r#"{"choices":[],"usage":null}"#).unwrap(); + assert_eq!(chunk.cost, None); + } + #[tokio::test] async fn test_streaming_usage_only_chunk_is_not_ignored() { use crate::test_utils::MockStreamingClient; diff --git a/crates/rig-core/src/providers/openai/embedding.rs b/crates/rig-core/src/providers/openai/embedding.rs index dabaad7dd9..2e6d6609d7 100644 --- a/crates/rig-core/src/providers/openai/embedding.rs +++ b/crates/rig-core/src/providers/openai/embedding.rs @@ -254,6 +254,7 @@ where cache_creation_input_tokens: 0, tool_use_prompt_tokens: 0, reasoning_tokens: 0, + cost: None, }, None if Ext::REQUIRES_USAGE => { return Err(EmbeddingError::MissingUsage { diff --git a/crates/rig-core/src/providers/openai/responses_api/mod.rs b/crates/rig-core/src/providers/openai/responses_api/mod.rs index 5c021fdae8..e126759880 100644 --- a/crates/rig-core/src/providers/openai/responses_api/mod.rs +++ b/crates/rig-core/src/providers/openai/responses_api/mod.rs @@ -978,6 +978,7 @@ impl GetTokenUsage for ResponsesUsage { .as_ref() .map(|details| details.reasoning_tokens) .unwrap_or(0), + cost: None, } } } diff --git a/crates/rig-core/src/providers/openrouter/client.rs b/crates/rig-core/src/providers/openrouter/client.rs index e76a512378..aae46fd399 100644 --- a/crates/rig-core/src/providers/openrouter/client.rs +++ b/crates/rig-core/src/providers/openrouter/client.rs @@ -135,6 +135,7 @@ impl GetTokenUsage for Usage { cache_creation_input_tokens: cache_creation, tool_use_prompt_tokens: 0, reasoning_tokens: 0, + cost: None, } } } diff --git a/crates/rig-core/src/providers/openrouter/completion.rs b/crates/rig-core/src/providers/openrouter/completion.rs index a0c4804ee5..f0d7937613 100644 --- a/crates/rig-core/src/providers/openrouter/completion.rs +++ b/crates/rig-core/src/providers/openrouter/completion.rs @@ -724,6 +724,7 @@ impl TryFrom for completion::CompletionResponse) -> CompatibleChunk response_model: None, choice: Some(choice), usage: None, + cost: None, } } @@ -71,7 +72,11 @@ impl CompatibleStreamProfile for ErrorAfterPendingToolCallProfile { } } - fn build_final_response(&self, _usage: Self::Usage) -> Self::FinalResponse { + fn build_final_response( + &self, + _usage: Self::Usage, + _cost: Option, + ) -> Self::FinalResponse { MockResponse::new() } } @@ -134,7 +139,11 @@ impl CompatibleStreamProfile for DistinctToolCallEvictionProfile { Ok(choice.map(test_chunk)) } - fn build_final_response(&self, _usage: Self::Usage) -> Self::FinalResponse { + fn build_final_response( + &self, + _usage: Self::Usage, + _cost: Option, + ) -> Self::FinalResponse { MockResponse::new() } @@ -176,7 +185,11 @@ impl CompatibleStreamProfile for FinishReasonCleanupProfile { Ok(choice.map(test_chunk)) } - fn build_final_response(&self, _usage: Self::Usage) -> Self::FinalResponse { + fn build_final_response( + &self, + _usage: Self::Usage, + _cost: Option, + ) -> Self::FinalResponse { MockResponse::new() } }