Repository navigation
Python: Preserve per-choice message boundaries when streaming chat completions (n>1) - #9780
Conversation
…mpletions (n>1) 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 <noreply@zcode.dev>
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
🟡 Changes recommended
Public client and agent finalizers still fail to preserve per-choice message grouping, while regrouping updates can overwrite the latest usage total with an older snapshot.
2 open issues
What changed in this PR
This PR addresses #9779 by preserving separate completion choices in Python’s OpenAI Chat Completions streaming adapter, bringing raw streamed responses closer to non-streaming behavior.
Changes:
- Emit one update per choice with distinct message IDs while preserving live arrival order.
- Group updates by message ID during finalization and retain the legacy single-update parser wrapper.
- Add regressions for multi-choice text boundaries, trailing usage, and single-choice message IDs.
| File | Description |
|---|---|
python/packages/openai/tests/openai/test_openai_chat_completion_client.py |
Adds multi-choice aggregation and single-choice compatibility tests. |
python/packages/openai/agent_framework_openai/_chat_completion_client.py |
Adds per-choice parsing and message grouping during stream finalization. |
Review effort: Balanced
| 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]] |
There was a problem hiding this comment.
Fixed in 0295ecd: the client-side finalize regrouping is removed entirely. Updates now keep flowing into ChatResponse.from_updates in arrival order, so replace-mode usage snapshots apply strictly in arrival order regardless of which choice's chunk carried them. Added test_streaming_multiple_choices_keep_usage_in_arrival_order covering exactly this shape (choice-1 usage+content chunk with total 10, then a usage-only chunk with total 15 → final total is 15), plus the core-level test_from_updates_keeps_usage_in_arrival_order_across_interleaved_message_ids.
| 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) |
There was a problem hiding this comment.
Fixed in 0295ecd: the grouping moved into core _process_update — when an update's message_id names a message that already exists in the response, 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, so the existing author/role boundary tests are unchanged. Because FunctionInvocationLayer and Agent both aggregate through from_updates, per-choice grouping now survives those outer finalizers. Added test_public_client_streaming_multiple_choices_preserve_message_boundaries (public client, function-invocation path) and test_from_updates_rejoins_interleaved_updates_that_share_a_message_id, parametrized over both ChatResponse and AgentResponse.
|
@microsoft-github-policy-service agree |
…nt finalizer Review follow-up on microsoft#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 <noreply@zcode.dev>
GitHub Actions
Inactive
GitHub Actions
Inactive
Fixes #9779
Description
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, soChatResponse.from_updatesconcatenated the alternatives' texts in arrival order into a single message, while the non-streaming parser emits one assistantMessageper choice.The fix has two halves:
agent_framework_openai):_parse_response_updates_from_openaiparses one chunk into one update per choice. Choice 0 keeps the provider chunk id as itsmessage_id; the other choices get a{chunk.id}-{index}id so their message boundaries stay distinct. The live stream keeps provider arrival order. The old_parse_response_update_from_openaistays as a thin back-compat wrapper that merges the per-choice updates (its single-choice behavior is unchanged).agent_framework_core):_process_updatenow returns an update to the message that already carries itsmessage_idinstead of starting a new message for every interleaved fragment. Author and role boundaries still start a new message regardless of the id, and contiguous id sequences keep the previous behavior, soFunctionInvocationLayer,Agent, and directfrom_updatescallers all preserve per-choice grouping. Updates keep flowing in arrival order, so response-level usage snapshots apply in arrival order.RegressionCheck
['A1B1A2B2']; after,['A1A2', 'B1B2'], matching_parse_response_from_openai. Verified through the raw client, the public client (function-invocation path), andChatResponse/AgentResponsefrom_updates.n=1unchanged: single update per chunk,message_id == chunk.id, one message.tool_call_choice_index) unaffected.pytest tests/core/andpytest tests/openai/on this branch: failure sets byte-identical tomain(pre-existing environment failures only; zero new). New tests:test_streaming_multiple_choices_preserve_message_boundaries,test_public_client_streaming_multiple_choices_preserve_message_boundaries,test_streaming_multiple_choices_keep_usage_in_arrival_order,test_streaming_single_choice_message_id_matches_provider_chunk_id,test_from_updates_rejoins_interleaved_updates_that_share_a_message_id(chat + agent),test_from_updates_keeps_usage_in_arrival_order_across_interleaved_message_ids.🤖 Generated with ZCode
Co-Authored-By: ZCode noreply@zcode.dev