FR: Copy OpenAI cache tokens to retry attempts

packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py

@@ -61,7 +61,10 @@
 )
 from opentelemetry.instrumentation.openai.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
-from opentelemetry.semconv_ai import SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY
+from 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)
 

packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py

@@ -135,7 +135,15 @@ def _make_response(status_code: int = 200, request_id: str = "req-abc",
         body = {
             "id": f"chatcmpl-{request_id}",
             "model": "gpt-4o-mini-2024-07-18",
-            "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
+            "usage": {
+                "prompt_tokens": 7,
+                "completion_tokens": 3,
+                "total_tokens": 10,
+                "prompt_tokens_details": {
+                    "cached_tokens": 5,
+                    "audio_tokens": 0,
+                },
+            },
         }
     headers = {"x-request-id": request_id, "openai-request-id": request_id}
     return SimpleNamespace(
@@ -294,6 +302,8 @@ def test_single_attempt_emits_one_span_with_marker(fresh_tracer):
     assert rs.attributes.get("gen_ai.response.model") == "gpt-4o-mini-2024-07-18"
     assert rs.attributes.get("gen_ai.usage.input_tokens") == 7
     assert rs.attributes.get("gen_ai.usage.output_tokens") == 3
+    assert rs.attributes.get("llm.usage.total_tokens") == 10
+    assert rs.attributes.get("gen_ai.usage.cache_read_input_tokens") == 5
     assert rs.attributes.get("server.address") == "api.openai.com"
     assert rs.attributes.get("server.port") == 443
     assert parent_exported.attributes.get(_FR_HAS_ATTEMPT_CHILD_KEY) is True