Python: Preserve per-choice message boundaries when streaming chat completions (n>1)

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":

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"),

python/packages/openai/agent_framework_openai/_chat_completion_client.py

@@ -38,6 +38,7 @@
     ChatResponseUpdate,
     Content,
     FinishReason,
+    FinishReasonLiteral,
     Message,
     ResponseStream,
     UsageDetails,
@@ -669,25 +670,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 +929,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,14 +1007,69 @@ 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,
+            # 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,