From f98e75241cb2396fbe57608ea9cc56dba8656a16 Mon Sep 17 00:00:00 2001 From: sayonfortify Date: Sat, 20 Jun 2026 13:41:19 +0530 Subject: [PATCH] fix(litellm): copy streaming latency attrs to attempt span --- .../instrumentation/litellm/__init__.py | 36 +++++++++++++++++++ .../litellm/streaming_safety.py | 27 ++++++++++++-- .../tests/test_retry_attempt_emission.py | 13 +++++++ .../tests/test_safety_hooks.py | 12 +++++++ 4 files changed, 85 insertions(+), 3 deletions(-) diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py index 4e4f98e85a..ad397de0e2 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py +++ b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py @@ -292,6 +292,38 @@ def _maybe_warn_retry_attempt_eviction(evicted: int) -> None: ) +def _set_active_retry_attempt_attribute(parent, key: str, value) -> None: + """Copy wrapper-owned attrs onto the canonical retry_attempt span. + + LiteLLM streaming latency is measured by FR's safety wrapper while backend + dedup makes ``fortifyroot.litellm.attempt_N`` the canonical usage span. + Storing these attrs on the attempt span avoids relying on wrapper/attempt + spans arriving in the same OTLP batch. + """ + if parent is None: + return + + with _FR_RETRY_ATTEMPT_MAP_LOCK: + entries = [ + entry + for entry in _FR_RETRY_ATTEMPT_MAP.values() + if entry.get("parent") is parent and not entry.get("ended") + ] + for entry in entries: + streaming_attrs = entry.setdefault("streaming_attrs", {}) + streaming_attrs[key] = value + span = entry.get("span") + if span is None: + continue + try: + span.set_attribute(key, value) + except Exception: + logger.debug( + "Failed to set streaming latency on retry_attempt span", + exc_info=True, + ) + + def _resolve_routed_provider(kwargs) -> Optional[str]: """Best-effort: derive the ROUTED provider (e.g. ``openai``) from LiteLLM kwargs for the retry_attempt span's ``gen_ai.system`` @@ -568,6 +600,8 @@ def _finalize_retry_attempt_span( try: if success: + for key, value in (entry.get("streaming_attrs") or {}).items(): + span.set_attribute(key, value) response_model = get_object_value(response_obj, "model") if response_model: span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, str(response_model)) @@ -882,6 +916,7 @@ def _invoke_completion(tracer, wrapped, args, kwargs, *, is_text_completion=Fals span_name, _set_response_attributes, token=span_token, + set_canonical_span_attribute=_set_active_retry_attempt_attribute, ) context_api.detach(span_token) @@ -942,6 +977,7 @@ async def _invoke_acompletion( span_name, _set_response_attributes, token=span_token, + set_canonical_span_attribute=_set_active_retry_attempt_attribute, ) context_api.detach(span_token) diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py index fea4a688dc..f350354f18 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py +++ b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py @@ -38,7 +38,15 @@ def is_async_streaming_response(kwargs, response) -> bool: return False -def wrap_sync_streaming_response(span, response, request_type, span_name, set_response_attributes, token=None): +def wrap_sync_streaming_response( + span, + response, + request_type, + span_name, + set_response_attributes, + token=None, + set_canonical_span_attribute=None, +): """Wrap a sync streaming response with per-chunk safety and span lifecycle. ``token`` is the OTel context token established by ``_invoke_completion`` @@ -66,6 +74,7 @@ def wrap_sync_streaming_response(span, response, request_type, span_name, set_re span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, first_token_at - stream_started_at, + set_canonical_span_attribute, ) yield chunk if first_token_at is not None: @@ -73,6 +82,7 @@ def wrap_sync_streaming_response(span, response, request_type, span_name, set_re span, FR_STREAMING_TIME_TO_GENERATE_MS, time.time() - first_token_at, + set_canonical_span_attribute, ) _finalize_streaming_span(span, complete_response, set_response_attributes) except Exception as exc: @@ -93,6 +103,7 @@ async def wrap_async_streaming_response( span_name, set_response_attributes, token=None, + set_canonical_span_attribute=None, ): """Wrap an async streaming response with per-chunk safety and span lifecycle. @@ -119,6 +130,7 @@ async def wrap_async_streaming_response( span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, first_token_at - stream_started_at, + set_canonical_span_attribute, ) yield chunk if first_token_at is not None: @@ -126,6 +138,7 @@ async def wrap_async_streaming_response( span, FR_STREAMING_TIME_TO_GENERATE_MS, time.time() - first_token_at, + set_canonical_span_attribute, ) _finalize_streaming_span(span, complete_response, set_response_attributes) except Exception as exc: @@ -248,9 +261,17 @@ async def _aclose_streaming_response(response) -> None: await aclose() -def _set_streaming_latency_attr(span, key: str, seconds: float) -> None: +def _set_streaming_latency_attr( + span, + key: str, + seconds: float, + set_canonical_span_attribute=None, +) -> None: + value = int(round(seconds * 1000)) if span.is_recording(): - span.set_attribute(key, int(round(seconds * 1000))) + span.set_attribute(key, value) + if set_canonical_span_attribute is not None: + set_canonical_span_attribute(span, key, value) def _chunk_has_output_text(chunk) -> bool: diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py b/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py index 3d4b962081..a07f6a8243 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_retry_attempt_emission.py @@ -38,6 +38,11 @@ _resolve_routed_provider, _start_retry_attempt_span, _finalize_retry_attempt_span, + _set_active_retry_attempt_attribute, +) +from opentelemetry.instrumentation.litellm.streaming_safety import ( + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + FR_STREAMING_TIME_TO_GENERATE_MS, ) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor @@ -237,6 +242,12 @@ def test_single_attempt_emits_one_retry_attempt_under_parent(fresh_tracer): "messages": [{"role": "user", "content": "hello masked [EMAIL]"}], } _start_retry_attempt_span(kwargs) + _set_active_retry_attempt_attribute( + parent, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 123 + ) + _set_active_retry_attempt_attribute( + parent, FR_STREAMING_TIME_TO_GENERATE_MS, 456 + ) # Mock response object with usage/model. class MockUsage: @@ -272,6 +283,8 @@ class MockResponse: assert retry_span.attributes.get("gen_ai.response.id") == "resp-abc" assert retry_span.attributes.get("gen_ai.usage.input_tokens") == 10 assert retry_span.attributes.get("gen_ai.usage.output_tokens") == 5 + assert retry_span.attributes.get(FR_STREAMING_TIME_TO_FIRST_TOKEN_MS) == 123 + assert retry_span.attributes.get(FR_STREAMING_TIME_TO_GENERATE_MS) == 456 assert retry_span.attributes.get("gen_ai.prompt.0.role") == "user" assert retry_span.attributes.get("gen_ai.prompt.0.content") == "hello masked [EMAIL]" diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py b/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py index cdbfaae1b3..37eb358567 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py @@ -769,6 +769,7 @@ def test_streaming_chunk_output_detection_handles_delta_message_and_text(): def test_sync_streaming_wrapper_sets_streaming_latency_attrs(): exporter, tracer = _test_tracer() span = tracer.start_span("litellm.completion") + canonical_attrs = [] chunks = [ SimpleNamespace( choices=[ @@ -793,18 +794,24 @@ def test_sync_streaming_wrapper_sets_streaming_latency_attrs(): "chat", "litellm.completion", lambda *_: None, + set_canonical_span_attribute=lambda *args: canonical_attrs.append(args), ) ) attrs = exporter.get_finished_spans()[0].attributes assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 250 assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 650 + assert canonical_attrs == [ + (span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 250), + (span, FR_STREAMING_TIME_TO_GENERATE_MS, 650), + ] @pytest.mark.asyncio async def test_async_streaming_wrapper_sets_streaming_latency_attrs(): exporter, tracer = _test_tracer() span = tracer.start_span("litellm.completion") + canonical_attrs = [] async def _response(): yield SimpleNamespace( @@ -823,6 +830,7 @@ async def _response(): "chat", "litellm.completion", lambda *_: None, + set_canonical_span_attribute=lambda *args: canonical_attrs.append(args), ) ] @@ -830,6 +838,10 @@ async def _response(): attrs = exporter.get_finished_spans()[0].attributes assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 125 assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 375 + assert canonical_attrs == [ + (span, FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 125), + (span, FR_STREAMING_TIME_TO_GENERATE_MS, 375), + ] def test_streaming_wrapper_leaves_latency_attrs_unset_for_empty_stream():