fix(litellm): set ERROR status on streaming spans that fail mid-stream

python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py

@@ -1134,6 +1134,7 @@ def _finalize_sync_streaming_span(span: trace_api.Span, stream: Any) -> Any:
             _set_token_counts_from_usage(span, SimpleNamespace(usage=usage_stats))
     except Exception as e:
         span.record_exception(e)
+        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
         raise
     else:
         _set_span_status(span, aggregated_output)
@@ -1197,6 +1198,7 @@ async def _finalize_streaming_span(span: trace_api.Span, stream: Any) -> Any:
             _set_token_counts_from_usage(span, SimpleNamespace(usage=usage_stats))
     except Exception as e:
         span.record_exception(e)
+        span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, description=str(e)))
         raise
     else:
         _set_span_status(span, aggregated_output)

python/instrumentation/openinference-instrumentation-litellm/tests/test_instrumentor.py

@@ -532,6 +532,95 @@ def __next__(self) -> ModelResponseStream:
     assert attributes.get(SpanAttributes.OUTPUT_VALUE) is None
 
 
+def test_completion_sync_streaming_sets_error_status_on_mid_stream_failure(
+    in_memory_span_exporter: InMemorySpanExporter,
+    setup_litellm_instrumentation: Any,
+) -> None:
+    in_memory_span_exporter.clear()
+
+    class MidStreamError(Exception):
+        pass
+
+    class FailingStreamWrapper:
+        def __iter__(self) -> "FailingStreamWrapper":
+            return self
+
+        def __next__(self) -> Any:
+            raise MidStreamError("connection dropped")
+
+    original_func = LiteLLMInstrumentor.original_litellm_funcs["completion"]
+    try:
+        LiteLLMInstrumentor.original_litellm_funcs["completion"] = (
+            lambda *args, **kwargs: FailingStreamWrapper()
+        )
+        with patch(
+            "litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper",
+            FailingStreamWrapper,
+        ):
+            response = litellm.completion(
+                model="gpt-3.5-turbo",
+                messages=[{"content": "hi", "role": "user"}],
+                stream=True,
+            )
+            with pytest.raises(MidStreamError):
+                list(response)
+    finally:
+        LiteLLMInstrumentor.original_litellm_funcs["completion"] = original_func
+
+    spans = in_memory_span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "completion"
+    assert span.status.status_code == StatusCode.ERROR
+    assert len([e for e in span.events if e.name == "exception"]) == 1
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_sets_error_status_on_mid_stream_failure(
+    in_memory_span_exporter: InMemorySpanExporter,
+    setup_litellm_instrumentation: Any,
+) -> None:
+    in_memory_span_exporter.clear()
+
+    class MidStreamError(Exception):
+        pass
+
+    class FailingAsyncStreamWrapper:
+        def __aiter__(self) -> "FailingAsyncStreamWrapper":
+            return self
+
+        async def __anext__(self) -> Any:
+            raise MidStreamError("connection dropped")
+
+    async def fake_acompletion(*args: Any, **kwargs: Any) -> Any:
+        return FailingAsyncStreamWrapper()
+
+    original_func = LiteLLMInstrumentor.original_litellm_funcs["acompletion"]
+    try:
+        LiteLLMInstrumentor.original_litellm_funcs["acompletion"] = fake_acompletion
+        with patch(
+            "litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper",
+            FailingAsyncStreamWrapper,
+        ):
+            response = await litellm.acompletion(
+                model="gpt-3.5-turbo",
+                messages=[{"content": "hi", "role": "user"}],
+                stream=True,
+            )
+            with pytest.raises(MidStreamError):
+                async for _ in response:
+                    pass
+    finally:
+        LiteLLMInstrumentor.original_litellm_funcs["acompletion"] = original_func
+
+    spans = in_memory_span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "acompletion"
+    assert span.status.status_code == StatusCode.ERROR
+    assert len([e for e in span.events if e.name == "exception"]) == 1
+
+
 @pytest.mark.asyncio
 async def test_acompletion_streaming_with_tool_calls(
     in_memory_span_exporter: InMemorySpanExporter,