diff --git a/netra/instrumentation/openai/utils.py b/netra/instrumentation/openai/utils.py index 740d7587..5bc6174e 100644 --- a/netra/instrumentation/openai/utils.py +++ b/netra/instrumentation/openai/utils.py @@ -137,27 +137,69 @@ def set_response_attributes(span: Span, response_dict: Dict[str, Any]) -> None: _set_response_message_attributes(span, response_dict) -def _set_usage_attributes(span: Span, usage: Dict[str, Any]) -> None: - """Helper to set usage-related attributes""" - prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") - completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") +def _first_present(usage: Dict[str, Any], *keys: str) -> Any: + """Return the first of ``keys`` whose value in ``usage`` is not None. - if prompt_tokens: - span.set_attribute(f"{SpanAttributes.LLM_USAGE_PROMPT_TOKENS}", prompt_tokens) + Used to resolve token fields that differ between OpenAI APIs but mean the + same thing (e.g. ``prompt_tokens`` in Chat Completions vs ``input_tokens`` + in the Responses API). Keys are checked in the order given, so pass the + preferred alias first. A present key holding ``0`` is returned as-is; only + missing or explicitly-None values are skipped. - if completion_tokens: - span.set_attribute(f"{SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}", completion_tokens) + Args: + usage: The usage payload to read from. + *keys: Candidate keys to try, in priority order. - if prompt_tokens_details := (usage.get("prompt_tokens_details") or usage.get("input_tokens_details")): - if cache_tokens := prompt_tokens_details.get("cached_tokens"): - span.set_attribute(f"{SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS}", cache_tokens) + Returns: + The value of the first key present with a non-None value, or None if + none of the keys are present with a non-None value. + """ + for key in keys: + if (value := usage.get(key)) is not None: + return value + return None - if completion_tokens_details := (usage.get("completion_tokens_details") or usage.get("output_tokens_details")): - if reasoning_tokens := completion_tokens_details.get("reasoning_tokens"): - span.set_attribute(f"{SpanAttributes.LLM_USAGE_REASONING_TOKENS}", reasoning_tokens) - if total_tokens := usage.get("total_tokens"): - span.set_attribute(f"{SpanAttributes.LLM_USAGE_TOTAL_TOKENS}", total_tokens) +def _set_usage_attributes(span: Span, usage: Dict[str, Any]) -> None: + """Set usage/token attributes from an OpenAI usage payload. + + Handles both the Chat Completions shape (``prompt_tokens``/``completion_tokens`` + with ``prompt_tokens_details``) and the Responses API shape + (``input_tokens``/``output_tokens`` with ``input_tokens_details``). Token + counts are compared with ``is not None`` rather than truthiness so that a + legitimate ``0`` (e.g. a cache hit that wrote nothing) is recorded instead of + silently dropped. + """ + prompt_tokens = _first_present(usage, "prompt_tokens", "input_tokens") + completion_tokens = _first_present(usage, "completion_tokens", "output_tokens") + + if prompt_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens) + + if completion_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens) + + input_tokens_details = usage.get("prompt_tokens_details") or usage.get("input_tokens_details") + if input_tokens_details: + cache_read_tokens = input_tokens_details.get("cached_tokens") + if cache_read_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS, cache_read_tokens) + + # cache_write_tokens landed in the OpenAI SDK's usage breakdown alongside GPT-5.x + # prompt caching; map it to the Anthropic-style cache-creation attribute for parity. + cache_write_tokens = input_tokens_details.get("cache_write_tokens") + if cache_write_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS, cache_write_tokens) + + output_tokens_details = usage.get("completion_tokens_details") or usage.get("output_tokens_details") + if output_tokens_details: + reasoning_tokens = output_tokens_details.get("reasoning_tokens") + if reasoning_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_REASONING_TOKENS, reasoning_tokens) + + total_tokens = usage.get("total_tokens") + if total_tokens is not None: + span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens) def _set_response_message_attributes(span: Span, response_dict: Dict[str, Any]) -> Any: diff --git a/tests/test_openai_instrumentation.py b/tests/test_openai_instrumentation.py index 9d38e5df..6a916f46 100644 --- a/tests/test_openai_instrumentation.py +++ b/tests/test_openai_instrumentation.py @@ -149,3 +149,92 @@ def test_should_suppress_instrumentation_false(self, mock_get_value): result = should_suppress_instrumentation() assert result is False + + +class TestUsageAttributes: + """Test _set_usage_attributes token capture across Chat and Responses shapes.""" + + @staticmethod + def _capture(usage): + """Run _set_usage_attributes against a recording span and return {attr: value}.""" + from netra.instrumentation.openai.utils import _set_usage_attributes + + span = Mock() + span.is_recording.return_value = True + captured: dict = {} + span.set_attribute.side_effect = lambda key, value: captured.__setitem__(key, value) + _set_usage_attributes(span, usage) + return captured + + def test_chat_completions_cache_read_and_write(self): + """Chat Completions usage maps cached/cache_write tokens to read/creation attributes.""" + from opentelemetry.semconv_ai import SpanAttributes + + attrs = self._capture( + { + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + "prompt_tokens_details": {"cached_tokens": 64, "cache_write_tokens": 30}, + "completion_tokens_details": {"reasoning_tokens": 8}, + } + ) + + assert attrs[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 100 + assert attrs[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 20 + assert attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 120 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 64 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] == 30 + assert attrs[SpanAttributes.LLM_USAGE_REASONING_TOKENS] == 8 + + def test_responses_api_cache_write(self): + """Responses API usage (input_tokens_details) also captures cache_write_tokens.""" + from opentelemetry.semconv_ai import SpanAttributes + + attrs = self._capture( + { + "input_tokens": 200, + "output_tokens": 40, + "total_tokens": 240, + "input_tokens_details": {"cached_tokens": 128, "cache_write_tokens": 50}, + "output_tokens_details": {"reasoning_tokens": 12}, + } + ) + + assert attrs[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 200 + assert attrs[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 40 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 128 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] == 50 + assert attrs[SpanAttributes.LLM_USAGE_REASONING_TOKENS] == 12 + + def test_zero_valued_cache_write_is_recorded(self): + """A cache_write_tokens of 0 must be recorded, not dropped as falsy.""" + from opentelemetry.semconv_ai import SpanAttributes + + attrs = self._capture( + { + "prompt_tokens": 10, + "completion_tokens": 0, + "total_tokens": 10, + "prompt_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + } + ) + + assert attrs[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 0 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 0 + assert attrs[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] == 0 + + def test_missing_cache_write_is_omitted(self): + """When cache_write_tokens is absent, the creation attribute is not set.""" + from opentelemetry.semconv_ai import SpanAttributes + + attrs = self._capture( + { + "prompt_tokens": 10, + "completion_tokens": 5, + "prompt_tokens_details": {"cached_tokens": 4}, + } + ) + + assert SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS not in attrs + assert attrs[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 4