diff --git a/src/fortifyroot/_vendor/opentelemetry/instrumentation/openai/retry_handler.py b/src/fortifyroot/_vendor/opentelemetry/instrumentation/openai/retry_handler.py index 851baa8..72dd3bd 100644 --- a/src/fortifyroot/_vendor/opentelemetry/instrumentation/openai/retry_handler.py +++ b/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) diff --git a/tests/openai/test_instrumentation.py b/tests/openai/test_instrumentation.py index b6102d6..ddefa55 100644 --- a/tests/openai/test_instrumentation.py +++ b/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",