From 6fb3f4256c4c3cab0f23673a7798e4854d337fca Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 08:46:17 -0700 Subject: [PATCH 1/6] refactor: align exporter fallback parity Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 142 ++++------ crates/core/src/observability/manual.rs | 151 ++++++++--- crates/core/src/observability/mod.rs | 52 +++- .../core/src/observability/openinference.rs | 54 ++-- crates/core/src/observability/otel.rs | 30 ++- crates/core/tests/unit/atif_tests.rs | 253 ++++++++++++++++++ .../tests/unit/observability/manual_tests.rs | 99 ++++++- .../unit/observability/openinference_tests.rs | 74 +++++ .../tests/unit/observability/otel_tests.rs | 14 + 9 files changed, 671 insertions(+), 198 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 311b6e9e5..ca8e37949 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -39,10 +39,12 @@ use crate::api::event::Event; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; -use crate::codec::response::{AnnotatedLlmResponse, Usage, estimate_cost_for_provider}; +use crate::codec::response::AnnotatedLlmResponse; use crate::error::Result; use crate::json::Json; +use super::{estimate_cost_for_response_or_model, manual, model_name_for_llm_event}; + /// The ATIF schema version string embedded in all exported trajectories. /// /// Currently `"ATIF-v1.7"`. This constant is used by [`AtifTrajectory`] @@ -666,15 +668,29 @@ fn collect_openai_responses_content_text( const TOKEN_USAGE_KNOWN_KEYS: &[&str] = &[ "prompt_tokens", "input_tokens", + "inputTokens", + "input", "completion_tokens", "output_tokens", + "completionTokens", + "outputTokens", + "output", "cached_tokens", + "cachedTokens", + "cache_read_tokens", + "cacheReadTokens", "cache_read_input_tokens", + "cacheReadInputTokens", + "cacheRead", "cache_creation_input_tokens", + "cacheCreationInputTokens", "cache_write_tokens", + "cacheWriteTokens", + "cacheWrite", "cost_usd", "cost", "prompt_tokens_details", + "input_tokens_details", "prompt_token_ids", "completion_token_ids", "logprobs", @@ -691,38 +707,34 @@ fn extract_metrics( model_name: Option<&str>, ) -> Option { let usage = token_usage_object(output)?; - let prompt = usage_u64(usage, &["prompt_tokens", "input_tokens"]); - let completion = usage_u64(usage, &["completion_tokens", "output_tokens"]); - let cache_read = usage_u64(usage, &["cached_tokens"]) - .or_else(|| prompt_tokens_detail_u64(usage, "cached_tokens")) - .or_else(|| input_tokens_detail_u64(usage, "cached_tokens")) - .or_else(|| usage_u64(usage, &["cache_read_input_tokens"])); - let cache_write = usage_u64( - usage, - &["cache_creation_input_tokens", "cache_write_tokens"], - ); + let fallback_usage = manual::usage_from_manual_llm_output(Some(output)); + let prompt = fallback_usage + .as_ref() + .and_then(|usage| usage.prompt_tokens); + let completion = fallback_usage + .as_ref() + .and_then(|usage| usage.completion_tokens); + let cache_read = fallback_usage + .as_ref() + .and_then(|usage| usage.cache_read_tokens); + let cache_write = fallback_usage + .as_ref() + .and_then(|usage| usage.cache_write_tokens); let cached = sum_options(cache_read, cache_write); - let explicit_cost = usage - .get("cost_usd") - .and_then(Json::as_f64) - .or_else(|| usage.get("cost").and_then(cost_usd_from_cost_object)); + let explicit_cost = + manual::cost_from_manual_llm_output(Some(output), manual::ManualCostPolicy::AtifUsdOnly) + .map(|(total, _)| total); let has_reported_cost = usage.get("cost").is_some(); let cost = if has_reported_cost { explicit_cost } else { explicit_cost.or_else(|| { - let model_name = model_name.or_else(|| response_model_name(output))?; - estimate_cost_for_provider( + let usage = fallback_usage.as_ref()?; + estimate_cost_for_response_or_model( provider, model_name, - &Usage { - prompt_tokens: prompt, - completion_tokens: completion, - total_tokens: usage_u64(usage, &["total_tokens"]), - cache_read_tokens: cache_read, - cache_write_tokens: cache_write, - cost: None, - }, + response_model_name(output), + usage, ) .and_then(|cost| cost.total_for_currency("USD")) }) @@ -765,30 +777,6 @@ fn extract_metrics( }) } -fn cost_usd_from_cost_object(cost: &Json) -> Option { - let cost = cost.as_object()?; - let currency = cost.get("currency").and_then(Json::as_str); - let is_relay_normalized_cost = cost - .get("source") - .and_then(Json::as_str) - .is_some_and(|source| matches!(source, "provider_reported" | "model_pricing")); - let has_legacy_provider_total = - currency.is_none() && cost.get("total").and_then(Json::as_f64).is_some(); - let is_usd_cost = currency.is_some_and(|currency| currency.eq_ignore_ascii_case("USD")) - || currency.is_none() && (is_relay_normalized_cost || has_legacy_provider_total); - if !is_usd_cost { - return None; - } - - cost.get("total").and_then(Json::as_f64).or_else(|| { - let (has_component, component_total) = ["input", "output", "cache_read", "cache_write"] - .iter() - .filter_map(|field| cost.get(*field).and_then(Json::as_f64)) - .fold((false, 0.0), |(_, total), value| (true, total + value)); - has_component.then_some(component_total) - }) -} - fn merge_metrics( primary: Option, supplemental: Option<&AtifMetrics>, @@ -854,11 +842,6 @@ fn token_usage_object(output: &Json) -> Option<&serde_json::Map> { .and_then(Json::as_object) } -fn usage_u64(usage: &serde_json::Map, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| usage.get(*key).and_then(Json::as_u64)) -} - fn response_model_name(output: &Json) -> Option<&str> { output .as_object() @@ -873,22 +856,6 @@ fn sum_options(left: Option, right: Option) -> Option { } } -fn prompt_tokens_detail_u64(usage: &serde_json::Map, key: &str) -> Option { - usage - .get("prompt_tokens_details") - .and_then(Json::as_object) - .and_then(|details| details.get(key)) - .and_then(Json::as_u64) -} - -fn input_tokens_detail_u64(usage: &serde_json::Map, key: &str) -> Option { - usage - .get("input_tokens_details") - .and_then(Json::as_object) - .and_then(|details| details.get(key)) - .and_then(Json::as_u64) -} - /// Extract `reasoning_effort` from an LLM request (string or number). /// /// The request content may have `reasoning_effort` (e.g. `"high"`, `"medium"`, @@ -1014,33 +981,13 @@ fn extract_tool_calls(output: &Json) -> Option> { .or_else(|| anthropic_messages_tool_use_items(output))?; let mut calls = Vec::with_capacity(arr.len()); for (index, tc) in arr.iter().enumerate() { - let tc_obj = tc.as_object()?; - let mut id = tc_obj - .get("id") - .or_else(|| tc_obj.get("tool_call_id")) - .or_else(|| tc_obj.get("call_id")) - .and_then(Json::as_str) - .unwrap_or("") - .to_string(); - // The function details live under "function". - let func = tc_obj.get("function").and_then(Json::as_object); - let name = func - .and_then(|f| f.get("name")) - .or_else(|| tc_obj.get("name")) - .or_else(|| tc_obj.get("tool_name")) - .or_else(|| tc_obj.get("function_name")) - .and_then(Json::as_str) - .unwrap_or("") - .to_string(); + let fields = manual::tool_call_fields(tc)?; + let mut id = fields.id.unwrap_or("").to_string(); + let name = fields.name.unwrap_or("").to_string(); if id.is_empty() && !name.is_empty() { id = format!("{name}:{}", index + 1); } - let raw_arguments = func - .and_then(|f| f.get("arguments")) - .or_else(|| tc_obj.get("arguments")) - .or_else(|| tc_obj.get("args")) - .or_else(|| tc_obj.get("input")); - let arguments = normalize_tool_arguments(raw_arguments); + let arguments = normalize_tool_arguments(fields.arguments); // Skip entries with no id and no name — they are not meaningful. if id.is_empty() && name.is_empty() { continue; @@ -1154,6 +1101,7 @@ fn tool_call_extra(tool_call: &Json) -> Option { | "type" | "function" | "name" + | "toolName" | "tool_name" | "function_name" | "arguments" @@ -1454,7 +1402,9 @@ impl LlmSpanCandidate { model_name: start .model_name() .or_else(|| end.model_name()) - .map(ToOwned::to_owned), + .map(ToOwned::to_owned) + .or_else(|| model_name_for_llm_event(start)) + .or_else(|| model_name_for_llm_event(end)), fidelity_score: llm_event_fidelity_score(start).max(llm_event_fidelity_score(end)), end_metrics: end .data() @@ -2292,7 +2242,7 @@ impl StepConversionState { .and_then(|response| atif_message_from_annotated_response(response)) .unwrap_or_else(|| extract_llm_response_message(output)), timestamp: Some(event.timestamp().to_rfc3339()), - model_name: event.model_name().map(ToOwned::to_owned), + model_name: model_name_for_llm_event(event), reasoning_effort, reasoning_content, tool_calls, diff --git a/crates/core/src/observability/manual.rs b/crates/core/src/observability/manual.rs index 0d2977ca8..d414ee362 100644 --- a/crates/core/src/observability/manual.rs +++ b/crates/core/src/observability/manual.rs @@ -2,17 +2,57 @@ // SPDX-License-Identifier: Apache-2.0 //! Shared best-effort extraction for non-provider/manual LLM output, used as the -//! fallback by the OpenTelemetry and OpenInference exporters. +//! fallback by observability exporters before they project into exporter-specific +//! schemas. use serde_json::Map; use crate::codec::response::Usage; use crate::json::Json; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ManualCostPolicy { + #[cfg(feature = "otel")] + AnyCurrency, + #[cfg(feature = "openinference")] + UsdOnly, + AtifUsdOnly, +} + pub(crate) fn model_name_from_manual_llm_output(output: Option<&Json>) -> Option<&str> { output?.as_object()?.get("model").and_then(Json::as_str) } +pub(crate) struct ManualToolCallFields<'a> { + pub(crate) id: Option<&'a str>, + pub(crate) name: Option<&'a str>, + pub(crate) arguments: Option<&'a Json>, +} + +pub(crate) fn tool_call_fields(tool_call: &Json) -> Option> { + let object = tool_call.as_object()?; + let function = object.get("function").and_then(Json::as_object); + Some(ManualToolCallFields { + id: object + .get("id") + .or_else(|| object.get("tool_call_id")) + .or_else(|| object.get("call_id")) + .and_then(Json::as_str), + name: function + .and_then(|function| function.get("name")) + .or_else(|| object.get("name")) + .or_else(|| object.get("toolName")) + .or_else(|| object.get("tool_name")) + .or_else(|| object.get("function_name")) + .and_then(Json::as_str), + arguments: function + .and_then(|function| function.get("arguments")) + .or_else(|| object.get("arguments")) + .or_else(|| object.get("args")) + .or_else(|| object.get("input")), + }) +} + pub(crate) fn usage_from_manual_llm_output(output: Option<&Json>) -> Option { let object = output?.as_object()?; let usage = object.get("usage").and_then(Json::as_object); @@ -42,35 +82,39 @@ pub(crate) fn usage_from_manual_llm_output(output: Option<&Json>) -> Option) -> Option, - usd_only: bool, + policy: ManualCostPolicy, ) -> Option<(f64, String)> { let object = output?.as_object()?; let usage = object.get("usage").and_then(Json::as_object); let token_usage = object.get("token_usage").and_then(Json::as_object); usage - .and_then(|usage| cost_from_manual_usage(usage, usd_only)) - .or_else(|| token_usage.and_then(|usage| cost_from_manual_usage(usage, usd_only))) + .and_then(|usage| cost_from_manual_usage(usage, policy)) + .or_else(|| token_usage.and_then(|usage| cost_from_manual_usage(usage, policy))) } -fn cost_from_manual_usage(usage: &Map, usd_only: bool) -> Option<(f64, String)> { +fn cost_from_manual_usage( + usage: &Map, + policy: ManualCostPolicy, +) -> Option<(f64, String)> { if let Some(total) = usage.get("cost_usd").and_then(Json::as_f64) { return Some((total, "USD".to_string())); } let cost = usage.get("cost")?.as_object()?; - let currency = cost - .get("currency") - .and_then(Json::as_str) - .unwrap_or("USD") - .to_string(); - if usd_only && !currency.eq_ignore_ascii_case("USD") { + let currency = cost.get("currency").and_then(Json::as_str); + if !policy.accepts_currency(currency, cost) { return None; } let total = cost.get("total").and_then(Json::as_f64).or_else(|| { @@ -136,7 +179,33 @@ fn cost_from_manual_usage(usage: &Map, usd_only: bool) -> Option<( .fold((false, 0.0), |(_, total), value| (true, total + value)); has_component.then_some(component_total) })?; - Some((total, currency)) + Some((total, currency.unwrap_or("USD").to_string())) +} + +impl ManualCostPolicy { + fn accepts_currency(self, currency: Option<&str>, cost: &Map) -> bool { + match self { + #[cfg(feature = "otel")] + Self::AnyCurrency => true, + #[cfg(feature = "openinference")] + Self::UsdOnly => currency + .map(|currency| currency.eq_ignore_ascii_case("USD")) + .unwrap_or(true), + Self::AtifUsdOnly => currency + .map(|currency| currency.eq_ignore_ascii_case("USD")) + .unwrap_or_else(|| { + let source_is_relay_normalized = cost + .get("source") + .and_then(Json::as_str) + .is_some_and(|source| { + matches!(source, "provider_reported" | "model_pricing") + }); + let has_legacy_provider_total = + cost.get("total").and_then(Json::as_f64).is_some(); + source_is_relay_normalized || has_legacy_provider_total + }), + } + } } // Keep a reported total only when it is internally consistent with the component diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index adbc2fc90..c92722486 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -13,7 +13,6 @@ pub(crate) fn test_mutex() -> &'static Mutex<()> { pub mod atif; pub mod atof; -#[cfg(any(feature = "otel", feature = "openinference"))] pub(crate) mod manual; #[cfg(feature = "openinference")] pub mod openinference; @@ -26,24 +25,59 @@ 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 { + estimate_cost_for_response_or_model( + Some(event.name()), + event.model_name(), + response_model, + usage, + ) +} + +pub(crate) fn estimate_cost_for_response_or_model( + provider: Option<&str>, + requested_model: Option<&str>, + 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, - ) + && let Some(cost) = + crate::codec::response::estimate_cost_for_provider(provider, model_name, usage) { return Some(cost); } - let event_model = event.model_name()?; - if response_model == Some(event_model) { + let requested_model = requested_model?; + if response_model == Some(requested_model) { + return None; + } + crate::codec::response::estimate_cost_for_provider(provider, requested_model, usage) +} + +pub(crate) fn model_name_for_llm_event(event: &crate::api::event::Event) -> Option { + if let Some(model_name) = event.model_name() { + return Some(model_name.to_string()); + } + if event.category().map(|category| category.as_str()) != Some("llm") { return None; } - crate::codec::response::estimate_cost_for_provider(Some(event.name()), event_model, usage) + event + .normalized_llm_response() + .and_then(|response| response.as_ref().model.clone()) + .or_else(|| { + event + .normalized_llm_request() + .and_then(|request| request.as_ref().model.clone()) + }) + .or_else(|| { + event + .output() + .or_else(|| event.input()) + .and_then(|payload| manual::model_name_from_manual_llm_output(Some(payload))) + .map(ToOwned::to_owned) + }) } #[cfg(any(feature = "otel", feature = "openinference"))] diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index c694166bc..1ffe2dfdd 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -20,7 +20,10 @@ use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use super::{estimate_cost_for_response_or_requested_model, manual}; +use super::{ + estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, + model_name_for_llm_event, +}; use crate::api::event::{Event, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; use crate::api::scope::ScopeType; @@ -28,9 +31,7 @@ use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_ use crate::codec::request::{ AnnotatedLlmRequest, ContentPart, Message, MessageContent, ToolDefinition, }; -use crate::codec::response::{ - AnnotatedLlmResponse, FinishReason, ResponseToolCall, Usage, estimate_cost_for_provider, -}; +use crate::codec::response::{AnnotatedLlmResponse, FinishReason, ResponseToolCall, Usage}; use crate::error::FlowError; use crate::json::Json; use chrono::{DateTime, Utc}; @@ -1100,13 +1101,16 @@ fn push_raw_output_tool_calls( tool_calls: &[Json], ) { for (call_index, tool_call) in tool_calls.iter().enumerate() { + let Some(fields) = manual::tool_call_fields(tool_call) else { + continue; + }; push_output_tool_call( attributes, message_index, call_index, - tool_call.get("id").and_then(Json::as_str), - raw_tool_call_name(tool_call), - raw_tool_call_arguments(tool_call).and_then(|value| { + fields.id, + fields.name, + fields.arguments.and_then(|value| { value .as_str() .map(str::to_string) @@ -1116,23 +1120,6 @@ fn push_raw_output_tool_calls( } } -fn raw_tool_call_name(tool_call: &Json) -> Option<&str> { - tool_call - .get("function") - .and_then(|function| function.get("name")) - .and_then(Json::as_str) - .or_else(|| tool_call.get("name").and_then(Json::as_str)) - .or_else(|| tool_call.get("toolName").and_then(Json::as_str)) -} - -fn raw_tool_call_arguments(tool_call: &Json) -> Option<&Json> { - tool_call - .get("function") - .and_then(|function| function.get("arguments")) - .or_else(|| tool_call.get("arguments")) - .or_else(|| tool_call.get("input")) -} - fn push_output_tool_call( attributes: &mut Vec, message_index: usize, @@ -1183,7 +1170,8 @@ fn cost_total_from_llm_event( fallback_usage: Option<&Usage>, ) -> Option { if let Some(cost) = - manual::cost_from_manual_llm_output(event.output(), true).map(|(total, _)| total) + manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::UsdOnly) + .map(|(total, _)| total) { return Some(cost); } @@ -1202,11 +1190,13 @@ fn cost_total_from_llm_event( } let usage = fallback_usage?; - let model_name = event - .model_name() - .or_else(|| manual::model_name_from_manual_llm_output(event.output()))?; - estimate_cost_for_provider(Some(event.name()), model_name, usage) - .and_then(|cost| cost.total_for_currency("USD")) + estimate_cost_for_response_or_model( + Some(event.name()), + event.model_name(), + manual::model_name_from_manual_llm_output(event.output()), + usage, + ) + .and_then(|cost| cost.total_for_currency("USD")) } fn mark_attributes(event: &Event) -> Vec { @@ -1255,8 +1245,8 @@ fn common_attributes(event: &Event) -> Vec { ), ]; - if let Some(model_name) = event.model_name() { - attributes.push(KeyValue::new(oi::llm::MODEL_NAME, model_name.to_string())); + if let Some(model_name) = model_name_for_llm_event(event) { + attributes.push(KeyValue::new(oi::llm::MODEL_NAME, model_name)); } if let Some(tool_call_id) = event.tool_call_id() { attributes.push(KeyValue::new(oi::tool_call::ID, tool_call_id.to_string())); diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index d85a7f48a..4427a591f 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -20,13 +20,16 @@ use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use super::{estimate_cost_for_response_or_requested_model, manual}; +use super::{ + estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, + model_name_for_llm_event, +}; use crate::api::event::Event; use crate::api::event::ScopeCategory; use crate::api::runtime::EventSubscriberFn; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; -use crate::codec::response::{CostEstimate, estimate_cost_for_provider}; +use crate::codec::response::CostEstimate; use crate::error::FlowError; use chrono::{DateTime, Utc}; use opentelemetry::trace::{ @@ -713,7 +716,9 @@ fn end_attributes(event: &Event) -> Vec { } fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { - if let Some(cost) = manual::cost_from_manual_llm_output(event.output(), false) { + if let Some(cost) = + manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::AnyCurrency) + { return Some(cost); } if let Some(response) = event.normalized_llm_response() { @@ -732,11 +737,13 @@ fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { } } let usage = manual::usage_from_manual_llm_output(event.output())?; - let model_name = event - .model_name() - .or_else(|| manual::model_name_from_manual_llm_output(event.output()))?; - estimate_cost_for_provider(Some(event.name()), model_name, &usage) - .and_then(|cost| cost_total_and_currency(&cost)) + estimate_cost_for_response_or_model( + Some(event.name()), + event.model_name(), + manual::model_name_from_manual_llm_output(event.output()), + &usage, + ) + .and_then(|cost| cost_total_and_currency(&cost)) } fn cost_total_and_currency(cost: &CostEstimate) -> Option<(f64, String)> { @@ -785,11 +792,8 @@ fn common_attributes(event: &Event) -> Vec { ), ]; - if let Some(model_name) = event.model_name() { - attributes.push(KeyValue::new( - "nemo_relay.model_name", - model_name.to_string(), - )); + if let Some(model_name) = model_name_for_llm_event(event) { + attributes.push(KeyValue::new("nemo_relay.model_name", model_name)); } if let Some(tool_call_id) = event.tool_call_id() { attributes.push(KeyValue::new( diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 660fe0ba8..fdbf8218c 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -60,6 +60,46 @@ fn install_test_pricing(model_id: &str) { set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); } +fn install_two_model_test_pricing(requested_model: &str, response_model: &str) { + let catalog = PricingCatalog::from_json_str( + &json!({ + "version": 1, + "entries": [ + { + "provider": "test", + "model_id": requested_model, + "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": "test", + "model_id": response_model, + "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(); +} + #[derive(Debug, Clone, Copy)] enum EventType { Start, @@ -738,6 +778,77 @@ fn test_exporter_llm_lifecycle() { assert_eq!(fm.total_steps, Some(2)); } +#[test] +fn test_exporter_uses_raw_output_model_name_without_profile() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let end = event_builder(llm_uuid, EventType::End) + .name("llm-call") + .scope_type(ScopeType::Llm) + .output(json!({ + "model": "raw-model", + "content": "done", + "usage": { + "prompt_tokens": 7, + "completion_tokens": 3 + } + })) + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let agent_step = &trajectory.steps[0]; + assert_eq!(agent_step.model_name, Some("raw-model".to_string())); + let metrics = agent_step.metrics.as_ref().unwrap(); + assert_eq!(metrics.prompt_tokens, Some(7)); + assert_eq!(metrics.completion_tokens, Some(3)); +} + +#[test] +fn test_exporter_prefers_explicit_end_model_over_raw_start_model() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let start = event_builder(llm_uuid, EventType::Start) + .name("llm-call") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "raw-start-model", + "messages": [{"role": "user", "content": "hello"}] + })) + .build(); + + let end = event_builder(llm_uuid, EventType::End) + .name("llm-call") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "done", + "usage": { + "prompt_tokens": 7, + "completion_tokens": 3 + } + })) + .model_name("profile-end-model") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(start); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + assert_eq!( + trajectory.steps[1].model_name, + Some("profile-end-model".to_string()) + ); +} + #[test] fn test_extract_metrics_supports_provider_usage_payloads() { let openai_metrics = extract_metrics( @@ -824,6 +935,78 @@ fn test_extract_metrics_supports_provider_usage_payloads() { assert_eq!(non_usd_metrics.cost_usd, None); } +#[test] +fn test_extract_metrics_merges_usage_and_token_usage_fields() { + let metrics = extract_metrics( + &json!({ + "usage": { + "prompt_tokens": 10 + }, + "token_usage": { + "completion_tokens": 20, + "total_tokens": 30 + } + }), + None, + None, + ) + .unwrap(); + + assert_eq!(metrics.prompt_tokens, Some(10)); + assert_eq!(metrics.completion_tokens, Some(20)); + assert_eq!(metrics.extra.as_ref().unwrap()["total_tokens"], json!(30)); +} + +#[test] +fn test_extract_metrics_uses_shared_cache_read_precedence() { + let metrics = extract_metrics( + &json!({ + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "cached_tokens": 12, + "prompt_tokens_details": {"cached_tokens": 42}, + "input_tokens_details": {"cached_tokens": 24}, + "cache_read_input_tokens": 77, + "cache_read_tokens": 99 + } + }), + None, + None, + ) + .unwrap(); + + assert_eq!(metrics.cached_tokens, Some(12)); +} + +#[test] +fn test_extract_metrics_filters_extracted_usage_aliases_from_extra() { + let metrics = extract_metrics( + &json!({ + "usage": { + "inputTokens": 10, + "outputTokens": 20, + "cacheReadInputTokens": 5, + "totalTokens": 35, + "reasoning_tokens": 3 + } + }), + None, + None, + ) + .unwrap(); + + assert_eq!(metrics.prompt_tokens, Some(10)); + assert_eq!(metrics.completion_tokens, Some(20)); + assert_eq!(metrics.cached_tokens, Some(5)); + let extra = metrics.extra.as_ref().unwrap().as_object().unwrap(); + assert!(!extra.contains_key("inputTokens")); + assert!(!extra.contains_key("outputTokens")); + assert!(!extra.contains_key("cacheReadInputTokens")); + assert_eq!(extra.get("totalTokens"), Some(&json!(35))); + assert_eq!(extra.get("reasoning_tokens"), Some(&json!(3))); +} + #[test] fn test_reported_cost_object_blocks_model_pricing_estimation() { let _pricing_guard = pricing_test_mutex().lock().unwrap(); @@ -945,6 +1128,76 @@ fn test_exporter_derives_llm_cost_from_model_pricing() { ); } +#[test] +fn test_exporter_prefers_response_model_before_requested_model_pricing() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_two_model_test_pricing("requested-model", "response-model"); + let _reset_guard = ResetPricingResolverGuard; + + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let end = event_builder(llm_uuid, EventType::End) + .name("test") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "priced response", + "model": "response-model", + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + "prompt_tokens_details": {"cached_tokens": 200} + } + })) + .model_name("requested-model") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let metrics = trajectory.steps[0].metrics.as_ref().unwrap(); + assert_eq!(metrics.cost_usd, Some(0.000_435)); +} + +#[test] +fn test_exporter_falls_back_to_requested_model_when_response_model_unpriced() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_test_pricing("requested-model"); + let _reset_guard = ResetPricingResolverGuard; + + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let end = event_builder(llm_uuid, EventType::End) + .name("test") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "priced response", + "model": "api-echoed-model", + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + "prompt_tokens_details": {"cached_tokens": 200} + } + })) + .model_name("requested-model") + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let metrics = trajectory.steps[0].metrics.as_ref().unwrap(); + assert_eq!(metrics.cost_usd, Some(0.000_435)); +} + #[test] fn test_exporter_uses_normalized_usage_cost_before_model_pricing() { let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); diff --git a/crates/core/tests/unit/observability/manual_tests.rs b/crates/core/tests/unit/observability/manual_tests.rs index 4e11c54c3..0c8a27e8d 100644 --- a/crates/core/tests/unit/observability/manual_tests.rs +++ b/crates/core/tests/unit/observability/manual_tests.rs @@ -2,54 +2,89 @@ // SPDX-License-Identifier: Apache-2.0 //! Unit tests for the shared manual/non-provider fallback helpers, focused on -//! the points where the OpenTelemetry and OpenInference copies diverged. +//! exporter-specific policy boundaries and shared scalar extraction. use super::*; use serde_json::json; #[test] +#[cfg(feature = "otel")] fn cost_otel_policy_emits_any_currency() { let output = json!({"usage": {"cost": {"total": 0.5, "currency": "EUR"}}}); assert_eq!( - cost_from_manual_llm_output(Some(&output), false), + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::AnyCurrency), Some((0.5, "EUR".to_string())) ); } #[test] +#[cfg(feature = "openinference")] fn cost_openinference_policy_drops_non_usd() { let output = json!({"usage": {"cost": {"total": 0.5, "currency": "EUR"}}}); - assert_eq!(cost_from_manual_llm_output(Some(&output), true), None); + assert_eq!( + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::UsdOnly), + None + ); } #[test] +#[cfg(feature = "otel")] fn cost_component_sum_emits_currency_for_otel() { let output = json!({"usage": {"cost": {"input": 0.5, "output": 0.375, "currency": "EUR"}}}); assert_eq!( - cost_from_manual_llm_output(Some(&output), false), + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::AnyCurrency), Some((0.875, "EUR".to_string())) ); } #[test] +#[cfg(feature = "openinference")] fn cost_usd_field_passes_usd_only() { let output = json!({"usage": {"cost_usd": 1.25}}); assert_eq!( - cost_from_manual_llm_output(Some(&output), true), + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::UsdOnly), Some((1.25, "USD".to_string())) ); } #[test] +#[cfg(feature = "openinference")] fn cost_absent_currency_treated_as_usd() { let output = json!({"usage": {"cost": {"total": 0.9}}}); assert_eq!( - cost_from_manual_llm_output(Some(&output), true), + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::UsdOnly), Some((0.9, "USD".to_string())) ); } #[test] +fn cost_atif_policy_rejects_component_only_cost_without_currency() { + let output = json!({"usage": {"cost": {"input": 0.5, "output": 0.375}}}); + assert_eq!( + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::AtifUsdOnly), + None + ); +} + +#[test] +fn cost_atif_policy_accepts_relay_normalized_component_cost_without_currency() { + let output = json!({ + "usage": { + "cost": { + "input": 0.5, + "output": 0.375, + "source": "provider_reported" + } + } + }); + assert_eq!( + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::AtifUsdOnly), + Some((0.875, "USD".to_string())) + ); +} + +#[test] +#[cfg(feature = "openinference")] fn cost_per_map_fallthrough_under_usd_only() { // A non-USD `usage` cost is skipped under usd_only; `token_usage` USD wins. let output = json!({ @@ -57,7 +92,7 @@ fn cost_per_map_fallthrough_under_usd_only() { "token_usage": {"cost_usd": 0.2} }); assert_eq!( - cost_from_manual_llm_output(Some(&output), true), + cost_from_manual_llm_output(Some(&output), ManualCostPolicy::UsdOnly), Some((0.2, "USD".to_string())) ); } @@ -94,6 +129,21 @@ fn usage_extracts_aliases_and_returns_none_without_tokens() { assert!(usage_from_manual_llm_output(Some(&json!({}))).is_none()); } +#[test] +fn usage_cache_read_prefers_legacy_atif_order_for_conflicts() { + let output = json!({ + "usage": { + "cache_read_tokens": 99, + "cache_read_input_tokens": 77, + "prompt_tokens_details": {"cached_tokens": 42}, + "input_tokens_details": {"cached_tokens": 24}, + "cached_tokens": 12 + } + }); + let usage = usage_from_manual_llm_output(Some(&output)).expect("has tokens"); + assert_eq!(usage.cache_read_tokens, Some(12)); +} + #[test] fn model_name_extraction() { assert_eq!( @@ -103,3 +153,38 @@ fn model_name_extraction() { assert_eq!(model_name_from_manual_llm_output(Some(&json!({}))), None); assert_eq!(model_name_from_manual_llm_output(None), None); } + +#[test] +fn tool_call_fields_prefers_function_fields_before_aliases() { + let tool_call = json!({ + "id": "call-1", + "call_id": "call-2", + "name": "top_level_name", + "toolName": "camel_name", + "function": { + "name": "function_name", + "arguments": {"query": "needle"} + }, + "arguments": {"query": "fallback"} + }); + + let fields = tool_call_fields(&tool_call).unwrap(); + assert_eq!(fields.id, Some("call-1")); + assert_eq!(fields.name, Some("function_name")); + assert_eq!(fields.arguments, Some(&json!({"query": "needle"}))); +} + +#[test] +fn tool_call_fields_supports_flat_aliases() { + let tool_call = json!({ + "call_id": "call-3", + "function_name": "search_docs", + "args": "{\"query\":\"docs\"}" + }); + + let fields = tool_call_fields(&tool_call).unwrap(); + assert_eq!(fields.id, Some("call-3")); + assert_eq!(fields.name, Some("search_docs")); + assert_eq!(fields.arguments, Some(&json!("{\"query\":\"docs\"}"))); + assert!(tool_call_fields(&json!("not an object")).is_none()); +} diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index a60533b97..27c04653a 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -1297,6 +1297,65 @@ fn openclaw_replay_payloads_emit_flattened_openinference_llm_attributes() { assert_no_attr_contains(&attributes, "secret-token"); } +#[test] +fn openclaw_replay_tool_call_alias_fields_emit_openinference_attributes() { + let (provider, exporter) = make_provider(); + let mut processor = + OpenInferenceEventProcessor::new(provider.clone(), "test-scope".to_string()); + let uuid = Uuid::now_v7(); + + processor.process(&make_start_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "content": { + "messages": [{"role": "user", "content": "Find docs."}], + "source": "openclaw.llm_output" + } + })), + )); + processor.process(&make_end_event( + uuid, + None, + "openclaw-model-call", + ScopeType::Llm, + Some(json!({ + "role": "assistant", + "tool_calls": [{ + "call_id": "call-search-docs", + "toolName": "search_docs", + "args": {"query": "docs"} + }], + "openclaw": { + "assistant_tool_call_names": ["search_docs"] + } + })), + )); + + processor.force_flush().unwrap(); + + let spans = exporter.get_finished_spans().unwrap(); + assert_eq!(spans.len(), 1); + let attributes = attr_map(&spans[0].attributes); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.0.tool_call.id", + "call-search-docs", + ); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.0.tool_call.function.name", + "search_docs", + ); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments", + "{\"query\":\"docs\"}", + ); +} + #[test] fn openclaw_subagent_scopes_preserve_nested_and_fallback_parent_linkage() { let (provider, exporter) = make_provider(); @@ -2404,6 +2463,21 @@ fn helper_functions_cover_additional_openinference_branches() { llm_attributes.get(oi::llm::MODEL_NAME.as_str()), Some(&"demo-model".to_string()) ); + let raw_model_end = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("raw-llm") + .data(json!({"model": "raw-model", "answer": "ok"})) + .build(), + ScopeCategory::End, + Vec::new(), + EventCategory::llm(), + None, + )); + let raw_model_attributes = attr_map(&common_attributes(&raw_model_end)); + assert_eq!( + raw_model_attributes.get(oi::llm::MODEL_NAME.as_str()), + Some(&"raw-model".to_string()) + ); assert_eq!( llm_attributes.get(oi::METADATA.as_str()), Some(&"{\"phase\":\"done\"}".to_string()) diff --git a/crates/core/tests/unit/observability/otel_tests.rs b/crates/core/tests/unit/observability/otel_tests.rs index 661fe660a..09d769ad4 100644 --- a/crates/core/tests/unit/observability/otel_tests.rs +++ b/crates/core/tests/unit/observability/otel_tests.rs @@ -1057,6 +1057,20 @@ fn helper_functions_cover_additional_otel_branches() { llm_attributes.get("nemo_relay.model_name"), Some(&"demo-model".to_string()) ); + let raw_model_event = make_scope_event_with_profile( + ScopeCategory::End, + Uuid::now_v7(), + None, + "chat", + ScopeType::Llm, + Some(json!({"model": "raw-model", "answer": "ok"})), + None, + ); + let raw_model_attributes = attr_map(&common_attributes(&raw_model_event)); + assert_eq!( + raw_model_attributes.get("nemo_relay.model_name"), + Some(&"raw-model".to_string()) + ); let tool_event = Event::Scope(ScopeEvent::new( BaseEvent::builder() From 00fa2587d81ff01e8e070a5e5e515d71c8154053 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 09:07:14 -0700 Subject: [PATCH 2/6] fix: support flat OpenInference tool aliases Signed-off-by: mnajafian-nv --- crates/core/src/observability/openinference.rs | 2 ++ .../unit/observability/openinference_tests.rs | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 1ffe2dfdd..a074f6c0f 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -1528,6 +1528,8 @@ fn tool_call_name(value: &Json) -> Option { .and_then(|function| function.get("name")) .and_then(Json::as_str) }) + .or_else(|| value.get("tool_name").and_then(Json::as_str)) + .or_else(|| value.get("function_name").and_then(Json::as_str)) .map(str::to_string) } diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 27c04653a..4dc84f561 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -1764,6 +1764,24 @@ fn output_value_extracts_chat_completion_display_text() { assert_eq!(attributes.get("llm.cost.total"), Some(&"0.001".to_string())); } +#[test] +fn display_text_from_tool_calls_supports_flat_aliases_without_changing_precedence() { + assert_eq!( + display_text_from_tool_calls(&json!([ + { + "name": "top_level", + "toolName": "camel", + "function": {"name": "nested"}, + "tool_name": "snake", + "function_name": "flat_function" + }, + {"tool_name": "search_docs"}, + {"function_name": "read_file"} + ])), + Some("Requested tools: top_level, search_docs, read_file".to_string()) + ); +} + #[test] fn output_value_extracts_openai_responses_display_text_and_usage() { let (provider, exporter) = make_provider(); From 865379a7c3ca71475e22db8846dfa18367ffef66 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 09:55:57 -0700 Subject: [PATCH 3/6] refactor: prefer normalized exporter fallbacks Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 127 +++++++++++++----- crates/core/src/observability/manual.rs | 30 ----- .../core/src/observability/openinference.rs | 55 +++++--- crates/core/src/observability/otel.rs | 10 +- crates/core/tests/unit/atif_tests.rs | 127 +++++++++++++++++- .../tests/unit/observability/manual_tests.rs | 35 ----- .../unit/observability/openinference_tests.rs | 14 +- 7 files changed, 272 insertions(+), 126 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index ca8e37949..a2f72be97 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -39,7 +39,7 @@ use crate::api::event::Event; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; -use crate::codec::response::AnnotatedLlmResponse; +use crate::codec::response::{AnnotatedLlmResponse, Usage}; use crate::error::Result; use crate::json::Json; @@ -705,58 +705,74 @@ fn extract_metrics( output: &Json, provider: Option<&str>, model_name: Option<&str>, + normalized_response: Option<&AnnotatedLlmResponse>, ) -> Option { - let usage = token_usage_object(output)?; + let raw_usage = token_usage_object(output); let fallback_usage = manual::usage_from_manual_llm_output(Some(output)); - let prompt = fallback_usage - .as_ref() - .and_then(|usage| usage.prompt_tokens); - let completion = fallback_usage + let normalized_usage = normalized_response.and_then(|response| response.usage.as_ref()); + let merged_usage = merge_usage(normalized_usage, fallback_usage.as_ref()); + let prompt = merged_usage.as_ref().and_then(|usage| usage.prompt_tokens); + let completion = merged_usage .as_ref() .and_then(|usage| usage.completion_tokens); - let cache_read = fallback_usage + let cache_read = merged_usage .as_ref() .and_then(|usage| usage.cache_read_tokens); - let cache_write = fallback_usage + let cache_write = merged_usage .as_ref() .and_then(|usage| usage.cache_write_tokens); let cached = sum_options(cache_read, cache_write); - let explicit_cost = + let normalized_cost_source = normalized_usage.and_then(|usage| usage.cost.as_ref()); + let normalized_cost = + normalized_cost_source.and_then(|cost| cost.total_or_component_sum_for_currency("USD")); + let manual_cost = manual::cost_from_manual_llm_output(Some(output), manual::ManualCostPolicy::AtifUsdOnly) .map(|(total, _)| total); - let has_reported_cost = usage.get("cost").is_some(); + let explicit_cost = if normalized_cost_source.is_some() { + normalized_cost + } else { + manual_cost + }; + let has_reported_cost = normalized_cost_source.is_some() + || raw_usage.is_some_and(|usage| usage.get("cost").is_some()); let cost = if has_reported_cost { explicit_cost } else { explicit_cost.or_else(|| { - let usage = fallback_usage.as_ref()?; + let usage = merged_usage.as_ref()?; estimate_cost_for_response_or_model( provider, model_name, - response_model_name(output), + normalized_response + .and_then(|response| response.model.as_deref()) + .or_else(|| response_model_name(output)), usage, ) .and_then(|cost| cost.total_for_currency("USD")) }) }; - let prompt_ids = usage - .get("prompt_token_ids") + let prompt_ids = raw_usage + .and_then(|usage| usage.get("prompt_token_ids")) .and_then(Json::as_array) .map(|a| a.iter().filter_map(Json::as_u64).collect()); - let completion_ids = usage - .get("completion_token_ids") + let completion_ids = raw_usage + .and_then(|usage| usage.get("completion_token_ids")) .and_then(Json::as_array) .map(|a| a.iter().filter_map(Json::as_u64).collect()); - let logprobs = usage - .get("logprobs") + let logprobs = raw_usage + .and_then(|usage| usage.get("logprobs")) .and_then(Json::as_array) .map(|a| a.iter().filter_map(Json::as_f64).collect()); let known: std::collections::HashSet<&str> = TOKEN_USAGE_KNOWN_KEYS.iter().copied().collect(); - let extra_map: serde_json::Map = usage - .iter() - .filter(|(k, _)| !known.contains(k.as_str())) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); + let extra_map: serde_json::Map = raw_usage + .map(|usage| { + usage + .iter() + .filter(|(k, _)| !known.contains(k.as_str())) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() + }) + .unwrap_or_default(); let extra = if extra_map.is_empty() { None } else { @@ -777,6 +793,21 @@ fn extract_metrics( }) } +fn merge_usage(primary: Option<&Usage>, secondary: Option<&Usage>) -> Option { + match (primary, secondary) { + (None, None) => None, + (None, Some(usage)) | (Some(usage), None) => Some(usage.clone()), + (Some(primary), Some(secondary)) => Some(Usage { + prompt_tokens: primary.prompt_tokens.or(secondary.prompt_tokens), + completion_tokens: primary.completion_tokens.or(secondary.completion_tokens), + total_tokens: primary.total_tokens.or(secondary.total_tokens), + cache_read_tokens: primary.cache_read_tokens.or(secondary.cache_read_tokens), + cache_write_tokens: primary.cache_write_tokens.or(secondary.cache_write_tokens), + cost: primary.cost.clone().or_else(|| secondary.cost.clone()), + }), + } +} + fn merge_metrics( primary: Option, supplemental: Option<&AtifMetrics>, @@ -981,13 +1012,32 @@ fn extract_tool_calls(output: &Json) -> Option> { .or_else(|| anthropic_messages_tool_use_items(output))?; let mut calls = Vec::with_capacity(arr.len()); for (index, tc) in arr.iter().enumerate() { - let fields = manual::tool_call_fields(tc)?; - let mut id = fields.id.unwrap_or("").to_string(); - let name = fields.name.unwrap_or("").to_string(); + let tc_obj = tc.as_object()?; + let mut id = tc_obj + .get("id") + .or_else(|| tc_obj.get("tool_call_id")) + .or_else(|| tc_obj.get("call_id")) + .and_then(Json::as_str) + .unwrap_or("") + .to_string(); + let func = tc_obj.get("function").and_then(Json::as_object); + let name = func + .and_then(|f| f.get("name")) + .or_else(|| tc_obj.get("name")) + .or_else(|| tc_obj.get("tool_name")) + .or_else(|| tc_obj.get("function_name")) + .and_then(Json::as_str) + .unwrap_or("") + .to_string(); if id.is_empty() && !name.is_empty() { id = format!("{name}:{}", index + 1); } - let arguments = normalize_tool_arguments(fields.arguments); + let raw_arguments = func + .and_then(|f| f.get("arguments")) + .or_else(|| tc_obj.get("arguments")) + .or_else(|| tc_obj.get("args")) + .or_else(|| tc_obj.get("input")); + let arguments = normalize_tool_arguments(raw_arguments); // Skip entries with no id and no name — they are not meaningful. if id.is_empty() && name.is_empty() { continue; @@ -1101,7 +1151,6 @@ fn tool_call_extra(tool_call: &Json) -> Option { | "type" | "function" | "name" - | "toolName" | "tool_name" | "function_name" | "arguments" @@ -1406,9 +1455,15 @@ impl LlmSpanCandidate { .or_else(|| model_name_for_llm_event(start)) .or_else(|| model_name_for_llm_event(end)), fidelity_score: llm_event_fidelity_score(start).max(llm_event_fidelity_score(end)), - end_metrics: end - .data() - .and_then(|output| extract_metrics(output, Some(end.name()), end.model_name())), + end_metrics: end.data().and_then(|output| { + let normalized_response = end.normalized_llm_response(); + extract_metrics( + output, + Some(end.name()), + end.model_name(), + normalized_response.as_deref(), + ) + }), hook_instrumentation: is_hook_instrumented_llm_event(start) || is_hook_instrumented_llm_event(end), gateway_instrumentation: is_gateway_instrumented_llm_event(start) @@ -2230,7 +2285,15 @@ impl StepConversionState { ); let metrics = merge_metrics( - extract_metrics(output, Some(event.name()), event.model_name()), + { + let normalized_response = event.normalized_llm_response(); + extract_metrics( + output, + Some(event.name()), + event.model_name(), + normalized_response.as_deref(), + ) + }, lookups.supplemental_llm_metrics.get(&event.uuid()), ); diff --git a/crates/core/src/observability/manual.rs b/crates/core/src/observability/manual.rs index d414ee362..289dd3ea3 100644 --- a/crates/core/src/observability/manual.rs +++ b/crates/core/src/observability/manual.rs @@ -23,36 +23,6 @@ pub(crate) fn model_name_from_manual_llm_output(output: Option<&Json>) -> Option output?.as_object()?.get("model").and_then(Json::as_str) } -pub(crate) struct ManualToolCallFields<'a> { - pub(crate) id: Option<&'a str>, - pub(crate) name: Option<&'a str>, - pub(crate) arguments: Option<&'a Json>, -} - -pub(crate) fn tool_call_fields(tool_call: &Json) -> Option> { - let object = tool_call.as_object()?; - let function = object.get("function").and_then(Json::as_object); - Some(ManualToolCallFields { - id: object - .get("id") - .or_else(|| object.get("tool_call_id")) - .or_else(|| object.get("call_id")) - .and_then(Json::as_str), - name: function - .and_then(|function| function.get("name")) - .or_else(|| object.get("name")) - .or_else(|| object.get("toolName")) - .or_else(|| object.get("tool_name")) - .or_else(|| object.get("function_name")) - .and_then(Json::as_str), - arguments: function - .and_then(|function| function.get("arguments")) - .or_else(|| object.get("arguments")) - .or_else(|| object.get("args")) - .or_else(|| object.get("input")), - }) -} - pub(crate) fn usage_from_manual_llm_output(output: Option<&Json>) -> Option { let object = output?.as_object()?; let usage = object.get("usage").and_then(Json::as_object); diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index a074f6c0f..3bbcd00e5 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -1101,16 +1101,13 @@ fn push_raw_output_tool_calls( tool_calls: &[Json], ) { for (call_index, tool_call) in tool_calls.iter().enumerate() { - let Some(fields) = manual::tool_call_fields(tool_call) else { - continue; - }; push_output_tool_call( attributes, message_index, call_index, - fields.id, - fields.name, - fields.arguments.and_then(|value| { + raw_tool_call_id(tool_call), + raw_tool_call_name(tool_call), + raw_tool_call_arguments(tool_call).and_then(|value| { value .as_str() .map(str::to_string) @@ -1120,6 +1117,36 @@ fn push_raw_output_tool_calls( } } +// Raw replay payloads are an OpenInference-local fallback. Provider-shaped +// responses should use codec-normalized response tool calls instead. +fn raw_tool_call_id(tool_call: &Json) -> Option<&str> { + tool_call + .get("id") + .or_else(|| tool_call.get("tool_call_id")) + .or_else(|| tool_call.get("call_id")) + .and_then(Json::as_str) +} + +fn raw_tool_call_name(tool_call: &Json) -> Option<&str> { + tool_call + .get("function") + .and_then(|function| function.get("name")) + .and_then(Json::as_str) + .or_else(|| tool_call.get("name").and_then(Json::as_str)) + .or_else(|| tool_call.get("toolName").and_then(Json::as_str)) + .or_else(|| tool_call.get("tool_name").and_then(Json::as_str)) + .or_else(|| tool_call.get("function_name").and_then(Json::as_str)) +} + +fn raw_tool_call_arguments(tool_call: &Json) -> Option<&Json> { + tool_call + .get("function") + .and_then(|function| function.get("arguments")) + .or_else(|| tool_call.get("arguments")) + .or_else(|| tool_call.get("args")) + .or_else(|| tool_call.get("input")) +} + fn push_output_tool_call( attributes: &mut Vec, message_index: usize, @@ -1169,13 +1196,6 @@ fn cost_total_from_llm_event( normalized_response: Option<&AnnotatedLlmResponse>, fallback_usage: Option<&Usage>, ) -> Option { - if let Some(cost) = - manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::UsdOnly) - .map(|(total, _)| total) - { - return Some(cost); - } - if let Some(response) = normalized_response && let Some(usage) = response.usage.as_ref() { @@ -1189,6 +1209,13 @@ fn cost_total_from_llm_event( } } + if let Some(cost) = + manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::UsdOnly) + .map(|(total, _)| total) + { + return Some(cost); + } + let usage = fallback_usage?; estimate_cost_for_response_or_model( Some(event.name()), @@ -1528,8 +1555,6 @@ fn tool_call_name(value: &Json) -> Option { .and_then(|function| function.get("name")) .and_then(Json::as_str) }) - .or_else(|| value.get("tool_name").and_then(Json::as_str)) - .or_else(|| value.get("function_name").and_then(Json::as_str)) .map(str::to_string) } diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index 4427a591f..79b924c98 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -716,11 +716,6 @@ fn end_attributes(event: &Event) -> Vec { } fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { - if let Some(cost) = - manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::AnyCurrency) - { - return Some(cost); - } if let Some(response) = event.normalized_llm_response() { let response = response.as_ref(); if let Some(usage) = response.usage.as_ref() { @@ -736,6 +731,11 @@ fn cost_from_llm_event(event: &Event) -> Option<(f64, String)> { } } } + if let Some(cost) = + manual::cost_from_manual_llm_output(event.output(), manual::ManualCostPolicy::AnyCurrency) + { + return Some(cost); + } let usage = manual::usage_from_manual_llm_output(event.output())?; estimate_cost_for_response_or_model( Some(event.name()), diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index fdbf8218c..5e6db74e9 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -17,8 +17,8 @@ use crate::codec::openai_responses::OpenAIResponsesCodec; use crate::codec::pricing::pricing_test_mutex; use crate::codec::request::AnnotatedLlmRequest; use crate::codec::response::{ - AnnotatedLlmResponse, PricingCatalog, PricingResolver, reset_active_pricing_resolver, - set_active_pricing_resolver, + AnnotatedLlmResponse, CostEstimate, CostSource, PricingCatalog, PricingResolver, Usage, + reset_active_pricing_resolver, set_active_pricing_resolver, }; use crate::codec::traits::{LlmCodec, LlmResponseCodec}; use serde_json::json; @@ -100,6 +100,35 @@ fn install_two_model_test_pricing(requested_model: &str, response_model: &str) { set_active_pricing_resolver(PricingResolver::from_catalogs(vec![catalog])).unwrap(); } +fn annotated_response_with_usage(model: &str, usage: Usage) -> AnnotatedLlmResponse { + AnnotatedLlmResponse { + id: None, + model: Some(model.to_string()), + message: None, + tool_calls: None, + finish_reason: None, + usage: Some(usage), + api_specific: None, + extra: serde_json::Map::new(), + } +} + +fn provider_reported_cost(total: f64, currency: &str, model: &str) -> CostEstimate { + CostEstimate { + total: Some(total), + currency: currency.to_string(), + input: None, + output: None, + cache_read: None, + cache_write: None, + source: CostSource::ProviderReported, + pricing_provider: None, + pricing_model: Some(model.to_string()), + pricing_as_of: None, + pricing_source: None, + } +} + #[derive(Debug, Clone, Copy)] enum EventType { Start, @@ -865,6 +894,7 @@ fn test_extract_metrics_supports_provider_usage_payloads() { }), None, None, + None, ) .unwrap(); assert_eq!(openai_metrics.prompt_tokens, Some(10)); @@ -890,6 +920,7 @@ fn test_extract_metrics_supports_provider_usage_payloads() { }), None, None, + None, ) .unwrap(); assert_eq!(responses_metrics.prompt_tokens, Some(75)); @@ -913,6 +944,7 @@ fn test_extract_metrics_supports_provider_usage_payloads() { }), None, None, + None, ) .unwrap(); assert_eq!(anthropic_metrics.prompt_tokens, Some(11)); @@ -930,6 +962,7 @@ fn test_extract_metrics_supports_provider_usage_payloads() { }), None, None, + None, ) .unwrap(); assert_eq!(non_usd_metrics.cost_usd, None); @@ -949,6 +982,7 @@ fn test_extract_metrics_merges_usage_and_token_usage_fields() { }), None, None, + None, ) .unwrap(); @@ -973,6 +1007,7 @@ fn test_extract_metrics_uses_shared_cache_read_precedence() { }), None, None, + None, ) .unwrap(); @@ -993,6 +1028,7 @@ fn test_extract_metrics_filters_extracted_usage_aliases_from_extra() { }), None, None, + None, ) .unwrap(); @@ -1027,6 +1063,7 @@ fn test_reported_cost_object_blocks_model_pricing_estimation() { }), Some("test"), Some("priced-model"), + None, ) .unwrap(); @@ -1045,6 +1082,7 @@ fn test_reported_cost_object_blocks_model_pricing_estimation() { }), Some("test"), Some("priced-model"), + None, ) .unwrap(); @@ -1063,6 +1101,7 @@ fn test_reported_cost_object_blocks_model_pricing_estimation() { }), Some("test"), Some("priced-model"), + None, ) .unwrap(); @@ -1084,6 +1123,7 @@ fn test_reported_cost_object_blocks_model_pricing_estimation() { }), Some("test"), Some("priced-model"), + None, ) .unwrap(); @@ -1236,6 +1276,88 @@ fn test_exporter_uses_normalized_usage_cost_before_model_pricing() { ); } +#[test] +fn test_exporter_prefers_annotated_usage_over_raw_usage_conflicts() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let end = event_builder(llm_uuid, EventType::End) + .name("test") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "priced response", + "model": "raw-model", + "usage": { + "prompt_tokens": 999, + "completion_tokens": 888, + "cost": { + "total": 9.99, + "currency": "USD" + } + } + })) + .annotated_response(annotated_response_with_usage( + "annotated-model", + Usage { + prompt_tokens: Some(12), + completion_tokens: Some(34), + total_tokens: Some(46), + cache_read_tokens: Some(5), + cache_write_tokens: None, + cost: Some(provider_reported_cost(0.42, "USD", "annotated-model")), + }, + )) + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let metrics = trajectory.steps[0].metrics.as_ref().unwrap(); + assert_eq!(metrics.prompt_tokens, Some(12)); + assert_eq!(metrics.completion_tokens, Some(34)); + assert_eq!(metrics.cached_tokens, Some(5)); + assert_eq!(metrics.cost_usd, Some(0.42)); +} + +#[test] +fn test_extract_metrics_does_not_mix_raw_cost_when_annotated_cost_is_non_usd() { + let normalized_response = annotated_response_with_usage( + "annotated-model", + Usage { + prompt_tokens: Some(12), + completion_tokens: Some(34), + total_tokens: Some(46), + cache_read_tokens: None, + cache_write_tokens: None, + cost: Some(provider_reported_cost(0.42, "EUR", "annotated-model")), + }, + ); + + let metrics = extract_metrics( + &json!({ + "usage": { + "prompt_tokens": 999, + "completion_tokens": 888, + "cost": { + "total": 9.99, + "currency": "USD" + } + } + }), + None, + None, + Some(&normalized_response), + ) + .unwrap(); + + assert_eq!(metrics.prompt_tokens, Some(12)); + assert_eq!(metrics.completion_tokens, Some(34)); + assert_eq!(metrics.cost_usd, None); +} + #[test] fn test_exporter_omits_cost_for_unknown_model_pricing() { let _pricing_guard = pricing_test_mutex().lock().unwrap(); @@ -1249,6 +1371,7 @@ fn test_exporter_omits_cost_for_unknown_model_pricing() { }), None, Some("unknown-model"), + None, ) .unwrap(); diff --git a/crates/core/tests/unit/observability/manual_tests.rs b/crates/core/tests/unit/observability/manual_tests.rs index 0c8a27e8d..55fafc564 100644 --- a/crates/core/tests/unit/observability/manual_tests.rs +++ b/crates/core/tests/unit/observability/manual_tests.rs @@ -153,38 +153,3 @@ fn model_name_extraction() { assert_eq!(model_name_from_manual_llm_output(Some(&json!({}))), None); assert_eq!(model_name_from_manual_llm_output(None), None); } - -#[test] -fn tool_call_fields_prefers_function_fields_before_aliases() { - let tool_call = json!({ - "id": "call-1", - "call_id": "call-2", - "name": "top_level_name", - "toolName": "camel_name", - "function": { - "name": "function_name", - "arguments": {"query": "needle"} - }, - "arguments": {"query": "fallback"} - }); - - let fields = tool_call_fields(&tool_call).unwrap(); - assert_eq!(fields.id, Some("call-1")); - assert_eq!(fields.name, Some("function_name")); - assert_eq!(fields.arguments, Some(&json!({"query": "needle"}))); -} - -#[test] -fn tool_call_fields_supports_flat_aliases() { - let tool_call = json!({ - "call_id": "call-3", - "function_name": "search_docs", - "args": "{\"query\":\"docs\"}" - }); - - let fields = tool_call_fields(&tool_call).unwrap(); - assert_eq!(fields.id, Some("call-3")); - assert_eq!(fields.name, Some("search_docs")); - assert_eq!(fields.arguments, Some(&json!("{\"query\":\"docs\"}"))); - assert!(tool_call_fields(&json!("not an object")).is_none()); -} diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 4dc84f561..0401c5076 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -1765,20 +1765,20 @@ fn output_value_extracts_chat_completion_display_text() { } #[test] -fn display_text_from_tool_calls_supports_flat_aliases_without_changing_precedence() { +fn display_text_from_tool_calls_preserves_legacy_name_precedence() { assert_eq!( display_text_from_tool_calls(&json!([ { "name": "top_level", "toolName": "camel", - "function": {"name": "nested"}, - "tool_name": "snake", - "function_name": "flat_function" + "function": {"name": "nested"} }, - {"tool_name": "search_docs"}, - {"function_name": "read_file"} + { + "toolName": "camel_without_top_level", + "function": {"name": "nested_without_top_level"} + } ])), - Some("Requested tools: top_level, search_docs, read_file".to_string()) + Some("Requested tools: top_level, camel_without_top_level".to_string()) ); } From 57ec4b5a8a2911bf33e2d22f79630482e7caf1bf Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 11:26:45 -0700 Subject: [PATCH 4/6] fix: address exporter fallback review comments Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 56 +++++++------- crates/core/src/observability/mod.rs | 18 +++++ .../core/src/observability/openinference.rs | 30 ++------ crates/core/tests/unit/atif_tests.rs | 77 ++++++++++++++++++- .../unit/observability/openinference_tests.rs | 40 ++++++++-- 5 files changed, 163 insertions(+), 58 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index a2f72be97..c65dbbc0c 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -39,11 +39,11 @@ use crate::api::event::Event; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; -use crate::codec::response::{AnnotatedLlmResponse, Usage}; +use crate::codec::response::AnnotatedLlmResponse; use crate::error::Result; use crate::json::Json; -use super::{estimate_cost_for_response_or_model, manual, model_name_for_llm_event}; +use super::{estimate_cost_for_response_or_model, manual, merge_usage, model_name_for_llm_event}; /// The ATIF schema version string embedded in all exported trajectories. /// @@ -764,15 +764,18 @@ fn extract_metrics( .and_then(Json::as_array) .map(|a| a.iter().filter_map(Json::as_f64).collect()); let known: std::collections::HashSet<&str> = TOKEN_USAGE_KNOWN_KEYS.iter().copied().collect(); - let extra_map: serde_json::Map = raw_usage - .map(|usage| { - usage - .iter() - .filter(|(k, _)| !known.contains(k.as_str())) - .map(|(k, v)| (k.clone(), v.clone())) - .collect() + let extra_map: serde_json::Map = output + .as_object() + .into_iter() + .flat_map(|output| { + ["usage", "token_usage"] + .into_iter() + .filter_map(|key| output.get(key).and_then(Json::as_object)) }) - .unwrap_or_default(); + .flat_map(|usage| usage.iter()) + .filter(|(k, _)| !known.contains(k.as_str())) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); let extra = if extra_map.is_empty() { None } else { @@ -793,21 +796,6 @@ fn extract_metrics( }) } -fn merge_usage(primary: Option<&Usage>, secondary: Option<&Usage>) -> Option { - match (primary, secondary) { - (None, None) => None, - (None, Some(usage)) | (Some(usage), None) => Some(usage.clone()), - (Some(primary), Some(secondary)) => Some(Usage { - prompt_tokens: primary.prompt_tokens.or(secondary.prompt_tokens), - completion_tokens: primary.completion_tokens.or(secondary.completion_tokens), - total_tokens: primary.total_tokens.or(secondary.total_tokens), - cache_read_tokens: primary.cache_read_tokens.or(secondary.cache_read_tokens), - cache_write_tokens: primary.cache_write_tokens.or(secondary.cache_write_tokens), - cost: primary.cost.clone().or_else(|| secondary.cost.clone()), - }), - } -} - fn merge_metrics( primary: Option, supplemental: Option<&AtifMetrics>, @@ -1024,6 +1012,7 @@ fn extract_tool_calls(output: &Json) -> Option> { let name = func .and_then(|f| f.get("name")) .or_else(|| tc_obj.get("name")) + .or_else(|| tc_obj.get("toolName")) .or_else(|| tc_obj.get("tool_name")) .or_else(|| tc_obj.get("function_name")) .and_then(Json::as_str) @@ -1324,6 +1313,7 @@ fn json_string_at(value: &Json, path: &[&str]) -> Option { struct EventLookupMaps { name_map: std::collections::HashMap, start_ts_map: std::collections::HashMap>, + llm_start_model_names: HashMap, tool_call_ids: std::collections::HashMap, suppressed_llm_events: HashSet, supplemental_llm_metrics: HashMap, @@ -1350,16 +1340,23 @@ impl EventLookupMaps { ) -> Self { let mut name_map = std::collections::HashMap::new(); let mut start_ts_map = std::collections::HashMap::new(); + let mut llm_start_model_names = HashMap::new(); for event in events { if is_start_event(event) { name_map.insert(event.uuid(), event.name().to_string()); start_ts_map.insert(event.uuid(), *event.timestamp()); + if event.category().map(|category| category.as_str()) == Some("llm") + && let Some(model_name) = model_name_for_llm_event(event) + { + llm_start_model_names.insert(event.uuid(), model_name); + } } } let llm_dedupe = build_llm_dedupe(llm_dedupe_events); Self { name_map, start_ts_map, + llm_start_model_names, tool_call_ids: build_tool_call_correlations(tool_correlation_events), suppressed_llm_events: llm_dedupe.suppressed_events, supplemental_llm_metrics: llm_dedupe.supplemental_metrics, @@ -1457,10 +1454,11 @@ impl LlmSpanCandidate { fidelity_score: llm_event_fidelity_score(start).max(llm_event_fidelity_score(end)), end_metrics: end.data().and_then(|output| { let normalized_response = end.normalized_llm_response(); + let requested_model = end.model_name().or_else(|| start.model_name()); extract_metrics( output, Some(end.name()), - end.model_name(), + requested_model, normalized_response.as_deref(), ) }), @@ -2276,6 +2274,10 @@ impl StepConversionState { let reasoning_effort = self.current_reasoning_effort.take(); let reasoning_content = extract_reasoning_content(output); let start_ts = lookups.start_ts_map.get(&event.uuid()).cloned(); + let paired_start_model = lookups + .llm_start_model_names + .get(&event.uuid()) + .map(String::as_str); let ancestry = build_ancestry(event, &lookups.name_map); let invocation = build_invocation_info( start_ts, @@ -2290,7 +2292,7 @@ impl StepConversionState { extract_metrics( output, Some(event.name()), - event.model_name(), + event.model_name().or(paired_start_model), normalized_response.as_deref(), ) }, diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index c92722486..685a6d2bf 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -56,6 +56,24 @@ pub(crate) fn estimate_cost_for_response_or_model( crate::codec::response::estimate_cost_for_provider(provider, requested_model, usage) } +pub(crate) fn merge_usage( + primary: Option<&crate::codec::response::Usage>, + secondary: Option<&crate::codec::response::Usage>, +) -> Option { + match (primary, secondary) { + (None, None) => None, + (None, Some(usage)) | (Some(usage), None) => Some(usage.clone()), + (Some(primary), Some(secondary)) => Some(crate::codec::response::Usage { + prompt_tokens: primary.prompt_tokens.or(secondary.prompt_tokens), + completion_tokens: primary.completion_tokens.or(secondary.completion_tokens), + total_tokens: primary.total_tokens.or(secondary.total_tokens), + cache_read_tokens: primary.cache_read_tokens.or(secondary.cache_read_tokens), + cache_write_tokens: primary.cache_write_tokens.or(secondary.cache_write_tokens), + cost: primary.cost.clone().or_else(|| secondary.cost.clone()), + }), + } +} + pub(crate) fn model_name_for_llm_event(event: &crate::api::event::Event) -> Option { if let Some(model_name) = event.model_name() { return Some(model_name.to_string()); diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index 3bbcd00e5..3ed3fa090 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -22,7 +22,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use super::{ estimate_cost_for_response_or_model, estimate_cost_for_response_or_requested_model, manual, - model_name_for_llm_event, + merge_usage, model_name_for_llm_event, }; use crate::api::event::{Event, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; @@ -789,24 +789,6 @@ fn end_attributes(event: &Event) -> Vec { attributes } -// Merge two usage sources field by field, preferring `primary` (codec-normalized) -// and filling gaps from `secondary` (manual scraper). This keeps provider-derived -// fields without dropping anything either source alone would have reported. -fn merge_usage(primary: Option<&Usage>, secondary: Option<&Usage>) -> Option { - match (primary, secondary) { - (None, None) => None, - (None, Some(usage)) | (Some(usage), None) => Some(usage.clone()), - (Some(primary), Some(secondary)) => Some(Usage { - prompt_tokens: primary.prompt_tokens.or(secondary.prompt_tokens), - completion_tokens: primary.completion_tokens.or(secondary.completion_tokens), - total_tokens: primary.total_tokens.or(secondary.total_tokens), - cache_read_tokens: primary.cache_read_tokens.or(secondary.cache_read_tokens), - cache_write_tokens: primary.cache_write_tokens.or(secondary.cache_write_tokens), - cost: primary.cost.clone().or_else(|| secondary.cost.clone()), - }), - } -} - fn push_llm_usage_attributes(attributes: &mut Vec, usage: Option<&Usage>) { let Some(usage) = usage else { return; @@ -1129,12 +1111,16 @@ fn raw_tool_call_id(tool_call: &Json) -> Option<&str> { fn raw_tool_call_name(tool_call: &Json) -> Option<&str> { tool_call - .get("function") - .and_then(|function| function.get("name")) + .get("name") .and_then(Json::as_str) - .or_else(|| tool_call.get("name").and_then(Json::as_str)) .or_else(|| tool_call.get("toolName").and_then(Json::as_str)) .or_else(|| tool_call.get("tool_name").and_then(Json::as_str)) + .or_else(|| { + tool_call + .get("function") + .and_then(|function| function.get("name")) + .and_then(Json::as_str) + }) .or_else(|| tool_call.get("function_name").and_then(Json::as_str)) } diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 5e6db74e9..948713a29 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -973,11 +973,13 @@ fn test_extract_metrics_merges_usage_and_token_usage_fields() { let metrics = extract_metrics( &json!({ "usage": { - "prompt_tokens": 10 + "prompt_tokens": 10, + "reasoning_tokens": 7 }, "token_usage": { "completion_tokens": 20, - "total_tokens": 30 + "total_tokens": 30, + "audio_tokens": 4 } }), None, @@ -989,6 +991,11 @@ fn test_extract_metrics_merges_usage_and_token_usage_fields() { assert_eq!(metrics.prompt_tokens, Some(10)); assert_eq!(metrics.completion_tokens, Some(20)); assert_eq!(metrics.extra.as_ref().unwrap()["total_tokens"], json!(30)); + assert_eq!( + metrics.extra.as_ref().unwrap()["reasoning_tokens"], + json!(7) + ); + assert_eq!(metrics.extra.as_ref().unwrap()["audio_tokens"], json!(4)); } #[test] @@ -1238,6 +1245,46 @@ fn test_exporter_falls_back_to_requested_model_when_response_model_unpriced() { assert_eq!(metrics.cost_usd, Some(0.000_435)); } +#[test] +fn test_paired_lifecycle_uses_start_model_for_end_metrics_pricing() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_test_pricing("requested-model"); + let _reset_guard = ResetPricingResolverGuard; + + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + let llm_uuid = Uuid::now_v7(); + + let start = event_builder(llm_uuid, EventType::Start) + .name("test") + .scope_type(ScopeType::Llm) + .input(json!({"messages": [{"role": "user", "content": "price this"}]})) + .model_name("requested-model") + .build(); + let end = event_builder(llm_uuid, EventType::End) + .name("test") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "priced response", + "model": "api-echoed-model", + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500 + } + })) + .build(); + + { + let mut state = exporter.state.lock().unwrap(); + state.events.push(start); + state.events.push(end); + } + + let trajectory = exporter.export().unwrap(); + let metrics = trajectory.steps[1].metrics.as_ref().unwrap(); + assert_eq!(metrics.cost_usd, Some(0.000_45)); +} + #[test] fn test_exporter_uses_normalized_usage_cost_before_model_pricing() { let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); @@ -2090,6 +2137,32 @@ fn test_exporter_llm_tool_calls_promoted() { assert_eq!(extra.llm_response.unwrap()["role"], json!("assistant")); } +#[test] +fn test_exporter_promotes_tool_name_alias_tool_calls() { + let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); + + let end = event_builder(Uuid::now_v7(), EventType::End) + .name("gpt-4") + .scope_type(ScopeType::Llm) + .output(json!({ + "role": "assistant", + "tool_calls": [{ + "toolName": "search_docs", + "arguments": {"query": "docs"} + }] + })) + .build(); + + exporter.state.lock().unwrap().events.push(end); + + let trajectory = exporter.export().unwrap(); + let tool_calls = trajectory.steps[0].tool_calls.as_ref().unwrap(); + assert_eq!(tool_calls.len(), 1); + assert_eq!(tool_calls[0].tool_call_id, "search_docs:1"); + assert_eq!(tool_calls[0].function_name, "search_docs"); + assert_eq!(tool_calls[0].arguments, json!({"query": "docs"})); +} + #[test] fn test_exporter_uses_annotated_message_but_raw_tool_calls() { // The annotation supplies the normalized message text (winning over the raw diff --git a/crates/core/tests/unit/observability/openinference_tests.rs b/crates/core/tests/unit/observability/openinference_tests.rs index 0401c5076..dc30bdbe4 100644 --- a/crates/core/tests/unit/observability/openinference_tests.rs +++ b/crates/core/tests/unit/observability/openinference_tests.rs @@ -1323,13 +1323,24 @@ fn openclaw_replay_tool_call_alias_fields_emit_openinference_attributes() { ScopeType::Llm, Some(json!({ "role": "assistant", - "tool_calls": [{ - "call_id": "call-search-docs", - "toolName": "search_docs", - "args": {"query": "docs"} - }], + "tool_calls": [ + { + "call_id": "call-search-docs", + "name": "top_level_search", + "toolName": "camel_search", + "function": { + "name": "nested_search", + "arguments": {"query": "docs"} + } + }, + { + "call_id": "call-read-docs", + "toolName": "read_docs", + "args": {"path": "README.md"} + } + ], "openclaw": { - "assistant_tool_call_names": ["search_docs"] + "assistant_tool_call_names": ["top_level_search", "read_docs"] } })), )); @@ -1347,13 +1358,28 @@ fn openclaw_replay_tool_call_alias_fields_emit_openinference_attributes() { assert_attr( &attributes, "llm.output_messages.0.message.tool_calls.0.tool_call.function.name", - "search_docs", + "top_level_search", ); assert_attr( &attributes, "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments", "{\"query\":\"docs\"}", ); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.1.tool_call.id", + "call-read-docs", + ); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.1.tool_call.function.name", + "read_docs", + ); + assert_attr( + &attributes, + "llm.output_messages.0.message.tool_calls.1.tool_call.function.arguments", + "{\"path\":\"README.md\"}", + ); } #[test] From 8bb700b26d02868200812a46bf5cfb351b50386b Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 11:40:18 -0700 Subject: [PATCH 5/6] fix: use start payload model for ATIF metrics Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 8 +++++-- crates/core/tests/unit/atif_tests.rs | 34 +++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index c65dbbc0c..a37797a15 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -1454,11 +1454,15 @@ impl LlmSpanCandidate { fidelity_score: llm_event_fidelity_score(start).max(llm_event_fidelity_score(end)), end_metrics: end.data().and_then(|output| { let normalized_response = end.normalized_llm_response(); - let requested_model = end.model_name().or_else(|| start.model_name()); + let requested_model = end + .model_name() + .or_else(|| start.model_name()) + .map(ToOwned::to_owned) + .or_else(|| model_name_for_llm_event(start)); extract_metrics( output, Some(end.name()), - requested_model, + requested_model.as_deref(), normalized_response.as_deref(), ) }), diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 948713a29..c4d2df9c5 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -1285,6 +1285,40 @@ fn test_paired_lifecycle_uses_start_model_for_end_metrics_pricing() { assert_eq!(metrics.cost_usd, Some(0.000_45)); } +#[test] +fn test_llm_span_candidate_uses_start_payload_model_for_metrics_pricing() { + let _pricing_guard = pricing_test_mutex().lock().unwrap(); + install_test_pricing("requested-model"); + let _reset_guard = ResetPricingResolverGuard; + + let llm_uuid = Uuid::now_v7(); + let start = event_builder(llm_uuid, EventType::Start) + .name("test") + .scope_type(ScopeType::Llm) + .input(json!({ + "model": "requested-model", + "messages": [{"role": "user", "content": "price this"}] + })) + .build(); + let end = event_builder(llm_uuid, EventType::End) + .name("test") + .scope_type(ScopeType::Llm) + .output(json!({ + "content": "priced response", + "model": "api-echoed-model", + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500 + } + })) + .build(); + + let candidate = LlmSpanCandidate::from_events(llm_uuid, &start, &end).unwrap(); + + assert_eq!(candidate.end_metrics.unwrap().cost_usd, Some(0.000_45)); +} + #[test] fn test_exporter_uses_normalized_usage_cost_before_model_pricing() { let exporter = AtifExporter::new("session-1".to_string(), make_agent_info()); From e0fa16255cfc41dac04ac7598d1d4a9ea34f74b0 Mon Sep 17 00:00:00 2001 From: mnajafian-nv Date: Thu, 25 Jun 2026 12:04:01 -0700 Subject: [PATCH 6/6] fix: prefer paired start model for ATIF pricing Signed-off-by: mnajafian-nv --- crates/core/src/observability/atif.rs | 8 ++++---- crates/core/tests/unit/atif_tests.rs | 2 ++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index a37797a15..62544aadc 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -1454,11 +1454,11 @@ impl LlmSpanCandidate { fidelity_score: llm_event_fidelity_score(start).max(llm_event_fidelity_score(end)), end_metrics: end.data().and_then(|output| { let normalized_response = end.normalized_llm_response(); - let requested_model = end + let requested_model = start .model_name() - .or_else(|| start.model_name()) .map(ToOwned::to_owned) - .or_else(|| model_name_for_llm_event(start)); + .or_else(|| model_name_for_llm_event(start)) + .or_else(|| end.model_name().map(ToOwned::to_owned)); extract_metrics( output, Some(end.name()), @@ -2296,7 +2296,7 @@ impl StepConversionState { extract_metrics( output, Some(event.name()), - event.model_name().or(paired_start_model), + paired_start_model.or_else(|| event.model_name()), normalized_response.as_deref(), ) }, diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index c4d2df9c5..3c7d971ba 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -1263,6 +1263,7 @@ fn test_paired_lifecycle_uses_start_model_for_end_metrics_pricing() { let end = event_builder(llm_uuid, EventType::End) .name("test") .scope_type(ScopeType::Llm) + .model_name("api-echoed-model") .output(json!({ "content": "priced response", "model": "api-echoed-model", @@ -1303,6 +1304,7 @@ fn test_llm_span_candidate_uses_start_payload_model_for_metrics_pricing() { let end = event_builder(llm_uuid, EventType::End) .name("test") .scope_type(ScopeType::Llm) + .model_name("api-echoed-model") .output(json!({ "content": "priced response", "model": "api-echoed-model",