fix(openai): finalize stream exits

Author: hassiebp

langfuse/openai.py

@@ -21,7 +21,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime
-from inspect import isclass
+from inspect import isawaitable, isclass
 from typing import Any, Optional, cast
 
 from openai._types import NotGiven
@@ -841,6 +841,7 @@ def _install_openai_stream_iteration_hooks() -> None:
 
     if not _openai_stream_iter_hook_installed:
         original_iter = openai.Stream.__iter__
+        original_aiter = openai.AsyncStream.__aiter__
 
         def traced_iter(self: Any) -> Any:
             try:
@@ -850,7 +851,17 @@ def traced_iter(self: Any) -> Any:
                 if finalize_once is not None:
                     finalize_once()
 
+        async def traced_aiter(self: Any) -> Any:
+            try:
+                async for item in original_aiter(self):
+                    yield item
+            finally:
+                finalize_once = getattr(self, "_langfuse_finalize_once", None)
+                if finalize_once is not None:
+                    await finalize_once()
+
         setattr(openai.Stream, "__iter__", traced_iter)
+        setattr(openai.AsyncStream, "__aiter__", traced_aiter)
         _openai_stream_iter_hook_installed = True
 
 
@@ -973,6 +984,8 @@ async def finalize_once() -> None:
             completion_start_time=completion_start_time,
         )
 
+    response._langfuse_finalize_once = finalize_once  # type: ignore[attr-defined]
+
     async def traced_iterator() -> Any:
         nonlocal completion_start_time
         try:
@@ -1228,7 +1241,16 @@ def __enter__(self) -> Any:
         return self.__iter__()
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        pass
+        self.close()
+
+    def close(self) -> None:
+        close = getattr(self.response, "close", None)
+
+        try:
+            if callable(close):
+                close()
+        finally:
+            self._finalize()
 
     def _finalize(self) -> None:
         if self._is_finalized:
@@ -1290,7 +1312,7 @@ async def __aenter__(self) -> Any:
         return self.__aiter__()
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        pass
+        await self.aclose()
 
     async def _finalize(self) -> None:
         if self._is_finalized:
@@ -1309,11 +1331,37 @@ async def close(self) -> None:
 
         Automatically called if the response body is read to completion.
         """
-        await self.response.close()
+        close = getattr(self.response, "close", None)
+        aclose = getattr(self.response, "aclose", None)
+
+        try:
+            if callable(close):
+                result = close()
+                if isawaitable(result):
+                    await result
+            elif callable(aclose):
+                result = aclose()
+                if isawaitable(result):
+                    await result
+        finally:
+            await self._finalize()
 
     async def aclose(self) -> None:
         """Close the response and release the connection.
 
         Automatically called if the response body is read to completion.
         """
-        await self.response.aclose()
+        aclose = getattr(self.response, "aclose", None)
+        close = getattr(self.response, "close", None)
+
+        try:
+            if callable(aclose):
+                result = aclose()
+                if isawaitable(result):
+                    await result
+            elif callable(close):
+                result = close()
+                if isawaitable(result):
+                    await result
+        finally:
+            await self._finalize()

tests/unit/test_openai.py

@@ -1,3 +1,4 @@
+import asyncio
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -51,6 +52,18 @@ def end(self) -> None:
         self.end_calls += 1
 
 
+class DummyFallbackAsyncResponse:
+    def __init__(self) -> None:
+        self.close_calls = 0
+        self.aclose_calls = 0
+
+    async def close(self) -> None:
+        self.close_calls += 1
+
+    async def aclose(self) -> None:
+        self.aclose_calls += 1
+
+
 def _make_chat_stream_chunks():
     usage = SimpleNamespace(prompt_tokens=3, completion_tokens=1, total_tokens=4)
 
@@ -469,6 +482,42 @@ async def test_openai_async_stream_supports_anext(
     }
 
 
+@pytest.mark.asyncio
+async def test_openai_async_stream_break_still_finalizes_generation(
+    langfuse_memory_client, get_span
+):
+    openai_client = lf_openai.AsyncOpenAI(api_key="test")
+    raw_stream = DummyOpenAIAsyncStream(
+        _make_chat_stream_chunks(), DummyAsyncResponse()
+    )
+
+    with patch.object(openai_client.chat.completions, "_post", return_value=raw_stream):
+        stream = await openai_client.chat.completions.create(
+            name="unit-openai-native-async-stream-break",
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "1 + 1 = ?"}],
+            temperature=0,
+            stream=True,
+        )
+
+    async for chunk in stream:
+        assert chunk.choices[0].delta.content == "2"
+        break
+
+    # Async generator finalizers are scheduled across event-loop turns.
+    for _ in range(5):
+        await asyncio.sleep(0)
+
+    langfuse_memory_client.flush()
+    span = get_span("unit-openai-native-async-stream-break")
+
+    assert span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT] == "2"
+    assert (
+        span.attributes[LangfuseOtelSpanAttributes.OBSERVATION_COMPLETION_START_TIME]
+        is not None
+    )
+
+
 def test_fallback_sync_stream_finalizes_once():
     resource = SimpleNamespace(object="Completions", type="chat")
     generation = DummyGeneration()
@@ -490,6 +539,24 @@ def fallback_stream():
     assert generation.end_calls == 1
 
 
+def test_fallback_sync_stream_exit_finalizes_once():
+    resource = SimpleNamespace(object="Completions", type="chat")
+    generation = DummyGeneration()
+
+    def fallback_stream():
+        yield _make_single_chunk_stream()
+
+    wrapper = lf_openai_module.LangfuseResponseGeneratorSync(
+        resource=resource,
+        response=fallback_stream(),
+        generation=generation,
+    )
+
+    wrapper.__exit__(None, None, None)
+
+    assert generation.end_calls == 1
+
+
 @pytest.mark.asyncio
 async def test_fallback_async_stream_finalizes_once():
     resource = SimpleNamespace(object="Completions", type="chat")
@@ -513,6 +580,45 @@ async def fallback_stream():
     assert generation.end_calls == 1
 
 
+@pytest.mark.asyncio
+async def test_fallback_async_stream_close_and_exit_finalize_once():
+    resource = SimpleNamespace(object="Completions", type="chat")
+    generation = DummyGeneration()
+    response = DummyFallbackAsyncResponse()
+
+    wrapper = lf_openai_module.LangfuseResponseGeneratorAsync(
+        resource=resource,
+        response=response,
+        generation=generation,
+    )
+
+    await wrapper.close()
+    await wrapper.__aexit__(None, None, None)
+
+    assert generation.end_calls == 1
+    assert response.close_calls == 1
+    assert response.aclose_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_fallback_async_stream_aclose_finalizes_once():
+    resource = SimpleNamespace(object="Completions", type="chat")
+    generation = DummyGeneration()
+
+    async def fallback_stream():
+        yield _make_single_chunk_stream()
+
+    wrapper = lf_openai_module.LangfuseResponseGeneratorAsync(
+        resource=resource,
+        response=fallback_stream(),
+        generation=generation,
+    )
+
+    await wrapper.aclose()
+
+    assert generation.end_calls == 1
+
+
 def test_embedding_exports_dimensions_and_count(
     langfuse_memory_client, get_span, json_attr
 ):