From 636a3bc70c539fea1e4351a30583d9ca567bb70d Mon Sep 17 00:00:00 2001 From: sayonfortify Date: Wed, 17 Jun 2026 01:45:35 +0530 Subject: [PATCH] FR: Copy OpenAI cache tokens to retry attempts --- .../instrumentation/openai/retry_handler.py | 16 +++++++++++++++- .../tests/test_retry_attempt_emission.py | 12 +++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py index b30eb38198..a3e11f115f 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/retry_handler.py +++ b/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) diff --git a/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py index 046f8ac197..fb578afc2d 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/test_retry_attempt_emission.py +++ b/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