From 81c0d30153401a4ca9dc29f712a7386941372b7d Mon Sep 17 00:00:00 2001 From: boulea7 Date: Sat, 24 Oct 2026 05:33:41 +0800 Subject: [PATCH] fix(openai): record stream iteration errors Fixes #2039 --- langfuse/openai.py | 34 +++++- tests/unit/test_openai.py | 229 +++++++++++++++++++++++++++++++++++++- 2 files changed, 257 insertions(+), 6 deletions(-) diff --git a/langfuse/openai.py b/langfuse/openai.py index 32637569b..d10e50681 100644 --- a/langfuse/openai.py +++ b/langfuse/openai.py @@ -1037,8 +1037,12 @@ def _finalize_stream_response( generation: LangfuseGeneration, completion_start_time: Optional[datetime], model_parameters: Optional[Any] = None, + error: Optional[Exception] = None, ) -> None: try: + if error is not None: + generation.update(status_message=str(error), level="ERROR") + model, completion, usage, metadata, service_tier = ( _extract_streamed_response_api_response(items) if resource.object == "Responses" or resource.object == "AsyncResponses" @@ -1083,7 +1087,7 @@ def _instrument_openai_stream( is_finalized = False close = response.close - def finalize_once() -> None: + def finalize_once(error: Optional[Exception] = None) -> None: nonlocal is_finalized if is_finalized: return @@ -1095,6 +1099,7 @@ def finalize_once() -> None: generation=generation, completion_start_time=completion_start_time, model_parameters=model_parameters, + error=error, ) response._langfuse_finalize_once = finalize_once # type: ignore[attr-defined] @@ -1109,6 +1114,9 @@ def traced_iterator() -> Any: completion_start_time = _get_timestamp() yield item + except Exception as ex: + finalize_once(error=ex) + raise finally: finalize_once() @@ -1145,7 +1153,7 @@ def _instrument_openai_async_stream( is_finalized = False close = response.close - async def finalize_once() -> None: + async def finalize_once(error: Optional[Exception] = None) -> None: nonlocal is_finalized if is_finalized: return @@ -1157,6 +1165,7 @@ async def finalize_once() -> None: generation=generation, completion_start_time=completion_start_time, model_parameters=model_parameters, + error=error, ) response._langfuse_finalize_once = finalize_once # type: ignore[attr-defined] @@ -1171,6 +1180,9 @@ async def traced_iterator() -> Any: completion_start_time = _get_timestamp() yield item + except Exception as ex: + await finalize_once(error=ex) + raise finally: await finalize_once() @@ -1471,6 +1483,9 @@ def __iter__(self) -> Any: self.completion_start_time = _get_timestamp() yield i + except Exception as ex: + self._finalize(error=ex) + raise finally: self._finalize() @@ -1488,6 +1503,9 @@ def __next__(self) -> Any: self._finalize() raise + except Exception as ex: + self._finalize(error=ex) + raise def __enter__(self) -> Any: return self.__iter__() @@ -1504,7 +1522,7 @@ def close(self) -> None: finally: self._finalize() - def _finalize(self) -> None: + def _finalize(self, error: Optional[Exception] = None) -> None: if self._is_finalized: return @@ -1515,6 +1533,7 @@ def _finalize(self) -> None: generation=self.generation, completion_start_time=self.completion_start_time, model_parameters=self.model_parameters, + error=error, ) @@ -1545,6 +1564,9 @@ async def __aiter__(self) -> Any: self.completion_start_time = _get_timestamp() yield i + except Exception as ex: + await self._finalize(error=ex) + raise finally: await self._finalize() @@ -1562,6 +1584,9 @@ async def __anext__(self) -> Any: await self._finalize() raise + except Exception as ex: + await self._finalize(error=ex) + raise async def __aenter__(self) -> Any: return self.__aiter__() @@ -1569,7 +1594,7 @@ async def __aenter__(self) -> Any: async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: await self.aclose() - async def _finalize(self) -> None: + async def _finalize(self, error: Optional[Exception] = None) -> None: if self._is_finalized: return @@ -1580,6 +1605,7 @@ async def _finalize(self) -> None: generation=self.generation, completion_start_time=self.completion_start_time, model_parameters=self.model_parameters, + error=error, ) async def close(self) -> None: diff --git a/tests/unit/test_openai.py b/tests/unit/test_openai.py index a2a4f7f4f..0594bf062 100644 --- a/tests/unit/test_openai.py +++ b/tests/unit/test_openai.py @@ -49,8 +49,10 @@ async def _stream(self, items): class DummyGeneration: def __init__(self) -> None: self.end_calls = 0 + self.updates = [] def update(self, **kwargs): + self.updates.append(kwargs) return self def end(self) -> None: @@ -568,6 +570,7 @@ def fake_stream(): span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME] is not None ) + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" assert span.attributes["langfuse.observation.metadata.finish_reason"] == '"stop"' assert json_attr(span, LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS) == { "prompt_tokens": 3, @@ -632,6 +635,7 @@ def test_openai_stream_preserves_original_stream_contract( span = get_span("unit-openai-native-stream") assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" assert ( span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME] is not None @@ -705,6 +709,7 @@ def test_openai_stream_break_still_finalizes_generation( span = get_span("unit-openai-native-stream-break") assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" assert ( span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME] is not None @@ -791,6 +796,7 @@ async def test_openai_async_stream_preserves_original_stream_contract( span = get_span("unit-openai-native-async-stream") assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" assert ( span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME] is not None @@ -880,6 +886,44 @@ async def test_openai_async_stream_break_still_finalizes_generation( span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME] is not None ) + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" + + +@pytest.mark.asyncio +async def test_openai_async_stream_cancellation_does_not_mark_generation_error( + langfuse_memory_client, get_span +): + cancelled = asyncio.CancelledError() + + def chunks_then_cancel(): + yield _make_single_chunk_stream() + raise cancelled + + raw_stream = DummyOpenAIAsyncStream(chunks_then_cancel(), DummyAsyncResponse()) + async with lf_openai.AsyncOpenAI(api_key="test") as openai_client: + with patch.object( + openai_client.chat.completions, "_post", return_value=raw_stream + ): + stream = await openai_client.chat.completions.create( + name="unit-openai-stream-cancelled", + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + stream=True, + ) + + with pytest.raises(asyncio.CancelledError) as error: + async for _ in stream: + pass + + await stream.aclose() + + assert error.value is cancelled + langfuse_memory_client.flush() + span = get_span("unit-openai-stream-cancelled") + + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) != "ERROR" + assert LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE not in span.attributes def test_fallback_sync_stream_finalizes_once(): @@ -901,6 +945,7 @@ def fallback_stream(): next(wrapper) assert generation.end_calls == 1 + assert all(update.get("level") != "ERROR" for update in generation.updates) def test_fallback_sync_stream_exit_finalizes_once(): @@ -942,6 +987,7 @@ async def fallback_stream(): await wrapper.__anext__() assert generation.end_calls == 1 + assert all(update.get("level") != "ERROR" for update in generation.updates) @pytest.mark.asyncio @@ -983,6 +1029,107 @@ async def fallback_stream(): assert generation.end_calls == 1 +@pytest.mark.parametrize("direct_next", [False, True], ids=["iterate", "next"]) +def test_fallback_sync_stream_error_preserves_partial_response( + direct_next, langfuse_memory_client, get_span, json_attr +): + error = RuntimeError("stream failed") + chunks = _make_chat_stream_chunks() + + def response(): + yield from chunks + raise error + + generation = langfuse_memory_client.start_observation( + as_type="generation", name="unit-fallback-stream-error" + ) + wrapper = lf_openai_module.LangfuseResponseGeneratorSync( + resource=SimpleNamespace(object="Completions", type="chat"), + response=response(), + generation=generation, + ) + received = [] + + with patch.object(generation, "end", wraps=generation.end) as end: + with pytest.raises(RuntimeError) as caught: + if direct_next: + while True: + received.append(next(wrapper)) + else: + received.extend(wrapper) + + wrapper.close() + assert end.call_count == 1 + + assert caught.value is error + assert received == chunks + langfuse_memory_client.flush() + span = get_span("unit-fallback-stream-error") + + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_LEVEL] == "ERROR" + assert span.attributes[ + LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE + ] == str(error) + assert json_attr(span, LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS) == { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("direct_next", [False, True], ids=["iterate", "anext"]) +async def test_fallback_async_stream_error_preserves_partial_response( + direct_next, langfuse_memory_client, get_span, json_attr +): + error = RuntimeError("stream failed") + chunks = _make_chat_stream_chunks() + + async def response(): + for chunk in chunks: + yield chunk + raise error + + generation = langfuse_memory_client.start_observation( + as_type="generation", name="unit-fallback-async-stream-error" + ) + wrapper = lf_openai_module.LangfuseResponseGeneratorAsync( + resource=SimpleNamespace(object="Completions", type="chat"), + response=response(), + generation=generation, + ) + received = [] + + with patch.object(generation, "end", wraps=generation.end) as end: + with pytest.raises(RuntimeError) as caught: + if direct_next: + while True: + received.append(await wrapper.__anext__()) + else: + async for chunk in wrapper: + received.append(chunk) + + await wrapper.aclose() + assert end.call_count == 1 + + assert caught.value is error + assert received == chunks + langfuse_memory_client.flush() + span = get_span("unit-fallback-async-stream-error") + + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2" + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_LEVEL] == "ERROR" + assert span.attributes[ + LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE + ] == str(error) + assert json_attr(span, LangfuseOtelSpanAttributes.OBSERVATION_USAGE_DETAILS) == { + "prompt_tokens": 3, + "completion_tokens": 1, + "total_tokens": 4, + } + + def test_embedding_exports_dimensions_and_count( langfuse_memory_client, get_span, json_attr ): @@ -1403,14 +1550,31 @@ def _chat_completion_chunk_sse_body(): ) -def _mock_transport_openai_client(async_client: bool = False): +def _chat_completion_stream_error_sse_body(): + return ( + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk",' + '"created":1700000000,"model":"gpt-4o-mini-2024-07-18",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"},' + '"finish_reason":null}]}\n\n' + 'data: {"error":{"message":"stream failed","type":"server_error",' + '"code":"server_error"}}\n\n' + ) + + +def _mock_transport_openai_client( + async_client: bool = False, sse_body: str | None = None +): import httpx def handler(request: httpx.Request) -> httpx.Response: if b'"stream": true' in request.content or b'"stream":true' in request.content: return httpx.Response( 200, - content=_chat_completion_chunk_sse_body().encode(), + content=( + sse_body + if sse_body is not None + else _chat_completion_chunk_sse_body() + ).encode(), headers={"content-type": "text/event-stream"}, ) @@ -1428,6 +1592,67 @@ def handler(request: httpx.Request) -> httpx.Response: ) +def test_openai_stream_error_marks_generation_error(langfuse_memory_client, get_span): + chunks = [] + with _mock_transport_openai_client( + sse_body=_chat_completion_stream_error_sse_body() + ) as openai_client: + stream = openai_client.chat.completions.create( + name="unit-openai-stream-error", + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + stream=True, + ) + + with pytest.raises(lf_openai.APIError, match="stream failed") as error: + for chunk in stream: + chunks.append(chunk) + + stream.close() + + assert [chunk.choices[0].delta.content for chunk in chunks] == ["Hel"] + langfuse_memory_client.flush() + span = get_span("unit-openai-stream-error") + + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "Hel" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) == "ERROR" + assert span.attributes[ + LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE + ] == str(error.value) + + +@pytest.mark.asyncio +async def test_openai_async_stream_error_marks_generation_error( + langfuse_memory_client, get_span +): + chunks = [] + async with _mock_transport_openai_client( + async_client=True, sse_body=_chat_completion_stream_error_sse_body() + ) as openai_client: + stream = await openai_client.chat.completions.create( + name="unit-openai-async-stream-error", + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + stream=True, + ) + + with pytest.raises(lf_openai.APIError, match="stream failed") as error: + async for chunk in stream: + chunks.append(chunk) + + await stream.aclose() + + assert [chunk.choices[0].delta.content for chunk in chunks] == ["Hel"] + langfuse_memory_client.flush() + span = get_span("unit-openai-async-stream-error") + + assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "Hel" + assert span.attributes.get(LangfuseOtelSpanAttributes.OBSERVATION_LEVEL) == "ERROR" + assert span.attributes[ + LangfuseOtelSpanAttributes.OBSERVATION_STATUS_MESSAGE + ] == str(error.value) + + def test_with_raw_response_chat_completion_captures_output_and_usage( langfuse_memory_client, get_span, json_attr ):