diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index c0b58a5b74..535c9d9a6b 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -681,8 +681,13 @@ class RawOpenAIChatClient( # type: ignore[misc] except Exception as ex: self._handle_request_error(ex) - response_format = validated_options.get("response_format") if validated_options else None - return self._build_response_stream(_stream(), response_format=response_format) + return ResponseStream( + _stream(), + finalizer=lambda updates: self._finalize_response_updates( + updates, + response_format=validated_options.get("response_format") if validated_options else None, + ), + ) # Non-streaming async def _get_response() -> ChatResponse: diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index abf3aa51e2..170bcbdf95 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -841,6 +841,46 @@ async def test_served_model_header_not_captured_for_streaming_text_format() -> N assert update.model == "test-model" +async def test_streaming_text_format_preserves_final_structured_output() -> None: + """Streaming structured output should still parse into the final ChatResponse value.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + events = [ + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="text_item", + output_index=0, + sequence_number=1, + logprobs=[], + delta='{"location":"Seattle","weather":"Sunny"}', + ), + ] + + fake_stream_ctx = _FakeAsyncEventStreamContext(events) + + with ( + patch.object( + client, + "_prepare_request", + new=AsyncMock( + return_value=( + client.client, + {"text_format": OutputStruct}, + {"response_format": OutputStruct}, + ) + ), + ), + patch.object(client.client.responses, "stream", return_value=fake_stream_ctx), + patch.object(client, "_get_metadata_from_response", return_value={}), + ): + stream = client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True) + response = await stream.get_final_response() + + assert response.model == "test-model" + assert response.value == OutputStruct(location="Seattle", weather="Sunny") + + async def test_bad_request_error_non_content_filter() -> None: """Test get_response BadRequestError without content_filter.""" client = OpenAIChatClient(model="test-model", api_key="test-key")