Fix OpenAIGenerator and OpenAIChatGenerator to allow wrapped streaming objects usage (#9304)

* Fix for handling wrapped ChatCompletion instances in streaming (used by tools like weave)

* Add release note

* Applied same fix to OpenAIGenerator ; Refactoring ; Update release note

* Fix integration test error after refactoring

Author: mpangrazzi

haystack/components/generators/chat/openai.py

@@ -274,10 +274,7 @@ def run(
             **api_args
         )
 
-        is_streaming = isinstance(chat_completion, Stream)
-        assert is_streaming or streaming_callback is None
-
-        if is_streaming:
+        if streaming_callback is not None:
             completions = self._handle_stream_response(
                 chat_completion,  # type: ignore
                 streaming_callback,  # type: ignore
@@ -353,10 +350,7 @@ async def run_async(
             AsyncStream[ChatCompletionChunk], ChatCompletion
         ] = await self.async_client.chat.completions.create(**api_args)
 
-        is_streaming = isinstance(chat_completion, AsyncStream)
-        assert is_streaming or streaming_callback is None
-
-        if is_streaming:
+        if streaming_callback is not None:
             completions = await self._handle_async_stream_response(
                 chat_completion,  # type: ignore
                 streaming_callback,  # type: ignore

haystack/components/generators/openai.py

@@ -224,22 +224,25 @@ def run(
         )
 
         completions: List[ChatMessage] = []
-        if isinstance(completion, Stream):
+        if streaming_callback is not None:
             num_responses = generation_kwargs.pop("n", 1)
             if num_responses > 1:
                 raise ValueError("Cannot stream multiple responses, please set n=1.")
             chunks: List[StreamingChunk] = []
-            completion_chunk: Optional[ChatCompletionChunk] = None
-
-            # pylint: disable=not-an-iterable
-            for completion_chunk in completion:
-                if completion_chunk.choices and streaming_callback:
-                    chunk_delta: StreamingChunk = self._build_chunk(completion_chunk)
-                    chunks.append(chunk_delta)
-                    streaming_callback(chunk_delta)  # invoke callback with the chunk_delta
-            # Makes type checkers happy
-            assert completion_chunk is not None
-            completions = [self._create_message_from_chunks(completion_chunk, chunks)]
+            last_chunk: Optional[ChatCompletionChunk] = None
+
+            for chunk in completion:
+                if isinstance(chunk, ChatCompletionChunk):
+                    last_chunk = chunk
+
+                    if chunk.choices:
+                        chunk_delta: StreamingChunk = self._build_chunk(chunk)
+                        chunks.append(chunk_delta)
+                        streaming_callback(chunk_delta)
+
+            assert last_chunk is not None
+
+            completions = [self._create_message_from_chunks(last_chunk, chunks)]
         elif isinstance(completion, ChatCompletion):
             completions = [self._build_message(completion, choice) for choice in completion.choices]
 

releasenotes/notes/fix-openai-generators-wrapped-streaming-f3a08f83a7528c3f.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Fix an issue where OpenAIChatGenerator and OpenAIGenerator were not properly handling wrapped streaming responses from tools like Weave.

test/components/generators/chat/test_openai.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
 #
 # SPDX-License-Identifier: Apache-2.0
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock, AsyncMock
 import pytest
 
 
@@ -364,6 +364,40 @@ def streaming_callback(chunk: StreamingChunk) -> None:
         assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
         assert "Hello" in response["replies"][0].text  # see openai_mock_chat_completion_chunk
 
+    def test_run_with_wrapped_stream_simulation(self, chat_messages, openai_mock_stream):
+        streaming_callback_called = False
+
+        def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+            assert isinstance(chunk, StreamingChunk)
+
+        chunk = ChatCompletionChunk(
+            id="id",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[chat_completion_chunk.Choice(index=0, delta=chat_completion_chunk.ChoiceDelta(content="Hello"))],
+            created=int(datetime.now().timestamp()),
+        )
+
+        # Here we wrap the OpenAI stream in a MagicMock
+        # This is to simulate the behavior of some tools like Weave (https://github.com/wandb/weave)
+        # which wrap the OpenAI stream in their own stream
+        wrapped_openai_stream = MagicMock()
+        wrapped_openai_stream.__iter__.return_value = iter([chunk])
+
+        component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
+
+        with patch.object(
+            component.client.chat.completions, "create", return_value=wrapped_openai_stream
+        ) as mock_create:
+            response = component.run(chat_messages, streaming_callback=streaming_callback)
+
+            mock_create.assert_called_once()
+            assert streaming_callback_called
+            assert "replies" in response
+            assert "Hello" in response["replies"][0].text
+
     def test_check_abnormal_completions(self, caplog):
         caplog.set_level(logging.INFO)
         component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))

test/components/generators/chat/test_openai_async.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
 #
 # SPDX-License-Identifier: Apache-2.0
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, patch, MagicMock
 
 from openai import AsyncOpenAI, OpenAIError
 import pytest
@@ -365,3 +365,42 @@ async def test_live_run_with_tools_async(self, tools):
         assert tool_call.tool_name == "weather"
         assert tool_call.arguments == {"city": "Paris"}
         assert message.meta["finish_reason"] == "tool_calls"
+
+    @pytest.mark.asyncio
+    async def test_run_with_wrapped_stream_simulation_async(self, chat_messages, openai_mock_stream_async):
+        streaming_callback_called = False
+
+        async def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+            assert isinstance(chunk, StreamingChunk)
+
+        chunk = ChatCompletionChunk(
+            id="id",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[chat_completion_chunk.Choice(index=0, delta=chat_completion_chunk.ChoiceDelta(content="Hello"))],
+            created=int(datetime.now().timestamp()),
+        )
+
+        # Here we wrap the OpenAI async stream in an AsyncMock
+        # This is to simulate the behavior of some tools like Weave (https://github.com/wandb/weave)
+        # which wrap the OpenAI async stream in their own stream
+        wrapped_openai_async_stream = AsyncMock()
+        wrapped_openai_async_stream.__aiter__.return_value = iter([chunk])
+
+        component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
+
+        # Patch the async client's create method
+        with patch.object(
+            component.async_client.chat.completions,
+            "create",
+            return_value=wrapped_openai_async_stream,
+            new_callable=AsyncMock,
+        ) as mock_create:
+            response = await component.run_async(chat_messages, streaming_callback=streaming_callback)
+
+            mock_create.assert_called_once()
+            assert streaming_callback_called
+            assert "replies" in response
+            assert "Hello" in response["replies"][0].text

test/components/generators/test_openai.py

@@ -8,6 +8,8 @@
 
 import pytest
 from openai import OpenAIError
+from openai.types.chat import ChatCompletionChunk, chat_completion_chunk
+from unittest.mock import MagicMock, patch
 
 from haystack.components.generators import OpenAIGenerator
 from haystack.components.generators.utils import print_streaming_chunk
@@ -359,3 +361,42 @@ def __call__(self, chunk: StreamingChunk) -> None:
 
         assert callback.counter > 1
         assert "Paris" in callback.responses
+
+    def test_run_with_wrapped_stream_simulation(self, openai_mock_stream):
+        streaming_callback_called = False
+
+        def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+            assert isinstance(chunk, StreamingChunk)
+
+        chunk = ChatCompletionChunk(
+            id="id",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[
+                chat_completion_chunk.Choice(
+                    index=0, delta=chat_completion_chunk.ChoiceDelta(content="Hello"), finish_reason="stop"
+                )
+            ],
+            created=int(datetime.now().timestamp()),
+        )
+
+        # Here we wrap the OpenAI stream in a MagicMock
+        # This is to simulate the behavior of some tools like Weave (https://github.com/wandb/weave)
+        # which wrap the OpenAI stream in their own stream
+        wrapped_openai_stream = MagicMock()
+        wrapped_openai_stream.__iter__.return_value = iter([chunk])
+
+        component = OpenAIGenerator(api_key=Secret.from_token("test-api-key"))
+
+        with patch.object(
+            component.client.chat.completions, "create", return_value=wrapped_openai_stream
+        ) as mock_create:
+            response = component.run(prompt="test prompt", streaming_callback=streaming_callback)
+
+            mock_create.assert_called_once()
+            assert streaming_callback_called
+            assert "replies" in response
+            assert "Hello" in response["replies"][0]
+            assert response["meta"][0]["finish_reason"] == "stop"