fix(openai): finalize sync stream on early break

Author: hassiebp

langfuse/openai.py

@@ -830,6 +830,30 @@ def _is_streaming_response(response: Any) -> bool:
     )
 
 
+_openai_stream_iter_hook_installed = False
+
+
+def _install_openai_stream_iteration_hooks() -> None:
+    global _openai_stream_iter_hook_installed
+
+    if not _is_openai_v1():
+        return
+
+    if not _openai_stream_iter_hook_installed:
+        original_iter = openai.Stream.__iter__
+
+        def traced_iter(self: Any) -> Any:
+            try:
+                yield from original_iter(self)
+            finally:
+                finalize_once = getattr(self, "_langfuse_finalize_once", None)
+                if finalize_once is not None:
+                    finalize_once()
+
+        openai.Stream.__iter__ = traced_iter
+        _openai_stream_iter_hook_installed = True
+
+
 def _finalize_stream_response(
     *,
     resource: OpenAiDefinition,
@@ -858,21 +882,6 @@ def _finalize_stream_response(
         generation.end()
 
 
-async def _finalize_stream_response_async(
-    *,
-    resource: OpenAiDefinition,
-    items: list[Any],
-    generation: LangfuseGeneration,
-    completion_start_time: Optional[datetime],
-) -> None:
-    _finalize_stream_response(
-        resource=resource,
-        items=items,
-        generation=generation,
-        completion_start_time=completion_start_time,
-    )
-
-
 def _instrument_openai_stream(
     *,
     resource: OpenAiDefinition,
@@ -905,6 +914,8 @@ def finalize_once() -> None:
             completion_start_time=completion_start_time,
         )
 
+    response._langfuse_finalize_once = finalize_once  # type: ignore[attr-defined]
+
     def traced_iterator() -> Any:
         nonlocal completion_start_time
         try:
@@ -955,7 +966,7 @@ async def finalize_once() -> None:
             return
 
         is_finalized = True
-        await _finalize_stream_response_async(
+        _finalize_stream_response(
             resource=resource,
             items=items,
             generation=generation,
@@ -1167,6 +1178,7 @@ def register_tracing() -> None:
 
 
 register_tracing()
+_install_openai_stream_iteration_hooks()
 
 
 class LangfuseResponseGeneratorSync:
@@ -1275,7 +1287,7 @@ async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None
         pass
 
     async def _finalize(self) -> None:
-        await _finalize_stream_response_async(
+        _finalize_stream_response(
             resource=self.resource,
             items=self.items,
             generation=self.generation,

tests/unit/test_openai.py

@@ -272,6 +272,38 @@ def test_openai_stream_preserves_original_stream_contract(
     }
 
 
+def test_openai_stream_break_still_finalizes_generation(
+    langfuse_memory_client, get_span
+):
+    openai_client = lf_openai.OpenAI(api_key="test")
+    raw_response = DummySyncResponse()
+    raw_stream = DummyOpenAIStream(_make_chat_stream_chunks(), raw_response)
+
+    with patch.object(openai_client.chat.completions, "_post", return_value=raw_stream):
+        stream = openai_client.chat.completions.create(
+            name="unit-openai-native-stream-break",
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "1 + 1 = ?"}],
+            temperature=0,
+            stream=True,
+        )
+
+    for chunk in stream:
+        assert chunk.choices[0].delta.content == "2"
+        break
+
+    assert raw_response.closed is False
+
+    langfuse_memory_client.flush()
+    span = get_span("unit-openai-native-stream-break")
+
+    assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2"
+    assert (
+        span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME]
+        is not None
+    )
+
+
 @pytest.mark.asyncio
 async def test_async_chat_completion_exports_generation_span(
     langfuse_memory_client, get_span, json_attr