[NET-1379] feat: Add support for cache-write tokens in OpenAI

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:

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