From 853b8a1fb32328ed36a01d7bb6f3b3d5907a2875 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Wed, 24 Jun 2026 13:48:38 -0700 Subject: [PATCH 1/3] refactor: align exporter cost fallback with normalized responses Signed-off-by: mnajafian-nv --- .../core/src/observability/openinference.rs | 13 ++- crates/core/src/observability/otel.rs | 19 +++-- .../unit/observability/openinference_tests.rs | 83 +++++++++++++++++++ .../tests/unit/observability/otel_tests.rs | 57 ++++++++++++- 4 files changed, 159 insertions(+), 13 deletions(-) diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 65a86fa7a..37466c5c0 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -776,7 +776,10 @@ fn end_attributes(event: &Event) -> Vec { if is_llm { push_llm_usage_attributes(&mut attributes, usage.as_ref()); } - if is_llm && let Some(cost_total) = cost_total_from_llm_event(event, fallback_usage.as_ref()) { + if is_llm + && let Some(cost_total) = + cost_total_from_llm_event(event, normalized.as_deref(), fallback_usage.as_ref()) + { attributes.push(KeyValue::new(oi::llm::cost::TOTAL, cost_total)); } if is_llm { @@ -1174,14 +1177,18 @@ fn finish_reason_value(reason: &FinishReason) -> String { } } -fn cost_total_from_llm_event(event: &Event, fallback_usage: Option<&Usage>) -> Option { +fn cost_total_from_llm_event( + event: &Event, + normalized_response: Option<&AnnotatedLlmResponse>, + fallback_usage: Option<&Usage>, +) -> Option { if let Some(cost) = manual::cost_from_manual_llm_output(event.output(), true).map(|(total, _)| total) { return Some(cost); } - if let Some(response) = event.annotated_response() + if let Some(response) = normalized_response && let Some(usage) = response.usage.as_ref() { if let Some(cost) = usage.cost.as_ref() { diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index ea4400ec4..a0037d7e2 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -716,15 +716,16 @@ fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { if let Some(cost) = manual::cost_from_manual_llm_output(event.output(), false) { return Some(cost); } - if let Some(response) = event.annotated_response() - && let Some(usage) = response.usage.as_ref() - { - if let Some(cost) = usage.cost.as_ref() { - return cost_total_and_currency(cost); - } - if let Some(model_name) = response.model.as_deref().or_else(|| event.model_name()) { - return estimate_cost_for_provider(Some(event.name()), model_name, usage) - .and_then(|cost| cost_total_and_currency(&cost)); + if let Some(response) = event.normalized_llm_response() { + let response = response.as_ref(); + if let Some(usage) = response.usage.as_ref() { + if let Some(cost) = usage.cost.as_ref() { + return cost_total_and_currency(cost); + } + if let Some(model_name) = response.model.as_deref().or_else(|| event.model_name()) { + return estimate_cost_for_provider(Some(event.name()), model_name, usage) + .and_then(|cost| cost_total_and_currency(&cost)); + } } } let usage = manual::usage_from_manual_llm_output(event.output())?; diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index b2acb7447..7789a54eb 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -158,6 +158,65 @@ fn install_test_pricing(model_id: &str) { set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); } +fn install_openai_disambiguation_pricing(model_id: &str) { + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + { + "provider": "other", + "model_id": model_id, + "pricing_as_of": "2026-06-05", + "pricing_source": "test", + "rates": { + "input_per_million": 1000.0, + "output_per_million": 1000.0 + }, + "prompt_cache": { + "read_accounting": "included_in_prompt_tokens" + } + }, + { + "provider": "openai", + "model_id": model_id, + "pricing_as_of": "2026-06-05", + "pricing_source": "test", + "rates": { + "input_per_million": 0.15, + "output_per_million": 0.60, + "cache_read_per_million": 0.075 + }, + "prompt_cache": { + "read_accounting": "included_in_prompt_tokens" + } + } + ] + }) + .to_string(), + ) + .unwrap(); + set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); +} + +fn openai_chat_provider_response(model_id: &str) -> Json { + json!({ + "id": "chatcmpl-test", + "object": "chat.completion", + "model": model_id, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 1_000, + "completion_tokens": 500, + "total_tokens": 1_500, + "prompt_tokens_details": {"cached_tokens": 200} + } + }) +} + fn sample_openinference_annotated_request() -> AnnotatedLlmRequest { AnnotatedLlmRequest { messages: vec![ @@ -2655,6 +2714,30 @@ fn llm_end_with_known_model_usage_emits_derived_cost_attribute() { ); } +#[test] +fn llm_end_with_unannotated_openai_response_uses_codec_cost_attribute() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_openai_disambiguation_pricing("priced-model"); + let _reset_guard = ResetPricingResolverGuard; + + let event = make_end_event( + Uuid::now_v7(), + None, + "other", + ScopeType::Llm, + Some(openai_chat_provider_response("priced-model")), + ); + + assert!(event.annotated_response().is_none()); + assert!(event.normalized_llm_response().is_some()); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("llm.cost.total"), + Some(&"0.000435".to_string()) + ); +} + #[test] fn llm_end_with_manual_usage_and_output_model_emits_derived_cost_attribute() { let _pricing_guard = pricing_test_mutex().lock().unwrap(); diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index ee5017ae4..235071e83 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -78,6 +78,14 @@ fn install_test_pricing(model_id: &str) { } fn install_provider_disambiguation_pricing(model_id: &str) { + install_disambiguation_pricing(model_id, "test"); +} + +fn install_openai_disambiguation_pricing(model_id: &str) { + install_disambiguation_pricing(model_id, "openai"); +} + +fn install_disambiguation_pricing(model_id: &str, preferred_provider: &str) { let catalog = PricingCatalog::from_json_str( &json!({ "version": 1, @@ -96,7 +104,7 @@ fn install_provider_disambiguation_pricing(model_id: &str) { } }, { - "provider": "test", + "provider": preferred_provider, "model_id": model_id, "pricing_as_of": "2026-06-05", "pricing_source": "test", @@ -117,6 +125,25 @@ fn install_provider_disambiguation_pricing(model_id: &str) { set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); } +fn openai_chat_provider_response(model_id: &str) -> Json { + json!({ + "id": "chatcmpl-test", + "object": "chat.completion", + "model": model_id, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 1_000, + "completion_tokens": 500, + "total_tokens": 1_500, + "prompt_tokens_details": {"cached_tokens": 200} + } + }) +} + fn reset_global() { crate::shared_runtime::reset_runtime_owner_for_tests(); let context = global_context(); @@ -916,6 +943,34 @@ fn pre_epoch_timestamps_round_trip_through_system_time() { ); } +#[test] +fn llm_end_with_unannotated_openai_response_uses_codec_cost() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_openai_disambiguation_pricing("priced-model"); + let _reset_guard = ResetPricingResolverGuard; + + let event = make_end_event( + Uuid::now_v7(), + None, + "other", + ScopeType::Llm, + Some(openai_chat_provider_response("priced-model")), + ); + + assert!(event.annotated_response().is_none()); + assert!(event.normalized_llm_response().is_some()); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("nemo_relay.llm.cost.total"), + Some(&"0.000435".to_string()) + ); + assert_eq!( + attributes.get("nemo_relay.llm.cost.currency"), + Some(&"USD".to_string()) + ); +} + #[test] fn helper_functions_cover_additional_otel_branches() { let function_end = make_end_event(Uuid::now_v7(), None, "fn-scope", ScopeType::Function, None); From e6ac4157aeef627618d286bf35f8df562df2c56b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Wed, 24 Jun 2026 17:53:41 -0700 Subject: [PATCH 2/3] test: cover normalized cost fallback without usage Signed-off-by: mnajafian-nv --- .../unit/observability/openinference_tests.rs | 18 ++++++++++++++++++ .../tests/unit/observability/otel_tests.rs | 19 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 7789a54eb..d09c3e804 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -2738,6 +2738,24 @@ fn llm_end_with_unannotated_openai_response_uses_codec_cost_attribute() { ); } +#[test] +fn llm_end_with_unannotated_openai_response_without_usage_omits_cost_attribute() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + reset_active_pricing_resolver().unwrap(); + let _reset_guard = ResetPricingResolverGuard; + + let mut output = openai_chat_provider_response("priced-model"); + output.as_object_mut().unwrap().remove("usage"); + let event = make_end_event(Uuid::now_v7(), None, "openai", ScopeType::Llm, Some(output)); + + assert!(event.annotated_response().is_none()); + let normalized = event.normalized_llm_response().unwrap(); + assert!(normalized.usage.is_none()); + + let attributes = attr_map(&end_attributes(&event)); + assert!(!attributes.contains_key("llm.cost.total")); +} + #[test] fn llm_end_with_manual_usage_and_output_model_emits_derived_cost_attribute() { let _pricing_guard = pricing_test_mutex().lock().unwrap(); diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index 235071e83..c657183d7 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -971,6 +971,25 @@ fn llm_end_with_unannotated_openai_response_uses_codec_cost() { ); } +#[test] +fn llm_end_with_unannotated_openai_response_without_usage_omits_cost() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + reset_active_pricing_resolver().unwrap(); + let _reset_guard = ResetPricingResolverGuard; + + let mut output = openai_chat_provider_response("priced-model"); + output.as_object_mut().unwrap().remove("usage"); + let event = make_end_event(Uuid::now_v7(), None, "openai", ScopeType::Llm, Some(output)); + + assert!(event.annotated_response().is_none()); + let normalized = event.normalized_llm_response().unwrap(); + assert!(normalized.usage.is_none()); + + let attributes = attr_map(&end_attributes(&event)); + assert!(!attributes.contains_key("nemo_relay.llm.cost.total")); + assert!(!attributes.contains_key("nemo_relay.llm.cost.currency")); +} + #[test] fn helper_functions_cover_additional_otel_branches() { let function_end = make_end_event(Uuid::now_v7(), None, "fn-scope", ScopeType::Function, None); From 8c0b80cd0b3e8f73a8995e4c35fa63ee6a658729 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Wed, 24 Jun 2026 22:39:56 -0700 Subject: [PATCH 3/3] fix: preserve requested model cost fallback Signed-off-by: mnajafian-nv --- crates/core/src/observability/mod.rs | 25 +++++++++++++ .../core/src/observability/openinference.rs | 9 ++--- crates/core/src/observability/otel.rs | 11 +++--- .../unit/observability/openinference_tests.rs | 31 ++++++++++++++++ .../tests/unit/observability/otel_tests.rs | 35 +++++++++++++++++++ 5 files changed, 103 insertions(+), 8 deletions(-) diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index dd83c5cfb..adbc2fc90 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -21,6 +21,31 @@ pub mod openinference; pub mod otel; pub mod plugin_component; +#[cfg(any(feature = "otel", feature = "openinference"))] +pub(crate) fn estimate_cost_for_response_or_requested_model( + event: &crate::api::event::Event, + response_model: Option<&str>, + usage: &crate::codec::response::Usage, +) -> Option { + // Prefer the provider-echoed model, but fall back to the requested model + // when pricing does not recognize the echoed model alias. + if let Some(model_name) = response_model + && let Some(cost) = crate::codec::response::estimate_cost_for_provider( + Some(event.name()), + model_name, + usage, + ) + { + return Some(cost); + } + + let event_model = event.model_name()?; + if response_model == Some(event_model) { + return None; + } + crate::codec::response::estimate_cost_for_provider(Some(event.name()), event_model, usage) +} + #[cfg(any(feature = "otel", feature = "openinference"))] pub(crate) fn set_span_status_from_event_metadata(span: &mut S, event: &crate::api::event::Event) where diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 37466c5c0..c694166bc 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -20,7 +20,7 @@ use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use super::manual; +use super::{estimate_cost_for_response_or_requested_model, manual}; use crate::api::event::{Event, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; use crate::api::scope::ScopeType; @@ -1194,9 +1194,10 @@ fn cost_total_from_llm_event( if let Some(cost) = usage.cost.as_ref() { return cost.total_or_component_sum_for_currency("USD"); } - if let Some(model_name) = response.model.as_deref().or_else(|| event.model_name()) { - return estimate_cost_for_provider(Some(event.name()), model_name, usage) - .and_then(|cost| cost.total_for_currency("USD")); + if let Some(cost) = + estimate_cost_for_response_or_requested_model(event, response.model.as_deref(), usage) + { + return cost.total_for_currency("USD"); } } diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index a0037d7e2..d85a7f48a 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -20,7 +20,7 @@ use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use super::manual; +use super::{estimate_cost_for_response_or_requested_model, manual}; use crate::api::event::Event; use crate::api::event::ScopeCategory; use crate::api::runtime::EventSubscriberFn; @@ -722,9 +722,12 @@ fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { if let Some(cost) = usage.cost.as_ref() { return cost_total_and_currency(cost); } - if let Some(model_name) = response.model.as_deref().or_else(|| event.model_name()) { - return estimate_cost_for_provider(Some(event.name()), model_name, usage) - .and_then(|cost| cost_total_and_currency(&cost)); + if let Some(cost) = estimate_cost_for_response_or_requested_model( + event, + response.model.as_deref(), + usage, + ) { + return cost_total_and_currency(&cost); } } } diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index d09c3e804..a60533b97 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -2738,6 +2738,37 @@ fn llm_end_with_unannotated_openai_response_uses_codec_cost_attribute() { ); } +#[test] +fn llm_end_with_unpriced_response_model_uses_requested_model_cost_attribute() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_openai_disambiguation_pricing("priced-model"); + let _reset_guard = ResetPricingResolverGuard; + + let event = make_scope_event_with_profile( + ScopeCategory::End, + Uuid::now_v7(), + None, + "openai", + ScopeType::Llm, + Some(openai_chat_provider_response("api-echoed-model")), + Some( + CategoryProfile::builder() + .model_name("priced-model") + .build(), + ), + ); + + assert!(event.annotated_response().is_none()); + let normalized = event.normalized_llm_response().unwrap(); + assert_eq!(normalized.model.as_deref(), Some("api-echoed-model")); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("llm.cost.total"), + Some(&"0.000435".to_string()) + ); +} + #[test] fn llm_end_with_unannotated_openai_response_without_usage_omits_cost_attribute() { let _pricing_guard = pricing_test_mutex().lock().unwrap(); diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index c657183d7..661fe660a 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -971,6 +971,41 @@ fn llm_end_with_unannotated_openai_response_uses_codec_cost() { ); } +#[test] +fn llm_end_with_unpriced_response_model_uses_requested_model_cost() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_openai_disambiguation_pricing("priced-model"); + let _reset_guard = ResetPricingResolverGuard; + + let event = make_scope_event_with_profile( + ScopeCategory::End, + Uuid::now_v7(), + None, + "openai", + ScopeType::Llm, + Some(openai_chat_provider_response("api-echoed-model")), + Some( + CategoryProfile::builder() + .model_name("priced-model") + .build(), + ), + ); + + assert!(event.annotated_response().is_none()); + let normalized = event.normalized_llm_response().unwrap(); + assert_eq!(normalized.model.as_deref(), Some("api-echoed-model")); + + let attributes = attr_map(&end_attributes(&event)); + assert_eq!( + attributes.get("nemo_relay.llm.cost.total"), + Some(&"0.000435".to_string()) + ); + assert_eq!( + attributes.get("nemo_relay.llm.cost.currency"), + Some(&"USD".to_string()) + ); +} + #[test] fn llm_end_with_unannotated_openai_response_without_usage_omits_cost() { let _pricing_guard = pricing_test_mutex().lock().unwrap();