From 3a469b454dac05b97a318c6b9015ef301418c7c6 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Thu, 9 Jul 2026 21:20:58 +0530 Subject: [PATCH] CCM-33801: capture LiteLLM streaming responses in spans The LiteLLM instrumentation ended the litellm_request span as soon as completion()/acompletion() returned, which for stream=True is the not-yet-consumed CustomStreamWrapper. Response-side telemetry (usage, response id, finish reasons, content) is only available after the stream is iterated, so streaming spans were missing all of it. Defer span completion for streaming: wrap the returned stream in a transparent proxy (_StreamSpanWrapper) that forwards every chunk while accumulating them, then on exhaustion (or error) rebuilds the aggregated response via litellm.stream_chunk_builder, copies response metadata onto the span, and ends it. Mirrors the deferred-span pattern already used by the OpenAI/Anthropic instrumentors. Adds sync + async streaming tests asserting the span is not finished until the stream is consumed and that response metadata is captured. Co-authored-by: Cursor --- .../instrumentation/litellm/__init__.py | 169 ++++++++++++++++++ .../litellm/litellm_instrumentation_test.py | 51 ++++++ 2 files changed, 220 insertions(+) diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index a37a383..13e24fb 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -10,6 +10,14 @@ provider call. The wrapper enriches that span with response metadata before it ends. +Streaming (``stream=True``) is handled by deferring span completion: the +returned LiteLLM ``CustomStreamWrapper`` is wrapped in a transparent proxy that +forwards every chunk to the caller, accumulates them, and only when the stream +is fully consumed (or errors) rebuilds the aggregated response, copies response +metadata (usage, id, finish reasons) onto the span, and ends it. Without this +the span would close as soon as ``completion()`` returned the not-yet-consumed +stream object, losing all response-side telemetry. + Optional: ``pip install harness-sdk[litellm]`` """ @@ -463,6 +471,140 @@ def _set_response_attributes( ) +def _is_stream_response(response: Any) -> bool: + """Detect a LiteLLM streaming response (``CustomStreamWrapper``). + + Falls back to duck-typing on the async/sync iterator protocol so that the + check still works if LiteLLM renames or relocates the wrapper class. + """ + try: + from litellm import CustomStreamWrapper # pylint: disable=import-outside-toplevel + + if isinstance(response, CustomStreamWrapper): + return True + except Exception: # pylint: disable=broad-except + pass + # ModelResponse / EmbeddingResponse are not iterators; a streaming response + # exposes __anext__ (async) or __next__ (sync). + return hasattr(response, "__anext__") or hasattr(response, "__next__") + + +def _aggregate_stream_response(chunks: list[Any], messages: Any) -> Any: + """Rebuild a complete ``ModelResponse`` from streamed chunks. + + Uses ``litellm.stream_chunk_builder`` (the same helper LiteLLM uses + internally) so usage, choices, finish_reason and content are aggregated + exactly as they would be for a non-streaming call. Falls back to the last + chunk that carries usage if the builder is unavailable or fails. + """ + if not chunks: + return None + try: + import litellm # pylint: disable=import-outside-toplevel + + builder_messages = messages if isinstance(messages, list) else None + aggregated = litellm.stream_chunk_builder(chunks, messages=builder_messages) + if aggregated is not None: + return aggregated + except Exception as err: # pylint: disable=broad-except + logger.debug("LiteLLM: stream_chunk_builder failed: %s", err) + for chunk in reversed(chunks): + if _get_value(chunk, "usage") is not None: + return chunk + return chunks[-1] + + +class _StreamSpanWrapper(wrapt.ObjectProxy): + """Transparent proxy over a LiteLLM stream that defers span completion. + + Forwards every chunk to the caller unchanged while collecting them. When the + stream is exhausted (or raises) it aggregates the chunks, enriches the span + with response metadata, and ends the span. Attribute access and any other + protocol falls through to the wrapped ``CustomStreamWrapper`` via + ``wrapt.ObjectProxy``. + """ + + def __init__( + self, + wrapped: Any, + otel_logger: Any, + span: Any, + request_model: Optional[str], + messages: Any, + ) -> None: + super().__init__(wrapped) + self._self_otel_logger = otel_logger + self._self_span = span + self._self_request_model = request_model + self._self_messages = messages + self._self_chunks: list[Any] = [] + self._self_finished = False + + def __iter__(self) -> "_StreamSpanWrapper": + return self + + def __next__(self) -> Any: + try: + chunk = self.__wrapped__.__next__() + except StopIteration: + self._finalize(None) + raise + except BaseException as exc: # pylint: disable=broad-except + self._finalize(exc) + raise + self._self_chunks.append(chunk) + return chunk + + def __aiter__(self) -> "_StreamSpanWrapper": + return self + + async def __anext__(self) -> Any: + try: + chunk = await self.__wrapped__.__anext__() + except StopAsyncIteration: + self._finalize(None) + raise + except BaseException as exc: # pylint: disable=broad-except + self._finalize(exc) + raise + self._self_chunks.append(chunk) + return chunk + + def _finalize(self, exc: Optional[BaseException]) -> None: + if self._self_finished: + return + self._self_finished = True + try: + if exc is not None: + self._self_span.record_exception(exc) + self._self_span.set_status(Status(StatusCode.ERROR, str(exc))) + else: + aggregated = _aggregate_stream_response( + self._self_chunks, self._self_messages + ) + if aggregated is not None: + _set_response_attributes( + self._self_otel_logger, + self._self_span, + aggregated, + request_model=self._self_request_model, + ) + except Exception as err: # pylint: disable=broad-except + logger.debug("LiteLLM: failed to finalize stream span: %s", err) + finally: + self._self_span.end() + + def __del__(self) -> None: + # Safety net: if the consumer abandoned the stream before exhausting it, + # still end the span (with whatever chunks were collected) so it is not + # leaked/never exported. + try: + if not self._self_finished: + self._finalize(None) + except Exception: # pylint: disable=broad-except + pass + + def _build_traceable_otel_class() -> type: from litellm.integrations.opentelemetry import ( # pylint: disable=import-outside-toplevel OpenTelemetry, @@ -595,6 +737,29 @@ def set_response_attributes(self, response: Any) -> None: self.otel_logger, self.span, response, request_model=self.request_model ) + def wrap_stream(self, response: Any) -> Any: + """Hand span ownership to a stream proxy and stop managing it here. + + Detaches the active context and clears the re-dispatch guard without + ending the span, then returns a proxy that ends the span once the + stream is consumed. The context manager ``__exit__`` becomes a no-op for + the span because ``token``/``guard`` are cleared. + """ + if self.guard is not None: + _LITELLM_SPAN_ACTIVE.reset(self.guard) + self.guard = None + if self.token is not None: + otel_context.detach(self.token) + self.token = None + _, payload = _extract_model_and_input(self.args, self.kwargs) + return _StreamSpanWrapper( + response, + self.otel_logger, + self.span, + self.request_model, + payload, + ) + def __exit__( self, _exc_type: Optional[type[BaseException]], @@ -632,6 +797,8 @@ def _sync_wrapper( with _LiteLLMSpanRun(otel_logger, func_name, args, kwargs) as span_run: response = wrapped(*args, **kwargs) + if _is_stream_response(response): + return span_run.wrap_stream(response) span_run.set_response_attributes(response) return response @@ -649,6 +816,8 @@ async def _async_wrapper( with _LiteLLMSpanRun(otel_logger, func_name, args, kwargs) as span_run: response = await wrapped(*args, **kwargs) + if _is_stream_response(response): + return span_run.wrap_stream(response) span_run.set_response_attributes(response) return response diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index 328fa8f..800d19a 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -166,6 +166,57 @@ async def _fake_async(*_args, **_kwargs): assert attrs.get("gen_ai.usage.total_tokens") == 8 +def test_litellm_streaming_span_defers_until_consumed(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument + litellm_instrumentor.instrument() + stream = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + stream_options={"include_usage": True}, + mock_response="hello world", + ) + + # The span must NOT be finished before the stream is consumed: this is the + # core bug this fix addresses. + assert len(_litellm_spans(exporter.get_finished_spans())) == 0 + + chunks = list(stream) + assert len(chunks) > 0 + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.model") == "gpt-4o-mini" + assert attrs.get("gen_ai.request.streaming") == "True" + # Response-side metadata is aggregated from the streamed chunks. + assert attrs.get("gen_ai.response.finish_reasons") == "['stop']" + + +@pytest.mark.asyncio +async def test_litellm_async_streaming_span_defers_until_consumed(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument + litellm_instrumentor.instrument() + stream = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + stream_options={"include_usage": True}, + mock_response="hello world", + ) + + assert len(_litellm_spans(exporter.get_finished_spans())) == 0 + + chunks = [chunk async for chunk in stream] + assert len(chunks) > 0 + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.streaming") == "True" + assert attrs.get("gen_ai.response.finish_reasons") == "['stop']" + + def test_litellm_double_instrument_is_noop(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument with patch("litellm.main.completion", new=_fake_model_response): litellm_instrumentor.instrument()