From c6b6c5b6dacddccbb0d88d747996bb796d98a395 Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Tue, 28 Jul 2026 13:50:58 -0700 Subject: [PATCH] fix(aether-telemetry): Include reasoning tokens and pricing information --- crates/aether-cli/src/headless/run.rs | 1 + crates/aether-core/src/core/agent.rs | 3 +- crates/aether-core/src/events/turn_event.rs | 3 +- .../aether-telemetry/src/genai_constants.rs | 6 ++ crates/aether-telemetry/src/llm_call_state.rs | 16 ++-- crates/aether-telemetry/src/otel_observer.rs | 36 +++++++- .../aether-telemetry/tests/observer_tests.rs | 50 ++++++++++- .../tests/otlp_collector_tests.rs | 41 ++++++--- crates/llm-codegen/src/lib.rs | 83 ++++++++++++++++++- crates/llm/src/catalog/bedrock.rs | 9 +- crates/llm/src/catalog/mod.rs | 28 +++++++ crates/llm/src/catalog/pricing.rs | 11 +++ crates/llm/src/lib.rs | 2 +- crates/llm/src/llm_response.rs | 3 + .../llm/src/providers/anthropic/streaming.rs | 1 + crates/llm/src/providers/anthropic/types.rs | 1 + crates/llm/src/providers/bedrock/streaming.rs | 2 + crates/llm/src/providers/openai/streaming.rs | 1 + .../providers/openai_compatible/streaming.rs | 9 +- .../src/providers/openai_compatible/types.rs | 1 + .../providers/openai_responses/streaming.rs | 66 ++++++++++++--- .../03_mantle_cache_write.sse | 6 ++ 22 files changed, 343 insertions(+), 36 deletions(-) create mode 100644 crates/llm/src/catalog/pricing.rs create mode 100644 crates/llm/tests/fixtures/openai_responses/03_mantle_cache_write.sse diff --git a/crates/aether-cli/src/headless/run.rs b/crates/aether-cli/src/headless/run.rs index bcb4d73fb..b17766256 100644 --- a/crates/aether-cli/src/headless/run.rs +++ b/crates/aether-cli/src/headless/run.rs @@ -608,6 +608,7 @@ mod tests { provider: None, model: None, display_name: "test".to_string(), + pricing: None, attempt, max_attempts: 3, }) diff --git a/crates/aether-core/src/core/agent.rs b/crates/aether-core/src/core/agent.rs index 4cfe9e4b2..259142a6b 100644 --- a/crates/aether-core/src/core/agent.rs +++ b/crates/aether-core/src/core/agent.rs @@ -732,7 +732,8 @@ impl Agent { AgentEvent::Turn(TurnEvent::LlmCallStarted { purpose, provider: model.as_ref().map(|m| m.provider().to_string()), - model: model.map(|m| m.model_id().into_owned()), + model: model.as_ref().map(|m| m.model_id().into_owned()), + pricing: model.and_then(|m| m.pricing()), display_name: self.llm.display_name(), attempt, max_attempts: self.retry_config.max_attempts, diff --git a/crates/aether-core/src/events/turn_event.rs b/crates/aether-core/src/events/turn_event.rs index 5d3047592..e1c251be8 100644 --- a/crates/aether-core/src/events/turn_event.rs +++ b/crates/aether-core/src/events/turn_event.rs @@ -1,4 +1,4 @@ -use llm::{ContentBlock, StopReason, TokenUsage}; +use llm::{ContentBlock, ModelPricing, StopReason, TokenUsage}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -61,6 +61,7 @@ pub enum TurnEvent { provider: Option, model: Option, display_name: String, + pricing: Option, /// 0 for the initial call, incrementing per retry. attempt: u32, max_attempts: u32, diff --git a/crates/aether-telemetry/src/genai_constants.rs b/crates/aether-telemetry/src/genai_constants.rs index 77caed202..fbf6e8e3e 100644 --- a/crates/aether-telemetry/src/genai_constants.rs +++ b/crates/aether-telemetry/src/genai_constants.rs @@ -41,6 +41,12 @@ pub const TOOL_CALL_ID: &str = "tool_call.id"; pub const TOOL_CALL_NAME: &str = "tool_call.name"; pub const LLM_ATTEMPT: &str = "aether.llm.attempt"; pub const LLM_PURPOSE: &str = "aether.llm.purpose"; +pub const AI_INPUT_TOKEN_PRICE: &str = "$ai_input_token_price"; +pub const AI_OUTPUT_TOKEN_PRICE: &str = "$ai_output_token_price"; +pub const AI_CACHE_READ_TOKEN_PRICE: &str = "$ai_cache_read_token_price"; +pub const AI_CACHE_WRITE_TOKEN_PRICE: &str = "$ai_cache_write_token_price"; +pub const AI_CACHE_REPORTING_EXCLUSIVE: &str = "$ai_cache_reporting_exclusive"; +pub const AI_REASONING_TOKENS: &str = "$ai_reasoning_tokens"; pub fn genai_instrumentation_scope(version: impl Into) -> InstrumentationScope { InstrumentationScope::builder("aether.genai") diff --git a/crates/aether-telemetry/src/llm_call_state.rs b/crates/aether-telemetry/src/llm_call_state.rs index 309eb70d3..21703864a 100644 --- a/crates/aether-telemetry/src/llm_call_state.rs +++ b/crates/aether-telemetry/src/llm_call_state.rs @@ -2,10 +2,11 @@ use crate::content_capture::{ContentBuffer, ContentCapture}; use crate::content_json::output_messages_json; use crate::gen_ai_metrics::GenAiMetrics; use crate::genai_constants::{ - ERROR_TYPE, GEN_AI_OUTPUT_MESSAGES, GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_RESPONSE_TIME_TO_FIRST_CHUNK, - GEN_AI_TOKEN_TYPE, GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, - GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, - GENAI_RESPONSE_START_EVENT, GENAI_TOOL_CALL_START_EVENT, MESSAGE_ID, TOOL_CALL_ID, TOOL_CALL_NAME, + AI_CACHE_REPORTING_EXCLUSIVE, AI_REASONING_TOKENS, ERROR_TYPE, GEN_AI_OUTPUT_MESSAGES, + GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_RESPONSE_TIME_TO_FIRST_CHUNK, GEN_AI_TOKEN_TYPE, + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, GENAI_RESPONSE_START_EVENT, + GENAI_TOOL_CALL_START_EVENT, MESSAGE_ID, TOOL_CALL_ID, TOOL_CALL_NAME, }; use crate::span_guard::{ErrorKind, SpanGuard}; use aether_core::events::{LlmCallOutcome, LlmCallPurpose}; @@ -122,8 +123,13 @@ impl LlmCallState { self.span .set_attribute(KeyValue::new(GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, i64::from(tokens))); } + if let Some(exclusive) = usage.cache_reporting_exclusive { + self.span.set_attribute(KeyValue::new(AI_CACHE_REPORTING_EXCLUSIVE, exclusive)); + } if let Some(tokens) = usage.reasoning_tokens { - self.span.set_attribute(KeyValue::new(GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, i64::from(tokens))); + let tokens = i64::from(tokens); + self.span.set_attribute(KeyValue::new(GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, tokens)); + self.span.set_attribute(KeyValue::new(AI_REASONING_TOKENS, tokens)); } self.record_token_usage(*usage); } diff --git a/crates/aether-telemetry/src/otel_observer.rs b/crates/aether-telemetry/src/otel_observer.rs index 4f008d4ec..e10f2161f 100644 --- a/crates/aether-telemetry/src/otel_observer.rs +++ b/crates/aether-telemetry/src/otel_observer.rs @@ -6,7 +6,7 @@ use crate::llm_call_state::LlmCallState; use crate::span_guard::{ErrorKind, SpanGuard}; use aether_core::events::{AgentEvent, AgentObserver, LlmCallPurpose, MessageEvent, ToolEvent, TurnEvent, TurnOutcome}; use llm::catalog::Provider; -use llm::{ContentBlock, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition}; +use llm::{ContentBlock, ModelPricing, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition}; use opentelemetry::trace::{SpanBuilder, SpanKind, TraceContextExt, Tracer as _}; use opentelemetry::{Context, KeyValue}; use opentelemetry_sdk::trace::SdkTracer; @@ -121,13 +121,22 @@ impl TurnState { fn on_event(&mut self, message: &AgentEvent, instrumentation: &OtelInstrumentation, tools: &[ToolDefinition]) { match message { - AgentEvent::Turn(TurnEvent::LlmCallStarted { purpose, provider, model, display_name, attempt, .. }) => { + AgentEvent::Turn(TurnEvent::LlmCallStarted { + purpose, + provider, + model, + display_name, + pricing, + attempt, + .. + }) => { self.start_llm_call( LlmCallStart { purpose: *purpose, provider: provider.as_deref(), model: model.as_deref(), display_name, + pricing: *pricing, attempt: *attempt, }, instrumentation, @@ -207,6 +216,9 @@ impl TurnState { let mut attributes = metric_attributes.clone(); attributes.push(KeyValue::new(semconv::GEN_AI_REQUEST_STREAM, true)); attributes.push(KeyValue::new(semconv::LLM_ATTEMPT, i64::from(call.attempt))); + if let Some(pricing) = call.pricing { + attributes.extend(pricing_attributes(pricing)); + } // Only chat calls carry the turn's input and tool definitions; a // compaction call's actual input is the internal summarization prompt. if call.purpose == LlmCallPurpose::Chat { @@ -308,5 +320,25 @@ struct LlmCallStart<'a> { provider: Option<&'a str>, model: Option<&'a str>, display_name: &'a str, + pricing: Option, attempt: u32, } + +fn pricing_attributes(pricing: ModelPricing) -> Vec { + let mut attributes = vec![ + KeyValue::new(semconv::AI_INPUT_TOKEN_PRICE, pricing.input_per_million / TOKENS_PER_MILLION), + KeyValue::new(semconv::AI_OUTPUT_TOKEN_PRICE, pricing.output_per_million / TOKENS_PER_MILLION), + ]; + + if let Some(price) = pricing.cache_read_per_million { + attributes.push(KeyValue::new(semconv::AI_CACHE_READ_TOKEN_PRICE, price / TOKENS_PER_MILLION)); + } + + if let Some(price) = pricing.cache_write_per_million { + attributes.push(KeyValue::new(semconv::AI_CACHE_WRITE_TOKEN_PRICE, price / TOKENS_PER_MILLION)); + } + + attributes +} + +const TOKENS_PER_MILLION: f64 = 1_000_000.0; diff --git a/crates/aether-telemetry/tests/observer_tests.rs b/crates/aether-telemetry/tests/observer_tests.rs index 5cea5dc0a..fa91b2ab5 100644 --- a/crates/aether-telemetry/tests/observer_tests.rs +++ b/crates/aether-telemetry/tests/observer_tests.rs @@ -9,7 +9,7 @@ use aether_telemetry::{ GENAI_SEMCONV_SCHEMA_URL, GenAiMetrics, OtelInstrumentation, OtelObserver, genai_instrumentation_scope, }; use llm::testing::llm_response; -use llm::{LlmError, LlmResponse, StopReason, TokenUsage}; +use llm::{LlmError, LlmResponse, ModelPricing, StopReason, TokenUsage}; use opentelemetry::metrics::MeterProvider as _; use opentelemetry::trace::{Status, TracerProvider as _}; use opentelemetry::{Array, Value}; @@ -179,6 +179,53 @@ async fn completed_llm_calls_capture_finish_reasons() -> Result<(), Box Result<(), Box> { + let usage = TokenUsage { + cache_read_tokens: Some(40), + cache_creation_tokens: Some(10), + cache_reporting_exclusive: Some(true), + reasoning_tokens: Some(7), + ..TokenUsage::new(100, 20) + }; + let events = AgentTrace::from_events(vec![ + AgentEvent::Turn(TurnEvent::Started { content: vec![] }), + AgentEvent::Turn(TurnEvent::LlmCallStarted { + purpose: LlmCallPurpose::Chat, + provider: Some("anthropic".to_string()), + model: Some("priced-model".to_string()), + display_name: "priced-model".to_string(), + pricing: Some(ModelPricing { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + cache_write_per_million: Some(3.75), + }), + attempt: 0, + max_attempts: 1, + }), + AgentEvent::Turn(TurnEvent::LlmCallEnded { + purpose: LlmCallPurpose::Chat, + outcome: LlmCallOutcome::Completed { stop_reason: None, usage: Some(usage) }, + }), + AgentEvent::turn_ended(TurnOutcome::Completed), + ]); + + let spans = otel_test().redacting().observe_trace(&events).spans(); + let chat = spans.named("chat priced-model"); + chat.assert_attr("$ai_input_token_price", 0.000_003); + chat.assert_attr("$ai_output_token_price", 0.000_015); + chat.assert_attr("$ai_cache_read_token_price", 0.000_000_3); + chat.assert_attr("$ai_cache_write_token_price", 0.000_003_75); + chat.assert_attr("$ai_cache_reporting_exclusive", true); + chat.assert_attr("$ai_reasoning_tokens", 7); + // PostHog derives these from the semconv cache token attributes; duplicating + // them would be redundant. + chat.assert_no_attr("$ai_cache_read_input_tokens"); + chat.assert_no_attr("$ai_cache_creation_input_tokens"); + Ok(()) +} + #[tokio::test] async fn completed_llm_calls_capture_token_usage_breakdown() -> Result<(), Box> { let usage = TokenUsage { @@ -372,6 +419,7 @@ fn chat_call(provider: &str, model: &str, outcome: LlmCallOutcome) -> [AgentEven provider: Some(provider.to_string()), model: Some(model.to_string()), display_name: model.to_string(), + pricing: None, attempt: 0, max_attempts: 1, }), diff --git a/crates/aether-telemetry/tests/otlp_collector_tests.rs b/crates/aether-telemetry/tests/otlp_collector_tests.rs index 8aaa6a64d..28327e270 100644 --- a/crates/aether-telemetry/tests/otlp_collector_tests.rs +++ b/crates/aether-telemetry/tests/otlp_collector_tests.rs @@ -8,7 +8,7 @@ use axum::body::Bytes; use axum::extract::State; use axum::http::{HeaderMap, StatusCode}; use axum::routing::post; -use llm::TokenUsage; +use llm::{ModelPricing, TokenUsage}; use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use opentelemetry_proto::tonic::trace::v1::Span; @@ -154,14 +154,8 @@ async fn runtime_exports_genai_spans_and_metrics_to_an_otlp_collector() { runtime.shutdown().expect("runtime flushes both signals"); let exports = collector.exports(); - let span_names = exports - .traces - .iter() - .flat_map(|request| &request.resource_spans) - .flat_map(|resource| &resource.scope_spans) - .flat_map(|scope| &scope.spans) - .map(|span| span.name.as_str()) - .collect::>(); + let spans = trace_spans(&exports); + let span_names = spans.iter().map(|span| span.name.as_str()).collect::>(); let metric_names = exports .metrics @@ -174,6 +168,18 @@ async fn runtime_exports_genai_spans_and_metrics_to_an_otlp_collector() { assert!(span_names.contains(&"invoke_agent")); assert!(span_names.contains(&"chat test-model")); + let chat = spans.iter().find(|span| span.name == "chat test-model").expect("chat span exported"); + let attribute_keys = chat.attributes.iter().map(|attribute| attribute.key.as_str()).collect::>(); + for expected in [ + "$ai_input_token_price", + "$ai_output_token_price", + "$ai_cache_read_token_price", + "$ai_cache_write_token_price", + "$ai_cache_reporting_exclusive", + "$ai_reasoning_tokens", + ] { + assert!(attribute_keys.contains(&expected), "OTLP attribute {expected} missing from {attribute_keys:?}"); + } assert!(metric_names.contains(&"gen_ai.client.operation.duration")); assert!(metric_names.contains(&"gen_ai.client.token.usage")); @@ -235,12 +241,27 @@ fn events() -> Vec { provider: Some("anthropic".to_string()), model: Some("test-model".to_string()), display_name: "test-model".to_string(), + pricing: Some(ModelPricing { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + cache_write_per_million: Some(3.75), + }), attempt: 0, max_attempts: 1, }), AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Chat, - outcome: LlmCallOutcome::Completed { stop_reason: None, usage: Some(TokenUsage::new(10, 5)) }, + outcome: LlmCallOutcome::Completed { + stop_reason: None, + usage: Some(TokenUsage { + cache_read_tokens: Some(4), + cache_creation_tokens: Some(2), + cache_reporting_exclusive: Some(true), + reasoning_tokens: Some(3), + ..TokenUsage::new(10, 5) + }), + }, }), AgentEvent::turn_ended(TurnOutcome::Completed), ] diff --git a/crates/llm-codegen/src/lib.rs b/crates/llm-codegen/src/lib.rs index ed4738dd5..3be593891 100644 --- a/crates/llm-codegen/src/lib.rs +++ b/crates/llm-codegen/src/lib.rs @@ -66,8 +66,7 @@ struct ModalitiesData { input: Vec, } -#[derive(Debug, Deserialize)] -#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Deserialize)] struct CostData { #[serde(default)] input: f64, @@ -293,6 +292,7 @@ struct ModelInfo { context_window: u32, reasoning_levels: Vec, input_modalities: Vec, + pricing: Option, supports_prompt_caching: bool, transport: Option, } @@ -422,6 +422,7 @@ fn collect_models_from( let source_context_window = m.limit.as_ref().map_or(0, |l| l.context); let context_window = cfg.explicit_model(&m.id).map_or(source_context_window, |explicit| explicit.context_window); + let pricing = if cfg.dev_id == "codex" { None } else { m.cost.clone() }; Ok(ModelInfo { variant_name: model_id_to_variant(&m.id), model_id: m.id.clone(), @@ -430,6 +431,7 @@ fn collect_models_from( reasoning_levels, input_modalities, supports_prompt_caching: m.cost.as_ref().is_some_and(CostData::has_prompt_caching), + pricing, transport: transport_for_model(cfg, m)?, }) }) @@ -721,6 +723,8 @@ fn emit_provider_impls(provider_models: &ProviderModels) -> TokenStream { }, ); + let pricing_arms = emit_pricing_arms(models); + let modality_methods = ["image", "audio"].iter().map(|modality| { let method = format_ident!("supports_{}", modality); let mod_owned = (*modality).to_string(); @@ -776,6 +780,11 @@ fn emit_provider_impls(provider_models: &ProviderModels) -> TokenStream { match self { #prompt_caching_arms } } + #[allow(clippy::too_many_lines, clippy::match_same_arms, clippy::unreadable_literal)] + pub fn pricing(self) -> Option { + match self { #pricing_arms } + } + #(#modality_methods)* #[allow(clippy::too_many_lines)] @@ -792,6 +801,28 @@ fn emit_provider_impls(provider_models: &ProviderModels) -> TokenStream { quote! { #(#impls)* } } +fn emit_pricing_arms(models: &[ModelInfo]) -> TokenStream { + let arms = models.iter().map(|model| { + let variant = format_ident!("{}", model.variant_name); + let Some(pricing) = &model.pricing else { + return quote! { Self::#variant => None, }; + }; + let input = pricing.input; + let output = pricing.output; + let cache_read = pricing.cache_read.map_or_else(|| quote! { None }, |value| quote! { Some(#value) }); + let cache_write = pricing.cache_write.map_or_else(|| quote! { None }, |value| quote! { Some(#value) }); + quote! { + Self::#variant => Some(ModelPricing { + input_per_million: #input, + output_per_million: #output, + cache_read_per_million: #cache_read, + cache_write_per_million: #cache_write, + }), + } + }); + quote! { #(#arms)* } +} + fn emit_from_str_impl(enum_ident: &proc_macro2::Ident, parser_name: &str, models: &[ModelInfo]) -> TokenStream { let arms = models.iter().map(|m| { let id = &m.model_id; @@ -928,6 +959,7 @@ fn emit_llm_model_impl() -> TokenStream { let reasoning_levels = emit_llm_reasoning_levels(); let supports_reasoning = emit_llm_supports_reasoning(); let supports_prompt_caching = emit_llm_supports_prompt_caching(); + let pricing = emit_llm_pricing(); let modality_methods = ["image", "audio"].iter().map(|m| emit_llm_supports_modality(m)); let transport = emit_llm_transport(); let all = emit_llm_all(); @@ -946,6 +978,7 @@ fn emit_llm_model_impl() -> TokenStream { #reasoning_levels #supports_reasoning #supports_prompt_caching + #pricing #(#modality_methods)* #transport #all @@ -1134,6 +1167,15 @@ fn emit_llm_supports_prompt_caching() -> TokenStream { } } +fn emit_llm_pricing() -> TokenStream { + let body = llm_delegate_with_dynamic_default("pricing", "e! { None }); + quote! { + pub fn pricing(&self) -> Option { + #body + } + } +} + fn emit_llm_transport() -> TokenStream { let body = llm_delegate_with_dynamic_default("transport", "e! { None }); quote! { @@ -1715,6 +1757,43 @@ mod tests { assert!(matches!(error, CodegenError::UnsupportedReasoningEffort { .. })); } + #[test] + fn build_preserves_model_pricing_and_omits_codex_subscription_pricing() { + let mut data = minimal_models_dev_json(); + anthropic_models( + &mut data, + json!({ + "priced": { + "id": "priced", "name": "Priced", "tool_call": true, + "limit": {"context": 200_000, "output": 0}, + "cost": {"input": 3.0, "output": 15.0, "cache_read": 0.3, "cache_write": 3.75} + } + }), + ); + insert_models( + &mut data, + "openai", + json!({ + "gpt-5.5": { + "id": "gpt-5.5", "name": "GPT-5.5", "tool_call": true, + "limit": {"context": 1_050_000, "output": 128_000}, + "cost": {"input": 1.25, "output": 10.0, "cache_read": 0.125} + } + }), + ); + + let models = build_from_value(&data); + let priced = models["anthropic"].iter().find(|model| model.model_id == "priced").unwrap(); + assert_eq!(priced.pricing.as_ref().map(|pricing| pricing.input), Some(3.0)); + assert_eq!(priced.pricing.as_ref().map(|pricing| pricing.output), Some(15.0)); + assert_eq!(priced.pricing.as_ref().and_then(|pricing| pricing.cache_read), Some(0.3)); + assert_eq!(priced.pricing.as_ref().and_then(|pricing| pricing.cache_write), Some(3.75)); + + let codex = models["codex"].iter().find(|model| model.model_id == "gpt-5.5").unwrap(); + assert_eq!(codex.pricing, None); + assert!(codex.supports_prompt_caching); + } + #[test] fn build_derives_prompt_caching_from_cost_fields() { let mut data = minimal_models_dev_json(); diff --git a/crates/llm/src/catalog/bedrock.rs b/crates/llm/src/catalog/bedrock.rs index b70d35463..fa1e8c910 100644 --- a/crates/llm/src/catalog/bedrock.rs +++ b/crates/llm/src/catalog/bedrock.rs @@ -2,8 +2,8 @@ use std::borrow::Cow; use std::str::FromStr; use crate::ReasoningEffort; -use crate::catalog::BedrockFoundationModel; use crate::catalog::transport::ModelTransport; +use crate::catalog::{BedrockFoundationModel, ModelPricing}; #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum BedrockModel { @@ -65,6 +65,13 @@ impl BedrockModel { } } + pub fn pricing(&self) -> Option { + match self { + Self::Foundation(m) => m.pricing(), + Self::Profile(_) => None, + } + } + pub fn transport(&self) -> Option { match self { Self::Foundation(m) => m.transport(), diff --git a/crates/llm/src/catalog/mod.rs b/crates/llm/src/catalog/mod.rs index 93a082d0e..ea0b89034 100644 --- a/crates/llm/src/catalog/mod.rs +++ b/crates/llm/src/catalog/mod.rs @@ -4,10 +4,12 @@ use crate::providers::local::discovery::discover_local_models; mod bedrock; mod model_spec; +mod pricing; pub mod transport; pub use bedrock::BedrockModel; pub use model_spec::{ModelSpec, ModelSpecError, ReasoningEffortError, validate_reasoning_effort}; +pub use pricing::ModelPricing; pub use transport::{ModelTransport, TemplateError, expand_api_template}; include!(concat!(env!("OUT_DIR"), "/generated.rs")); @@ -52,6 +54,32 @@ mod tests { } } + #[test] + fn model_pricing_comes_from_provider_catalog() { + let model: LlmModel = "anthropic:claude-sonnet-4-5".parse().unwrap(); + + assert_eq!( + model.pricing(), + Some(ModelPricing { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: Some(0.3), + cache_write_per_million: Some(3.75), + }) + ); + } + + #[test] + fn models_without_usage_pricing_do_not_inherit_unrelated_prices() { + let codex: LlmModel = "codex:gpt-5.5".parse().unwrap(); + let local: LlmModel = "ollama:local-model".parse().unwrap(); + let bedrock_profile: LlmModel = "bedrock:us.anthropic.claude-future-model-v99:0".parse().unwrap(); + + assert_eq!(codex.pricing(), None); + assert_eq!(local.pricing(), None); + assert_eq!(bedrock_profile.pricing(), None); + } + #[test] fn codex_gpt56_subscription_models_have_expected_metadata() { let expected_levels = &[ diff --git a/crates/llm/src/catalog/pricing.rs b/crates/llm/src/catalog/pricing.rs new file mode 100644 index 000000000..0ebff6182 --- /dev/null +++ b/crates/llm/src/catalog/pricing.rs @@ -0,0 +1,11 @@ +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Provider/model prices sourced from models.dev, denominated in USD per million tokens. +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct ModelPricing { + pub input_per_million: f64, + pub output_per_million: f64, + pub cache_read_per_million: Option, + pub cache_write_per_million: Option, +} diff --git a/crates/llm/src/lib.rs b/crates/llm/src/lib.rs index 6ab5cd2d4..333f2a067 100644 --- a/crates/llm/src/lib.rs +++ b/crates/llm/src/lib.rs @@ -18,7 +18,7 @@ mod tool_schema; mod tools; pub mod types; -pub use catalog::LlmModel; +pub use catalog::{LlmModel, ModelPricing}; pub use chat_message::{AssistantReasoning, ChatMessage, ContentBlock, EncryptedReasoningContent}; pub use context::Context; pub use credential::ProviderCredential; diff --git a/crates/llm/src/llm_response.rs b/crates/llm/src/llm_response.rs index 40ea6dd7a..cedc0b16f 100644 --- a/crates/llm/src/llm_response.rs +++ b/crates/llm/src/llm_response.rs @@ -25,6 +25,9 @@ pub struct TokenUsage { pub cache_read_tokens: Option, #[serde(default)] pub cache_creation_tokens: Option, + /// Whether cache token counts are separate from, rather than included in, `input_tokens`. + #[serde(default)] + pub cache_reporting_exclusive: Option, #[serde(default)] pub input_audio_tokens: Option, #[serde(default)] diff --git a/crates/llm/src/providers/anthropic/streaming.rs b/crates/llm/src/providers/anthropic/streaming.rs index 3233c89d6..c21d3553b 100644 --- a/crates/llm/src/providers/anthropic/streaming.rs +++ b/crates/llm/src/providers/anthropic/streaming.rs @@ -376,6 +376,7 @@ mod tests { output_tokens: 25, cache_read_tokens: Some(60), cache_creation_tokens: Some(40), + cache_reporting_exclusive: Some(true), ..TokenUsage::default() }) ); diff --git a/crates/llm/src/providers/anthropic/types.rs b/crates/llm/src/providers/anthropic/types.rs index 487d6a304..360db1f21 100644 --- a/crates/llm/src/providers/anthropic/types.rs +++ b/crates/llm/src/providers/anthropic/types.rs @@ -267,6 +267,7 @@ impl From<&Usage> for TokenUsage { output_tokens: usage.output_tokens, cache_read_tokens: usage.cache_read_input_tokens, cache_creation_tokens: usage.cache_creation_input_tokens, + cache_reporting_exclusive: Some(true), ..TokenUsage::default() } } diff --git a/crates/llm/src/providers/bedrock/streaming.rs b/crates/llm/src/providers/bedrock/streaming.rs index 5361a41e8..6d2718bb0 100644 --- a/crates/llm/src/providers/bedrock/streaming.rs +++ b/crates/llm/src/providers/bedrock/streaming.rs @@ -19,6 +19,7 @@ impl From<&BedrockTokenUsage> for TokenUsage { output_tokens: u32::try_from(usage.output_tokens).unwrap_or(0), cache_read_tokens: usage.cache_read_input_tokens().and_then(|v| u32::try_from(v).ok()), cache_creation_tokens: usage.cache_write_input_tokens().and_then(|v| u32::try_from(v).ok()), + cache_reporting_exclusive: Some(true), ..TokenUsage::default() } } @@ -368,6 +369,7 @@ mod tests { assert_eq!(sample.output_tokens, 50); assert_eq!(sample.cache_read_tokens, Some(40)); assert_eq!(sample.cache_creation_tokens, Some(20)); + assert_eq!(sample.cache_reporting_exclusive, Some(true)); } _ => panic!("expected Emit(Usage{{..}})"), } diff --git a/crates/llm/src/providers/openai/streaming.rs b/crates/llm/src/providers/openai/streaming.rs index c40aad55b..ad2f7e398 100644 --- a/crates/llm/src/providers/openai/streaming.rs +++ b/crates/llm/src/providers/openai/streaming.rs @@ -16,6 +16,7 @@ impl From for TokenUsage { input_tokens: usage.prompt_tokens, output_tokens: usage.completion_tokens, cache_read_tokens: prompt.cached_tokens, + cache_reporting_exclusive: Some(false), input_audio_tokens: prompt.audio_tokens, reasoning_tokens: completion.reasoning_tokens, output_audio_tokens: completion.audio_tokens, diff --git a/crates/llm/src/providers/openai_compatible/streaming.rs b/crates/llm/src/providers/openai_compatible/streaming.rs index 07b8637b3..6986c7c53 100644 --- a/crates/llm/src/providers/openai_compatible/streaming.rs +++ b/crates/llm/src/providers/openai_compatible/streaming.rs @@ -266,7 +266,13 @@ mod tests { assert_eq!( tokens, - TokenUsage { input_tokens: 100, output_tokens: 50, cache_read_tokens: Some(30), ..TokenUsage::default() } + TokenUsage { + input_tokens: 100, + output_tokens: 50, + cache_read_tokens: Some(30), + cache_reporting_exclusive: Some(false), + ..TokenUsage::default() + } ); } @@ -299,6 +305,7 @@ mod tests { output_tokens: 500, cache_read_tokens: Some(100), cache_creation_tokens: Some(50), + cache_reporting_exclusive: Some(false), input_audio_tokens: Some(10), input_video_tokens: Some(5), reasoning_tokens: Some(300), diff --git a/crates/llm/src/providers/openai_compatible/types.rs b/crates/llm/src/providers/openai_compatible/types.rs index de25357d6..3835301fd 100644 --- a/crates/llm/src/providers/openai_compatible/types.rs +++ b/crates/llm/src/providers/openai_compatible/types.rs @@ -266,6 +266,7 @@ impl From for TokenUsage { output_tokens: u32::try_from(usage.completion_tokens.max(0)).unwrap_or(0), cache_read_tokens: prompt.cached_tokens, cache_creation_tokens: prompt.cache_write_tokens, + cache_reporting_exclusive: Some(false), input_audio_tokens: prompt.audio_tokens, input_video_tokens: prompt.video_tokens, reasoning_tokens: completion.reasoning_tokens, diff --git a/crates/llm/src/providers/openai_responses/streaming.rs b/crates/llm/src/providers/openai_responses/streaming.rs index 84919365e..f723d947c 100644 --- a/crates/llm/src/providers/openai_responses/streaming.rs +++ b/crates/llm/src/providers/openai_responses/streaming.rs @@ -1,18 +1,39 @@ use async_openai::types::responses::{OutputItem, ResponseUsage, Status}; use futures::Stream; -use serde::Deserialize; +use serde::{Deserialize, Deserializer, de::Error as _}; use tokio_stream::StreamExt; use crate::providers::tool_call_collector::ToolCallCollector; use crate::{LlmError, LlmResponse, Result, StopReason, TokenUsage}; -impl From for TokenUsage { - fn from(usage: ResponseUsage) -> Self { +#[derive(Debug)] +pub struct ResponsesUsage { + usage: ResponseUsage, + cache_write_tokens: Option, +} + +impl<'de> Deserialize<'de> for ResponsesUsage { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + let value = serde_json::Value::deserialize(deserializer)?; + let extension = serde_json::from_value::(value.clone()).map_err(D::Error::custom)?; + let usage = serde_json::from_value(value).map_err(D::Error::custom)?; + + Ok(Self { usage, cache_write_tokens: extension.input_tokens_details.cache_write_tokens }) + } +} + +impl From for TokenUsage { + fn from(usage: ResponsesUsage) -> Self { TokenUsage { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cache_read_tokens: Some(usage.input_tokens_details.cached_tokens), - reasoning_tokens: Some(usage.output_tokens_details.reasoning_tokens), + input_tokens: usage.usage.input_tokens, + output_tokens: usage.usage.output_tokens, + cache_read_tokens: Some(usage.usage.input_tokens_details.cached_tokens), + cache_creation_tokens: usage.cache_write_tokens, + cache_reporting_exclusive: Some(false), + reasoning_tokens: Some(usage.usage.output_tokens_details.reasoning_tokens), ..TokenUsage::default() } } @@ -108,7 +129,7 @@ pub struct ResponsesCompletedEvent { #[derive(Debug, Deserialize)] pub struct ResponsesCompleted { #[serde(default)] - pub usage: Option, + pub usage: Option, #[serde(default)] pub status: Option, } @@ -179,6 +200,18 @@ where } } +#[derive(Deserialize, Default)] +struct ResponsesUsageExtension { + #[serde(default)] + input_tokens_details: ResponsesInputTokenDetailsExtension, +} + +#[derive(Deserialize, Default)] +struct ResponsesInputTokenDetailsExtension { + #[serde(default)] + cache_write_tokens: Option, +} + fn process_event( event: ResponsesStreamEvent, tool_collector: &mut ToolCallCollector, @@ -448,6 +481,16 @@ mod tests { assert!(usage.reasoning_tokens.is_some_and(|tokens| tokens > 0), "{usage:?}"); } + #[tokio::test] + async fn captured_mantle_fixture_preserves_cache_write_usage() { + let responses = + process_fixture(include_str!("../../../tests/fixtures/openai_responses/03_mantle_cache_write.sse")).await; + + assert!(responses.iter().all(Result::is_ok), "{responses:?}"); + let usage = fixture_usage(&responses).expect("fixture should report usage"); + assert_eq!(usage.cache_creation_tokens, Some(1024)); + } + /// Decode a captured SSE body and run it through the shared processor. async fn process_fixture(sse: &str) -> Vec> { let events = sse @@ -498,6 +541,7 @@ mod tests { input_tokens: 120, output_tokens: 80, cache_read_tokens: Some(50), + cache_reporting_exclusive: Some(false), reasoning_tokens: Some(30), ..TokenUsage::default() }) @@ -562,7 +606,7 @@ mod tests { }) } - fn completed(status: Status, usage: Option) -> ResponsesStreamEvent { + fn completed(status: Status, usage: Option) -> ResponsesStreamEvent { ResponsesStreamEvent::Completed(ResponsesCompletedEvent { response: ResponsesCompleted { usage, status: Some(status) }, }) @@ -590,7 +634,7 @@ mod tests { ) } - fn make_usage(input_tokens: u32, output_tokens: u32) -> ResponseUsage { + fn make_usage(input_tokens: u32, output_tokens: u32) -> ResponsesUsage { make_usage_full(input_tokens, output_tokens, 0, 0) } @@ -599,7 +643,7 @@ mod tests { output_tokens: u32, cached_tokens: u32, reasoning_tokens: u32, - ) -> ResponseUsage { + ) -> ResponsesUsage { serde_json::from_value(make_usage_json(input_tokens, output_tokens, cached_tokens, reasoning_tokens)).unwrap() } diff --git a/crates/llm/tests/fixtures/openai_responses/03_mantle_cache_write.sse b/crates/llm/tests/fixtures/openai_responses/03_mantle_cache_write.sse new file mode 100644 index 000000000..12bf9d00e --- /dev/null +++ b/crates/llm/tests/fixtures/openai_responses/03_mantle_cache_write.sse @@ -0,0 +1,6 @@ +event: response.created +data: {"type":"response.created","response":{"id":"resp_mantle_cache_write"},"sequence_number":0} + +event: response.completed +data: {"type":"response.completed","response":{"id":"resp_mantle_cache_write","status":"completed","usage":{"input_tokens":1024,"input_tokens_details":{"cached_tokens":0,"cache_write_tokens":1024},"output_tokens":8,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":1032}},"sequence_number":1} +