fix(openai): copy cache tokens onto retry attempts

src/fortifyroot/_vendor/opentelemetry/instrumentation/openai/retry_handler.py

@@ -61,7 +61,10 @@
 )
 from fortifyroot._vendor.opentelemetry.instrumentation.openai.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
-from fortifyroot._vendor.opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+from fortifyroot._vendor.opentelemetry.semconv_ai import (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    SpanAttributes,
+)
 from opentelemetry.trace import SpanKind, Status, StatusCode
 from wrapt import wrap_function_wrapper
 
@@ -413,6 +416,17 @@ def _extract_usage_from_body(span: "trace.Span", response: Any) -> None:
                 span.set_attribute("gen_ai.usage.input_tokens", pt)
             if isinstance(ct, int):
                 span.set_attribute("gen_ai.usage.output_tokens", ct)
+            total_tokens = usage.get("total_tokens")
+            if isinstance(total_tokens, int):
+                span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens)
+            prompt_tokens_details = usage.get("prompt_tokens_details")
+            if isinstance(prompt_tokens_details, dict):
+                cached_tokens = prompt_tokens_details.get("cached_tokens")
+                if isinstance(cached_tokens, int):
+                    span.set_attribute(
+                        SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
+                        cached_tokens,
+                    )
     except Exception:
         logger.debug("failed to extract usage from openai response body", exc_info=True)
 

tests/openai/test_instrumentation.py

@@ -24,6 +24,9 @@
 from fortifyroot._vendor.opentelemetry.instrumentation.openai.shared.config import (
     Config as OpenAIConfig,
 )
+from fortifyroot._vendor.opentelemetry.instrumentation.openai.retry_handler import (
+    _extract_usage_from_body,
+)
 from fortifyroot._vendor.opentelemetry.instrumentation.openai.utils import (
     should_send_prompts,
 )
@@ -152,6 +155,50 @@ def _assert_no_content_attrs(span) -> None:
     assert all(key not in span.attributes for key in _content_attr_keys())
 
 
+class _FakeHTTPResponse:
+    def __init__(self, body: dict):
+        self._body = body
+
+    def json(self) -> dict:
+        return self._body
+
+
+class _FakeSpan:
+    def __init__(self):
+        self.attributes: dict[str, object] = {}
+
+    def set_attribute(self, key: str, value: object) -> None:
+        self.attributes[key] = value
+
+
+def test_retry_attempt_extracts_cache_tokens_from_response_body():
+    span = _FakeSpan()
+
+    _extract_usage_from_body(
+        span,
+        _FakeHTTPResponse(
+            {
+                "id": "chatcmpl-cache-test",
+                "model": "gpt-4o-mini",
+                "usage": {
+                    "prompt_tokens": 2909,
+                    "completion_tokens": 8,
+                    "total_tokens": 2917,
+                    "prompt_tokens_details": {
+                        "cached_tokens": 2816,
+                        "audio_tokens": 0,
+                    },
+                },
+            }
+        ),
+    )
+
+    assert span.attributes[GenAI.GEN_AI_USAGE_INPUT_TOKENS] == 2909
+    assert span.attributes[GenAI.GEN_AI_USAGE_OUTPUT_TOKENS] == 8
+    assert span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 2917
+    assert span.attributes[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 2816
+
+
 def test_chat_completion_creates_span(init_openai_sdk, span_exporter):
     with patch(
         "openai.resources.chat.completions.Completions.create",