From 5627b74a6a11bbd9aaf0d4da8b42b2b53a1a679d Mon Sep 17 00:00:00 2001 From: Wilson0511 <1795618152@qq.com> Date: Wed, 28 Oct 2026 12:56:52 +0800 Subject: [PATCH 1/3] Python: Preserve per-choice message boundaries when streaming chat completions (n>1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With n>1, chat completion chunks interleave deltas from independent alternative completions. The streaming parser merged every choice of a chunk into one update tagged with the shared chunk id, so ChatResponse.from_updates concatenated the alternatives' texts in arrival order into a single message, while the non-streaming parser emits one assistant Message per choice. Parse each chunk into one update per choice (choice-scoped message id for choices beyond the first), and group the accumulated updates by message id at finalize time so each choice aggregates into its own message. The live stream keeps provider arrival order; n=1 behavior is unchanged. 🤖 Generated with ZCode Co-Authored-By: ZCode --- .../_chat_completion_client.py | 181 ++++++++++++++---- .../test_openai_chat_completion_client.py | 88 +++++++++ 2 files changed, 231 insertions(+), 38 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index 27da4fea330..c19a3e1d955 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -669,25 +669,25 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: async for chunk in completion_stream: if len(chunk.choices) == 0 and chunk.usage is None: continue - update = self._parse_response_update_from_openai(chunk) - for content in update.contents: - if content.type != "function_call": - continue - choice_index = content.additional_properties.get("tool_call_choice_index") - tool_index = content.additional_properties.get("tool_call_index") - if not isinstance(choice_index, int) or not isinstance(tool_index, int): - continue - index_key = (choice_index, tool_index) - identity = tool_call_identities.get(index_key) - if identity is None: - identity = (f"af-call-{uuid4().hex}", content.call_id or "") - occurrence_id, provider_call_id = identity - if content.call_id: - provider_call_id = content.call_id - tool_call_identities[index_key] = (occurrence_id, provider_call_id) - content.id = occurrence_id - content.call_id = provider_call_id - yield update + for update in self._parse_response_updates_from_openai(chunk): + for content in update.contents: + if content.type != "function_call": + continue + choice_index = content.additional_properties.get("tool_call_choice_index") + tool_index = content.additional_properties.get("tool_call_index") + if not isinstance(choice_index, int) or not isinstance(tool_index, int): + continue + index_key = (choice_index, tool_index) + identity = tool_call_identities.get(index_key) + if identity is None: + identity = (f"af-call-{uuid4().hex}", content.call_id or "") + occurrence_id, provider_call_id = identity + if content.call_id: + provider_call_id = content.call_id + tool_call_identities[index_key] = (occurrence_id, provider_call_id) + content.id = occurrence_id + content.call_id = provider_call_id + yield update except BadRequestError as ex: if ex.code == "content_filter": raise OpenAIContentFilterException( @@ -928,39 +928,69 @@ def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping response_format=options.get("response_format"), ) - def _parse_response_update_from_openai( + def _parse_response_updates_from_openai( self, chunk: ChatCompletionChunk, - ) -> ChatResponseUpdate: - """Parse a streaming response update from OpenAI.""" - chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk) - contents: list[Content] = [] - finish_reason: FinishReason | None = None - + ) -> list[ChatResponseUpdate]: + """Parse one streaming chunk into one update per choice. + + Each choice of a chunk is an independent alternative completion (the ``n`` + request parameter), so a chunk with ``n>1`` interleaves deltas that belong + to different assistant messages. Emitting one update per choice, each + carrying a choice-scoped ``message_id``, lets ``ChatResponse.from_updates`` + keep the same per-choice message boundaries that the non-streaming path + produces; merging them into a single update would concatenate the + alternatives' texts in arrival order instead. + """ + base_metadata = self._get_metadata_from_streaming_chat_response(chunk) # Process usage data (may coexist with text/tool content in providers like Gemini). # See https://github.com/microsoft/agent-framework/issues/3434 - if chunk.usage: - contents.append( + usage_contents: list[Content] = ( + [ Content.from_usage( usage_details=self._parse_usage_from_openai(chunk.usage), merge_mode="replace", raw_representation=chunk, ) - ) + ] + if chunk.usage + else [] + ) - for choice in chunk.choices: + def _choice_update(choice: ChunkChoice, contents: list[Content]) -> ChatResponseUpdate: + choice_metadata = dict(base_metadata) # Missing per-chunk logprobs must not clear the most recent token metadata. if choice.logprobs is not None: - chunk_metadata.update(self._get_metadata_from_chat_choice(choice)) + choice_metadata.update(self._get_metadata_from_chat_choice(choice)) + finish_reason: FinishReason | None = None if choice.finish_reason: - finish_reason = "tool_calls" if choice.finish_reason == "function_call" else choice.finish_reason # type: ignore[assignment] + finish_reason = ( + "tool_calls" if choice.finish_reason == "function_call" else choice.finish_reason # type: ignore[assignment] + ) + # The provider chunk id is shared by every choice, so only choice 0 can + # use it as its message id; the other choices get a suffixed id to keep + # their message boundaries distinct. + message_id = chunk.id if choice.index == 0 else f"{chunk.id}-{choice.index}" + return ChatResponseUpdate( + created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + contents=contents, + role="assistant", + model=chunk.model, + additional_properties=choice_metadata, + finish_reason=finish_reason, + raw_representation=chunk, + response_id=chunk.id, + message_id=message_id, + ) + updates: list[ChatResponseUpdate] = [] + for choice in chunk.choices: # Some OpenAI-compatible providers (e.g. Azure) send `"delta": null` # on finish chunks instead of the spec-compliant `"delta": {}`. # Guard here so all content-parsing below can assume delta is present. if choice.delta is None: # pyright: ignore[reportUnnecessaryComparison] + updates.append(_choice_update(choice, [])) continue - choice_contents: list[Content] = [] choice_contents.extend(self._parse_tool_calls_from_openai(choice)) if text_content := self._parse_text_from_openai(choice): @@ -976,19 +1006,94 @@ def _parse_response_update_from_openai( ) if self.response_parser is not None: choice_contents = list(self.response_parser(choice.delta, choice_contents)) - contents.extend(choice_contents) + updates.append(_choice_update(choice, choice_contents)) + + if not updates: + if not usage_contents: + return [] + return [ + ChatResponseUpdate( + created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + contents=usage_contents, + role="assistant", + model=chunk.model, + additional_properties=dict(base_metadata), + raw_representation=chunk, + response_id=chunk.id, + message_id=chunk.id, + ) + ] + # Usage is response-level, so it rides on the first choice's update (which is + # where a single-choice chunk carried it before). + updates[0].contents = [*usage_contents, *updates[0].contents] + return updates + + def _parse_response_update_from_openai( + self, + chunk: ChatCompletionChunk, + ) -> ChatResponseUpdate: + """Parse a streaming response update from OpenAI. + + Multi-choice chunks (``n>1``) collapse into a single merged update here for + backwards compatibility; use :meth:`_parse_response_updates_from_openai` + to preserve per-choice message boundaries. + """ + updates = self._parse_response_updates_from_openai(chunk) + if not updates: + return ChatResponseUpdate( + created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + contents=[], + role="assistant", + model=chunk.model, + additional_properties=self._get_metadata_from_streaming_chat_response(chunk), + raw_representation=chunk, + response_id=chunk.id, + message_id=chunk.id, + ) + if len(updates) == 1: + return updates[0] + merged_metadata: dict[str, Any] = {} + for update in updates: + merged_metadata.update(update.additional_properties or {}) return ChatResponseUpdate( - created_at=datetime.fromtimestamp(chunk.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), - contents=contents, + created_at=updates[0].created_at, + contents=[content for update in updates for content in update.contents], role="assistant", model=chunk.model, - additional_properties=chunk_metadata, - finish_reason=finish_reason, + additional_properties=merged_metadata, + finish_reason=next((update.finish_reason for update in reversed(updates) if update.finish_reason), None), raw_representation=chunk, response_id=chunk.id, message_id=chunk.id, ) + @override + def _finalize_response_updates( + self, + updates: Sequence[ChatResponseUpdate], + *, + response_format: Any | None = None, + ) -> ChatResponse[Any]: + """Finalize response updates into a single ChatResponse. + + With ``n>1``, the provider interleaves the choices' deltas in one stream, so + the per-choice updates arrive as e.g. ``A1, B1, A2, B2`` even though ``A1`` + and ``A2`` belong to the same assistant message. ``ChatResponse.from_updates`` + starts a new message whenever the message id changes, so the updates are + first grouped by message id (stable, first-seen order) to aggregate each + choice into its own message, matching the non-streaming path. Streams where + every update shares one id (``n=1``) or any update has no id are left as-is. + """ + message_ids = [update.message_id for update in updates] + if len(message_ids) > 1 and all(message_id is not None for message_id in message_ids): + first_seen_order = dict.fromkeys(message_ids) + if len(first_seen_order) > 1: + grouped: dict[str, list[ChatResponseUpdate]] = {message_id: [] for message_id in first_seen_order} + for update in updates: + grouped[cast("str", update.message_id)].append(update) + updates = [update for message_id in first_seen_order for update in grouped[message_id]] + return super()._finalize_response_updates(updates, response_format=response_format) + def _parse_usage_from_openai(self, usage: CompletionUsage) -> UsageDetails: details = UsageDetails( input_token_count=usage.prompt_tokens, diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index 7d5fc75b077..2270169cbe8 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -2900,6 +2900,94 @@ def test_streaming_tool_call_preserves_choice_local_index_scope( assert [content.additional_properties["tool_call_choice_index"] for content in function_calls] == [0, 1] +async def test_streaming_multiple_choices_preserve_message_boundaries( + openai_unit_test_env: dict[str, str], +) -> None: + """Text deltas from different choices (n>1) must not interleave inside one message. + + Non-streamed parsing emits one assistant Message per choice, so the streamed + aggregation must preserve the same per-choice boundaries instead of + concatenating interleaved deltas from all choices into a single message. + """ + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + + client = RawOpenAIChatCompletionClient() + common: dict[str, Any] = {"object": "chat.completion.chunk", "created": 1234567890, "model": "test-model"} + chunks = [ + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "A1"}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": "B1"}, "finish_reason": None}, + ], + }), + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + {"index": 0, "delta": {"content": "A2"}, "finish_reason": None}, + {"index": 1, "delta": {"content": "B2"}, "finish_reason": "stop"}, + ], + }), + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }), + ] + + async def create(**kwargs: Any) -> Any: + return _FakeAsyncStream(chunks) + + with patch.object(client.client.chat.completions, "create", side_effect=create): + response_stream = client._inner_get_response( + messages=[Message(role="user", contents=["test"])], stream=True, options={} + ) + assert isinstance(response_stream, ResponseStream) + streamed_texts: list[str | None] = [] + async for update in response_stream: + streamed_texts.extend(content.text for content in update.contents if content.type == "text") + final_response = await response_stream.get_final_response() + + # The live stream keeps provider arrival order (interleaved) with per-choice ids. + assert streamed_texts == ["A1", "B1", "A2", "B2"] + # The final response keeps one message per choice, like the non-streaming path. + assert [message.text for message in final_response.messages] == ["A1A2", "B1B2"] + assert final_response.finish_reason == "stop" + assert final_response.usage_details is not None + assert final_response.usage_details["total_token_count"] == 15 + + +def test_streaming_single_choice_message_id_matches_provider_chunk_id( + openai_unit_test_env: dict[str, str], +) -> None: + """n=1 keeps today's behavior: one message whose id is the provider chunk id.""" + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + + client = OpenAIChatCompletionClient() + chunk = ChatCompletionChunk.model_validate({ + "id": "chatcmpl-single", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "test-model", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + }) + + updates = client._parse_response_updates_from_openai(chunk) + + assert len(updates) == 1 + assert updates[0].message_id == "chatcmpl-single" + response = ChatResponse.from_updates(updates) + assert [message.text for message in response.messages] == ["hi"] + assert response.messages[0].message_id == "chatcmpl-single" + + async def test_streaming_tool_call_identity_is_request_local_and_scoped_by_choice_index( openai_unit_test_env: dict[str, str], ) -> None: From 143d6a7e1c8f91be0384badfd356337741f193b2 Mon Sep 17 00:00:00 2001 From: Wilson0511 <1795618152@qq.com> Date: Wed, 28 Oct 2026 13:17:44 +0800 Subject: [PATCH 2/3] Python: Rejoin interleaved message ids in from_updates, not in a client finalizer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up on #9780: grouping updates inside the chat client's _finalize_response_updates fixed the direct finalizer path only. FunctionInvocationLayer and Agent re-aggregate the streamed updates with ChatResponse.from_updates directly, so interleaved choices still shattered into one message per fragment on the public path, and regrouping updates reordered response-level usage snapshots so an older replace-mode total could win. Move the grouping into _process_update: when an update's message id names a message that already exists, the update returns to that message instead of starting a new one. Author and role boundaries still start a new message regardless of the id, and contiguous id sequences keep today's behavior. Updates keep flowing in arrival order, so usage snapshots apply in arrival order again. The client-side finalize regrouping is dropped; the per-choice updates stay. 🤖 Generated with ZCode Co-Authored-By: ZCode --- .../packages/core/agent_framework/_types.py | 24 ++++- python/packages/core/tests/core/test_types.py | 53 ++++++++++ .../_chat_completion_client.py | 27 ----- .../test_openai_chat_completion_client.py | 100 ++++++++++++++++++ 4 files changed, 175 insertions(+), 29 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 00edbaa937a..14219fea374 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2615,8 +2615,28 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse is_new_message = True if is_new_message: - message = Message("assistant", []) - response.messages.append(message) + author_or_role_boundary = bool(response.messages) and ( + ( + update.author_name + and response.messages[-1].author_name + and response.messages[-1].author_name != update.author_name + ) + or (update.role and response.messages[-1].role != update.role) + ) + message = None + if update.message_id and not author_or_role_boundary: + # The update belongs to the message its id names. When alternatives + # interleave in one stream (e.g. `n>1` chat completions choices), + # return to the message that already carries this id instead of + # starting a new one for every interleaved fragment. Author and role + # boundaries still start a new message regardless of the id. + for existing_message in reversed(response.messages): + if existing_message.message_id == update.message_id: + message = existing_message + break + if message is None: + message = Message("assistant", []) + response.messages.append(message) else: message = response.messages[-1] if update.merge_mode == "replace": diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index d047c92bedc..04989821143 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2442,6 +2442,59 @@ def test_from_updates_starts_new_message_for_authors_sharing_a_message_id() -> N ] +@pytest.mark.parametrize( + ("response_type", "update_type"), + [(ChatResponse, ChatResponseUpdate), (AgentResponse, AgentResponseUpdate)], + ids=["chat", "agent"], +) +def test_from_updates_rejoins_interleaved_updates_that_share_a_message_id( + response_type: Any, + update_type: Any, +) -> None: + """Test that interleaved updates carrying the same message id aggregate into one message each. + + Streams with interleaved alternatives (e.g. ``n>1`` chat completions choices) tag + each fragment with the id of the message it belongs to; aggregation must return to + the message with that id instead of shattering the stream into per-fragment + messages. + """ + response = response_type.from_updates([ + update_type(contents=[Content.from_text("A1")], role="assistant", message_id="m0"), + update_type(contents=[Content.from_text("B1")], role="assistant", message_id="m1"), + update_type(contents=[Content.from_text("A2")], role="assistant", message_id="m0"), + update_type(contents=[Content.from_text("B2")], role="assistant", message_id="m1"), + ]) + + assert [(message.message_id, message.text) for message in response.messages] == [ + ("m0", "A1A2"), + ("m1", "B1B2"), + ] + + +def test_from_updates_keeps_usage_in_arrival_order_across_interleaved_message_ids() -> None: + """Test that response-level usage follows update arrival order, not message grouping.""" + response = ChatResponse.from_updates([ + ChatResponseUpdate(contents=[Content.from_text("A1")], role="assistant", message_id="m0"), + ChatResponseUpdate(contents=[Content.from_text("B1")], role="assistant", message_id="m1"), + ChatResponseUpdate( + contents=[ + Content.from_text("B2"), + Content.from_usage(UsageDetails(total_token_count=10), merge_mode="replace"), + ], + role="assistant", + message_id="m1", + ), + ChatResponseUpdate( + contents=[Content.from_usage(UsageDetails(total_token_count=15), merge_mode="replace")], + role="assistant", + message_id="m0", + ), + ]) + + assert response.usage_details is not None + assert response.usage_details["total_token_count"] == 15 + + def test_agent_run_response_from_updates_uses_last_non_none_agent_id() -> None: response = AgentResponse.from_updates([ AgentResponseUpdate(agent_id="first-agent"), diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index c19a3e1d955..5b7dd5bef61 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -1067,33 +1067,6 @@ def _parse_response_update_from_openai( message_id=chunk.id, ) - @override - def _finalize_response_updates( - self, - updates: Sequence[ChatResponseUpdate], - *, - response_format: Any | None = None, - ) -> ChatResponse[Any]: - """Finalize response updates into a single ChatResponse. - - With ``n>1``, the provider interleaves the choices' deltas in one stream, so - the per-choice updates arrive as e.g. ``A1, B1, A2, B2`` even though ``A1`` - and ``A2`` belong to the same assistant message. ``ChatResponse.from_updates`` - starts a new message whenever the message id changes, so the updates are - first grouped by message id (stable, first-seen order) to aggregate each - choice into its own message, matching the non-streaming path. Streams where - every update shares one id (``n=1``) or any update has no id are left as-is. - """ - message_ids = [update.message_id for update in updates] - if len(message_ids) > 1 and all(message_id is not None for message_id in message_ids): - first_seen_order = dict.fromkeys(message_ids) - if len(first_seen_order) > 1: - grouped: dict[str, list[ChatResponseUpdate]] = {message_id: [] for message_id in first_seen_order} - for update in updates: - grouped[cast("str", update.message_id)].append(update) - updates = [update for message_id in first_seen_order for update in grouped[message_id]] - return super()._finalize_response_updates(updates, response_format=response_format) - def _parse_usage_from_openai(self, usage: CompletionUsage) -> UsageDetails: details = UsageDetails( input_token_count=usage.prompt_tokens, diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index 2270169cbe8..3fcea939a22 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -2964,6 +2964,106 @@ async def create(**kwargs: Any) -> Any: assert final_response.usage_details["total_token_count"] == 15 +async def test_public_client_streaming_multiple_choices_preserve_message_boundaries( + openai_unit_test_env: dict[str, str], +) -> None: + """The public client (function-invocation middleware included) must keep per-choice messages. + + ``FunctionInvocationLayer`` re-aggregates the streamed updates with + ``ChatResponse.from_updates`` instead of the chat client's finalizer, so the + per-choice grouping must survive that outer aggregation too. + """ + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + + client = OpenAIChatCompletionClient() + common: dict[str, Any] = {"object": "chat.completion.chunk", "created": 1234567890, "model": "test-model"} + chunks = [ + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "A1"}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": "B1"}, "finish_reason": None}, + ], + }), + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + {"index": 0, "delta": {"content": "A2"}, "finish_reason": None}, + {"index": 1, "delta": {"content": "B2"}, "finish_reason": "stop"}, + ], + }), + ] + + async def create(**kwargs: Any) -> Any: + return _FakeAsyncStream(chunks) + + with patch.object(client.client.chat.completions, "create", side_effect=create): + response = await client.get_response( + messages=[Message(role="user", contents=["test"])], + stream=True, + ).get_final_response() + + assert [message.text for message in response.messages] == ["A1A2", "B1B2"] + + +async def test_streaming_multiple_choices_keep_usage_in_arrival_order( + openai_unit_test_env: dict[str, str], +) -> None: + """Response-level usage must follow chunk arrival order, not per-choice message grouping. + + A provider may attach usage to a choice-1-only chunk before the final + usage-only chunk; the last usage snapshot in arrival order wins. + """ + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + + client = RawOpenAIChatCompletionClient() + common: dict[str, Any] = {"object": "chat.completion.chunk", "created": 1234567890, "model": "test-model"} + chunks = [ + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": "A1"}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": "B1"}, "finish_reason": None}, + ], + }), + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [ + { + "index": 1, + "delta": {"content": "B2"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}, + }), + ChatCompletionChunk.model_validate({ + **common, + "id": "multi-choice", + "choices": [], + "usage": {"prompt_tokens": 12, "completion_tokens": 3, "total_tokens": 15}, + }), + ] + + async def create(**kwargs: Any) -> Any: + return _FakeAsyncStream(chunks) + + with patch.object(client.client.chat.completions, "create", side_effect=create): + response_stream = client._inner_get_response( + messages=[Message(role="user", contents=["test"])], stream=True, options={} + ) + assert isinstance(response_stream, ResponseStream) + final_response = await response_stream.get_final_response() + + assert [message.text for message in final_response.messages] == ["A1", "B1B2"] + assert final_response.usage_details is not None + assert final_response.usage_details["total_token_count"] == 15 + + def test_streaming_single_choice_message_id_matches_provider_chunk_id( openai_unit_test_env: dict[str, str], ) -> None: From 08ac9fa0b75d8c578647ff8f1111d0c2629993e0 Mon Sep 17 00:00:00 2001 From: Wilson0511 <1795618152@qq.com> Date: Sun, 1 Nov 2026 20:43:36 +0800 Subject: [PATCH 3/3] Python: Narrow merged finish_reason literal back for pyright Reading finish_reason back from a ChatResponseUpdate widens the literal union to str under pyright (literal widening on inferred instance attributes), which fails the strict type check on the OpenAI floor and package-check jobs. Cast the merged value back to the constructor parameter type, mirroring the existing cast in _types.from_updates. Co-Authored-By: ZCode Signed-off-by: Wilson Jiang --- .../agent_framework_openai/_chat_completion_client.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index 5b7dd5bef61..8a05eadc3ff 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -38,6 +38,7 @@ ChatResponseUpdate, Content, FinishReason, + FinishReasonLiteral, Message, ResponseStream, UsageDetails, @@ -1061,7 +1062,14 @@ def _parse_response_update_from_openai( role="assistant", model=chunk.model, additional_properties=merged_metadata, - finish_reason=next((update.finish_reason for update in reversed(updates) if update.finish_reason), None), + # Reading ``finish_reason`` back from an update widens the literal to ``str`` + # (pyright widens literal types on inferred instance attributes), so narrow + # it back to the constructor parameter type; the values only ever come from + # that same parameter in ``_choice_update``. + finish_reason=cast( + "FinishReasonLiteral | FinishReason | None", + next((update.finish_reason for update in reversed(updates) if update.finish_reason), None), + ), raw_representation=chunk, response_id=chunk.id, message_id=chunk.id,