diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 00edbaa937..14219fea37 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 d047c92bed..0498982114 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 27da4fea33..8a05eadc3f 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, @@ -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, 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 7d5fc75b07..3fcea939a2 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,194 @@ 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 + + +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: + """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: