diff --git a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py b/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py index e122a958bf..71e13b38e1 100644 --- a/python/instrumentation/openinference-instrumentation-litellm/src/openinference/instrumentation/litellm/__init__.py +++ b/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) diff --git a/python/instrumentation/openinference-instrumentation-litellm/tests/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-litellm/tests/test_instrumentor.py index 80338e4da5..c3a632da96 100644 --- a/python/instrumentation/openinference-instrumentation-litellm/tests/test_instrumentor.py +++ b/python/instrumentation/openinference-instrumentation-litellm/tests/test_instrumentor.py @@ -1,5 +1,6 @@ import json import os +from types import SimpleNamespace from typing import Any, Dict, Generator, List, Mapping, Optional, Union, cast from unittest.mock import patch @@ -532,6 +533,144 @@ def __next__(self) -> ModelResponseStream: assert attributes.get(SpanAttributes.OUTPUT_VALUE) is None +@pytest.mark.parametrize("fail_after_chunk", [False, True]) +def test_completion_sync_streaming_sets_error_status_on_mid_stream_failure( + in_memory_span_exporter: InMemorySpanExporter, + setup_litellm_instrumentation: Any, + fail_after_chunk: bool, +) -> None: + in_memory_span_exporter.clear() + + class MidStreamError(Exception): + pass + + class FailingStreamWrapper: + def __init__(self) -> None: + self._yield_chunk = fail_after_chunk + + def __iter__(self) -> "FailingStreamWrapper": + return self + + def __next__(self) -> Any: + if self._yield_chunk: + self._yield_chunk = False + return SimpleNamespace( + choices=[ + SimpleNamespace( + index=0, + delta=SimpleNamespace(role="assistant", content="partial"), + finish_reason=None, + ) + ] + ) + 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, + ) + received = [] + with pytest.raises(MidStreamError, match="connection dropped"): + for chunk in response: + received.append(chunk) + assert not in_memory_span_exporter.get_finished_spans() + assert len(received) == int(fail_after_chunk) + 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 span.status.description == "connection dropped" + exceptions = [e for e in span.events if e.name == "exception"] + assert len(exceptions) == 1 + exception_attributes = exceptions[0].attributes + assert exception_attributes is not None + assert exception_attributes["exception.message"] == "connection dropped" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_after_chunk", [False, True]) +async def test_acompletion_streaming_sets_error_status_on_mid_stream_failure( + in_memory_span_exporter: InMemorySpanExporter, + setup_litellm_instrumentation: Any, + fail_after_chunk: bool, +) -> None: + in_memory_span_exporter.clear() + + class MidStreamError(Exception): + pass + + class FailingAsyncStreamWrapper: + def __init__(self) -> None: + self._yield_chunk = fail_after_chunk + + def __aiter__(self) -> "FailingAsyncStreamWrapper": + return self + + async def __anext__(self) -> Any: + if self._yield_chunk: + self._yield_chunk = False + return SimpleNamespace( + choices=[ + SimpleNamespace( + index=0, + delta=SimpleNamespace(role="assistant", content="partial"), + finish_reason=None, + ) + ] + ) + 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, + ) + received = [] + with pytest.raises(MidStreamError, match="connection dropped"): + async for chunk in response: + received.append(chunk) + assert not in_memory_span_exporter.get_finished_spans() + assert len(received) == int(fail_after_chunk) + 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 span.status.description == "connection dropped" + exceptions = [e for e in span.events if e.name == "exception"] + assert len(exceptions) == 1 + exception_attributes = exceptions[0].attributes + assert exception_attributes is not None + assert exception_attributes["exception.message"] == "connection dropped" + + @pytest.mark.asyncio async def test_acompletion_streaming_with_tool_calls( in_memory_span_exporter: InMemorySpanExporter,