From b5bdffaeebe21e184f996ccb907398d34c2b51b9 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 27 Oct 2026 16:44:41 +0100 Subject: [PATCH 1/2] Python: Preserve Responses message identities --- .../foundry_hosting/tests/test_responses.py | 35 +++++++++++++++- .../agent_framework_openai/_chat_client.py | 9 ++++ .../tests/openai/test_openai_chat_client.py | 41 ++++++++++++++++--- 3 files changed, 78 insertions(+), 7 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 84e6ed6e69..d9adc22703 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -83,6 +83,7 @@ from azure.ai.agentserver.responses.streaming._checkpoint import ResponseCheckpointEvent from mcp import MCPError from openai import AsyncOpenAI, DefaultAsyncHttpxClient +from openai.types.responses import ResponseTextDeltaEvent, ResponseTextDoneEvent from openai.types.responses.response_input_item_param import ResponseInputItemParam from openai.types.responses.response_usage import ResponseUsage as OpenAIResponseUsage from pydantic import TypeAdapter @@ -7835,9 +7836,34 @@ async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage( assert loaded.function_call is not None assert loaded.function_call.name == "delete_file" - async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None: + @pytest.mark.parametrize("with_preamble", [False, True]) + async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self, with_preamble: bool) -> None: request_content = _make_function_approval_request_content(request_id="apr_streaming") - agent = _make_agent(stream_updates=[AgentResponseUpdate(contents=[request_content], role="assistant")]) + updates: list[AgentResponseUpdate] = [] + preambles = ("Checking the refund policy.", "The refund requires approval.") + if with_preamble: + client = OpenAIChatClient(model="test-model", api_key="test-key") + for output_index, text in enumerate(preambles): + event_args: dict[str, Any] = { + "content_index": 0, + "item_id": f"message_{output_index}", + "output_index": output_index, + "sequence_number": output_index * 2, + } + for event in ( + ResponseTextDeltaEvent(type="response.output_text.delta", delta=text, logprobs=[], **event_args), + ResponseTextDoneEvent(type="response.output_text.done", text=text, logprobs=[], **event_args), + ): + update = client._parse_chunk_from_openai(event, {}, {}) + updates.append( + AgentResponseUpdate( + contents=update.contents, + message_id=update.message_id, + role="assistant", + ) + ) + updates.append(AgentResponseUpdate(contents=[request_content], role="assistant")) + agent = _make_agent(stream_updates=updates) server = _make_server(agent) resp = await _post(server, stream=True) @@ -7848,6 +7874,11 @@ async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self assert types[0] == "response.created" assert types[-1] == "response.completed" + if with_preamble: + assert [event["data"]["text"] for event in events if event["event"] == "response.output_text.done"] == list( + preambles + ) + approval_request_id: str | None = None for e in events: if e["event"] != "response.output_item.added": diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 9647823b1a..ab0aa1551b 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -4577,8 +4577,17 @@ def parse_annotations( if not isinstance(event.type, str) or not event.type.startswith(_AZURE_AI_SEARCH_OUTPUT_EVENT_PREFIX): logger.debug("Unparsed event of type: %s: %s", event.type, event) + message_id: str | None = None + if any(content.type == "text" for content in contents): + item_id = getattr(event, "item_id", None) + if event.type == "response.output_item.done": + item_id = event.item.id + if isinstance(item_id, str): + message_id = item_id + update = ChatResponseUpdate( contents=contents, + message_id=message_id, conversation_id=conversation_id, response_id=response_id, role="assistant", 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 33162f9840..0a0e8ba8f0 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -60,6 +60,7 @@ ResponseFunctionShellToolCallOutput, ResponseOutputItemDoneEvent, ResponseOutputTextAnnotationAddedEvent, + ResponseStreamEvent, ) from openai.types.responses.response_reasoning_item import Summary from openai.types.responses.response_reasoning_summary_text_delta_event import ( @@ -1748,6 +1749,38 @@ def test_streaming_text_done_events_replace_deltas( assert response.messages[0].contents[0].text == expected_text +@pytest.mark.parametrize("refusal", [False, True]) +def test_streaming_text_snapshots_preserve_separate_output_messages(refusal: bool) -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + updates: list[ChatResponseUpdate] = [] + for output_index, text in enumerate(("Checking the refund policy.", "The refund requires approval.")): + event_args: dict[str, Any] = { + "content_index": 0, + "item_id": f"message_{output_index}", + "output_index": output_index, + "sequence_number": output_index * 2, + } + events: tuple[ResponseStreamEvent, ...] + if refusal: + events = ( + ResponseRefusalDeltaEvent(type="response.refusal.delta", delta=text, **event_args), + ResponseRefusalDoneEvent(type="response.refusal.done", refusal=text, **event_args), + ) + else: + events = ( + ResponseTextDeltaEvent(type="response.output_text.delta", delta=text, logprobs=[], **event_args), + ResponseTextDoneEvent(type="response.output_text.done", text=text, logprobs=[], **event_args), + ) + updates.extend(client._parse_chunk_from_openai(event, {}, {}) for event in events) + + response = ChatResponse.from_updates(updates) + + assert [message.text for message in response.messages] == [ + "Checking the refund policy.", + "The refund requires approval.", + ] + + def test_prepare_marked_refusal_text_uses_native_assistant_shape_and_input_text_fallback() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") refusal = Content.from_text( @@ -7373,11 +7406,9 @@ async def get_response() -> ChatResponse: responses = await asyncio.gather(get_response(), get_response()) for response in responses: - assert len(response.messages) == 1 - assert len(response.messages[0].contents) == 1 - annotations = response.messages[0].contents[0].annotations - assert annotations is not None - assert len(annotations) == 4 + assert len(response.messages) == 2 + assert [len(message.contents) for message in response.messages] == [1, 1] + assert [len(message.contents[0].annotations or []) for message in response.messages] == [3, 1] def test_streaming_completed_text_without_annotations_emits_no_content() -> None: From 26366eb0fe8f96a459973cfc7ff1c542cc1eba05 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 27 Oct 2026 17:48:38 +0100 Subject: [PATCH 2/2] Python: Align Responses message finalization --- .../agent_framework_openai/_chat_client.py | 55 ++++-- .../tests/openai/test_openai_chat_client.py | 181 +++++++++++++++++- 2 files changed, 223 insertions(+), 13 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index ab0aa1551b..6d6b08a9a5 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1073,8 +1073,28 @@ def _finalize_response_updates( self._enrich_mcp_search_citations([content for update in updates for content in update.contents]) response = super()._finalize_response_updates(updates, response_format=response_format) self._mark_completed_computer_calls([content for message in response.messages for content in message.contents]) - logprobs: list[Any] = [] + message_logprobs: list[list[Any]] = [] + current_message_id: str | None = None + current_author_name: str | None = None + current_role: str | None = None for update in updates: + is_new_message = ( + not message_logprobs + or (update.message_id and current_message_id and current_message_id != update.message_id) + or (update.author_name and current_author_name and current_author_name != update.author_name) + or (update.role and current_role != update.role) + ) + if is_new_message: + message_logprobs.append([]) + current_message_id = None + current_author_name = None + current_role = "assistant" + if update.message_id: + current_message_id = update.message_id + if update.author_name is not None: + current_author_name = update.author_name + if update.role is not None: + current_role = update.role for content in update.contents: if content.type != "text": continue @@ -1085,15 +1105,15 @@ def _finalize_response_updates( default=[], ) if isinstance(content_logprobs, list): - logprobs.extend(cast("list[Any]", content_logprobs)) - if logprobs: + message_logprobs[-1].extend(cast("list[Any]", content_logprobs)) + for message, logprobs in zip(response.messages, message_logprobs): + if not logprobs or message.role != "assistant": + continue assistant_text = next( ( content - for message in response.messages - if message.role == "assistant" for content in message.contents - if content.type == "text" + if content.type == "text" and not _is_refusal_text_content(content) ), None, ) @@ -3470,13 +3490,23 @@ def _parse_response_from_openai( structured_response: BaseModel | None = response.output_parsed if isinstance(response, ParsedResponse) else None # type: ignore[reportUnknownMemberType] metadata: dict[str, Any] = response.metadata or {} - contents: list[Content] = [] + messages: list[Message] = [] local_shell_tool_name = self._get_local_shell_tool_name(options.get("tools")) + + def contents_for_message(message_id: str | None = None) -> list[Content]: + if not messages or (message_id and messages[-1].message_id and messages[-1].message_id != message_id): + messages.append(Message(role="assistant", contents=[], message_id=message_id)) + elif message_id: + messages[-1].message_id = message_id + return messages[-1].contents + try: response_outputs = response.output # type: ignore[reportUnknownMemberType] except AttributeError: response_outputs = [] for item in response_outputs: # type: ignore[reportUnknownVariableType] + message_id = cast(Any, item).id if item.type == "message" else None + contents = contents_for_message(message_id if isinstance(message_id, str) else None) match item.type: # types: # ParsedResponseOutputMessage[Unknown] | @@ -3758,14 +3788,16 @@ def _parse_response_from_openai( contents.extend(self._shell_item_to_contents(item, local_shell_tool_name)) case _: logger.debug("Unparsed output of type: %s: %s", item.type, item) - self._mark_completed_computer_calls(contents) - response_message = Message(role="assistant", contents=contents) + if not messages: + messages.append(Message(role="assistant", contents=[])) + all_contents = [content for message in messages for content in message.contents] + self._mark_completed_computer_calls(all_contents) args: dict[str, Any] = { "response_id": response.id, "created_at": datetime.fromtimestamp(response.created_at, tz=timezone.utc).strftime( "%Y-%m-%dT%H:%M:%S.%fZ" ), - "messages": response_message, + "messages": messages, "model": response.model, "additional_properties": metadata, "raw_representation": response, @@ -3785,8 +3817,7 @@ def _parse_response_from_openai( if response.status and response.status in ("in_progress", "queued"): args["continuation_token"] = OpenAIContinuationToken(response_id=response.id) chat_response = ChatResponse(**args) - if chat_response.messages: - self._enrich_mcp_search_citations(chat_response.messages[0].contents) + self._enrich_mcp_search_citations(all_contents) return chat_response def _parse_annotation_from_openai(self, annotation: Any, annotation_index: int) -> Annotation | None: 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 0a0e8ba8f0..f5c7c58efe 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -58,11 +58,13 @@ ResponseErrorEvent, ResponseFunctionShellToolCall, ResponseFunctionShellToolCallOutput, + ResponseFunctionToolCall, ResponseOutputItemDoneEvent, + ResponseOutputMessage, ResponseOutputTextAnnotationAddedEvent, ResponseStreamEvent, ) -from openai.types.responses.response_reasoning_item import Summary +from openai.types.responses.response_reasoning_item import ResponseReasoningItem, Summary from openai.types.responses.response_reasoning_summary_text_delta_event import ( ResponseReasoningSummaryTextDeltaEvent, ) @@ -1781,6 +1783,136 @@ def test_streaming_text_snapshots_preserve_separate_output_messages(refusal: boo ] +def test_streaming_and_non_streaming_preserve_output_message_boundaries() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + first_message = ResponseOutputMessage.model_validate({ + "type": "message", + "id": "message_0", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Checking the refund policy.", + "annotations": [], + "logprobs": [], + } + ], + }) + second_message = ResponseOutputMessage.model_validate({ + "type": "message", + "id": "message_1", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "The refund requires approval.", + "annotations": [], + "logprobs": [], + } + ], + }) + reasoning_item = ResponseReasoningItem.model_validate({ + "type": "reasoning", + "id": "reasoning_0", + "content": [ + { + "type": "reasoning_text", + "text": "The amount exceeds the automatic refund threshold.", + } + ], + "summary": [], + "status": "completed", + }) + function_call = ResponseFunctionToolCall.model_validate({ + "type": "function_call", + "id": "function_0", + "call_id": "call_0", + "name": "refund", + "arguments": '{"amount": 125}', + "status": "completed", + }) + response = MagicMock( + output_parsed=None, + output=[first_message, reasoning_item, function_call, second_message], + metadata={}, + usage=None, + id="response_0", + created_at=1_000_000_000, + model="test-model", + status="completed", + incomplete_details=None, + conversation=None, + ) + + non_streaming = client._parse_response_from_openai(response, options={"store": False}) + streaming = client._finalize_response_updates([ + client._parse_chunk_from_openai( + ResponseTextDoneEvent( + type="response.output_text.done", + content_index=0, + item_id="message_0", + output_index=0, + sequence_number=0, + text="Checking the refund policy.", + logprobs=[], + ), + {}, + {}, + ), + ChatResponseUpdate( + contents=[ + Content.from_text_reasoning( + id="reasoning_0", + text="The amount exceeds the automatic refund threshold.", + ) + ], + role="assistant", + ), + ChatResponseUpdate( + contents=[ + Content.from_function_call( + call_id="call_0", + name="refund", + arguments='{"amount": 125}', + ) + ], + role="assistant", + ), + client._parse_chunk_from_openai( + ResponseTextDoneEvent( + type="response.output_text.done", + content_index=0, + item_id="message_1", + output_index=1, + sequence_number=1, + text="The refund requires approval.", + logprobs=[], + ), + {}, + {}, + ), + ]) + + expected = [ + ("message_0", ["text", "text_reasoning", "function_call"], "Checking the refund policy."), + ("message_1", ["text"], "The refund requires approval."), + ] + assert [ + (message.message_id, [content.type for content in message.contents], message.text) + for message in non_streaming.messages + ] == expected + assert [ + (message.message_id, [content.type for content in message.contents], message.text) + for message in streaming.messages + ] == expected + assert [ + (message.message_id, [content.type for content in message.contents], message.text) + for message in ChatResponse.from_dict(non_streaming.to_dict()).messages + ] == expected + + def test_prepare_marked_refusal_text_uses_native_assistant_shape_and_input_text_fallback() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") refusal = Content.from_text( @@ -3728,6 +3860,53 @@ def test_streaming_parse_and_replay_preserves_all_real_assistant_logprobs() -> N assert replayed["logprobs"] == [first_logprob, second_logprob] +def test_streaming_parse_and_replay_scopes_logprobs_to_output_messages() -> None: + """Streamed token logprobs remain attached to their provider output message.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + message_logprobs = [ + [ + {"token": "first", "bytes": [102], "logprob": -0.1, "top_logprobs": []}, + {"token": " message", "bytes": [109], "logprob": -0.2, "top_logprobs": []}, + ], + [ + {"token": "second", "bytes": [115], "logprob": -0.3, "top_logprobs": []}, + {"token": " message", "bytes": [109], "logprob": -0.4, "top_logprobs": []}, + ], + ] + updates = [ + client._parse_chunk_from_openai( + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id=f"message_{message_index}", + output_index=message_index, + sequence_number=message_index * 2 + chunk_index, + logprobs=[logprob], # type: ignore[list-item] + delta=cast(str, logprob["token"]), + ), + options={}, + function_call_ids={}, + ) + for message_index, logprobs in enumerate(message_logprobs) + for chunk_index, logprob in enumerate(logprobs) + ] + + response = client._finalize_response_updates(updates) + replayed = client._prepare_messages_for_openai( + response.messages, + request_uses_service_side_storage=False, + ) + + assert [(message.message_id, message.text) for message in response.messages] == [ + ("message_0", "first message"), + ("message_1", "second message"), + ] + assert [ + message.contents[0].additional_properties["openai"]["logprobs"] for message in response.messages + ] == message_logprobs + assert [message["content"][0]["logprobs"] for message in replayed] == message_logprobs + + def test_prepare_messages_for_openai_assistant_history_uses_output_text_with_annotations() -> None: """Assistant history should be output_text and include required annotations.""" client = OpenAIChatClient(model="test-model", api_key="test-key")