Python: Record actual served model from Azure OpenAI (#5910)

* Record actual served model as response model for Azure OpenAI * Formatting * Fix tests * Fix pipeline error * Comments * Address review: surface served model via ChatResponse.model Apply blocking review feedback from PR #5910: - Use ChatResponse.model / ChatResponseUpdate.model as the source of truth for the Azure x-ms-served-model header value, instead of stashing it in additional_properties and overriding it again in observability. Observability already reads response.model; the chat client now overwrites it post-parse when the served-model header is present. Empirically the Azure Responses API returns the deployment alias in body.model and the actual snapshot (e.g. gpt-5-nano-2025-08-07) in this header. - Move the AZURE_OPENAI_SERVED_MODEL_HEADER constant out of observability.py and into RawOpenAIChatClient (as the SERVED_MODEL_HEADER ClassVar). The header is Azure-OpenAI-Responses-API-specific so observability does not need to know about it. - Revert the streaming text_format path to client.responses.stream(...) and drop the _pydantic_model_to_text_format_param helper. That helper imported from openai.lib._parsing._responses (a private SDK path) and the swap to responses.create(stream=True) dropped client-side output_parsed for structured-output streaming. The streaming-with-text_format path is the only one that does not surface the served-model header - documented inline. - Wrap the raw streaming responses in async with so the underlying socket closes deterministically (continuation_token retrieve + create paths). - Fix the empty-string / whitespace-only header at the source by stripping in _extract_served_model and returning None when nothing remains. - Revert unrelated formatting-only churn in _skills.py and test_mcp.py. - Update unit tests to assert against chat_response.model / update.model and add an aggregated streaming assertion plus a pin that the streaming-with-text_format path does not get the header. Verified end-to-end against Azure OpenAI Responses API: deployment alias gpt-5-nano now reports gpt-5-nano-2025-08-07 as ChatResponse.model in both the non-streaming and streaming paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: preserve streaming structured output finalization Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/f62076ef-558d-49e8-8fe2-f38d527c9639 Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> * refactor: name streaming response finalizer Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/f62076ef-558d-49e8-8fe2-f38d527c9639 Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> * fix: capture streaming response format after prepare Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/f62076ef-558d-49e8-8fe2-f38d527c9639 Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> * refactor: clarify streaming response format capture Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/f62076ef-558d-49e8-8fe2-f38d527c9639 Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> * test: use public API for streaming structured output Agent-Logs-Url: https://github.com/microsoft/agent-framework/sessions/f62076ef-558d-49e8-8fe2-f38d527c9639 Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com> * Inline the served-model header override at its two call sites The `_apply_served_model_header` helper was a 1-line wrapper around `_extract_served_model`. Inlining the `if served_model is not None: ...` matches the pattern already used in the streaming paths and folds the explanatory docstring onto `_extract_served_model` (which is now the single place that knows about the header). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-05-18 23:38:53 -07:00
parent 3bbc81554b
commit 1b6f7d80fd
6 changed files with 471 additions and 54 deletions
@@ -72,9 +72,10 @@ class OutputStruct(BaseModel):


 class _FakeAsyncEventStream:
-    def __init__(self, events: list[object]) -> None:
+    def __init__(self, events: list[object], headers: dict[str, str] | None = None) -> None:
        self._events = events
        self._iterator = iter(())
+        self._headers = headers or {}

    def __aiter__(self) -> "_FakeAsyncEventStream":
        self._iterator = iter(self._events)
@@ -86,6 +87,45 @@ class _FakeAsyncEventStream:
        except StopIteration as exc:
            raise StopAsyncIteration from exc

+    # The chat client now consumes the streaming response via ``with_raw_response``,
+    # which returns a wrapper exposing ``.parse()`` (the underlying iterable) and
+    # ``.headers``. The chat client then ``async with``-s the parsed stream so the
+    # underlying socket is closed deterministically. Mimic both interfaces here so
+    # test mocks remain a single object.
+    def parse(self) -> "_FakeAsyncEventStream":
+        return self
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return self._headers
+
+    async def __aenter__(self) -> "_FakeAsyncEventStream":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: object | None,
+    ) -> None:
+        return None
+
+
+def _as_raw(mock_response: MagicMock, *, headers: dict[str, str] | None = None) -> MagicMock:
+    """Make ``mock_response`` look like an OpenAI ``with_raw_response`` wrapper.
+
+    The chat client now calls ``responses.with_raw_response.{create,parse,retrieve}``
+    and then ``.parse()`` on the returned wrapper to get the actual response payload,
+    plus ``.headers`` to surface the ``x-ms-served-model`` Azure header. Tests still
+    patch the underlying ``responses.{create,parse,retrieve}`` methods (the SDK's
+    raw-response wrapper internally delegates to these), so the patched return value
+    is what our code unwraps. Setting ``mock_response.parse`` to return the mock
+    itself lets the existing assertions on ``mock_response.id`` etc. continue to work.
+    """
+    mock_response.parse = MagicMock(return_value=mock_response)
+    mock_response.headers = headers or {}
+    return mock_response
+

 class _FakeAsyncEventStreamContext(_FakeAsyncEventStream):
    async def __aenter__(self) -> "_FakeAsyncEventStreamContext":
@@ -477,7 +517,7 @@ async def test_response_format_parse_path() -> None:
    mock_parsed_response.finish_reason = None
    mock_parsed_response.conversation = None  # No conversation object

-    with patch.object(client.client.responses, "parse", return_value=mock_parsed_response):
+    with patch.object(client.client.responses, "parse", return_value=_as_raw(mock_parsed_response)):
        response = await client.get_response(
            messages=[Message(role="user", contents=["Test message"])],
            options={"response_format": OutputStruct, "store": True},
@@ -504,7 +544,7 @@ async def test_response_format_parse_path_with_conversation_id() -> None:
    mock_parsed_response.conversation = MagicMock()
    mock_parsed_response.conversation.id = "conversation_456"

-    with patch.object(client.client.responses, "parse", return_value=mock_parsed_response):
+    with patch.object(client.client.responses, "parse", return_value=_as_raw(mock_parsed_response)):
        response = await client.get_response(
            messages=[Message(role="user", contents=["Test message"])],
            options={"response_format": OutputStruct, "store": True},
@@ -542,7 +582,7 @@ async def test_response_format_dict_parse_path() -> None:
    mock_message_item.content = [mock_message_content]
    mock_response.output = [mock_message_item]

-    with patch.object(client.client.responses, "create", return_value=mock_response):
+    with patch.object(client.client.responses, "create", return_value=_as_raw(mock_response)):
        response = await client.get_response(
            messages=[Message(role="user", contents=["Test message"])],
            options={"response_format": response_format},
@@ -554,6 +594,297 @@ async def test_response_format_dict_parse_path() -> None:
    assert response.value["answer"] == "Parsed"


+_SERVED_MODEL_HEADER = "x-ms-served-model"
+
+
+async def test_served_model_header_overrides_response_model() -> None:
+    """The ``x-ms-served-model`` Azure response header should overwrite ChatResponse.model."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_response = MagicMock()
+    mock_response.id = "response_123"
+    mock_response.model = "test-model"  # deployment alias returned in the body
+    mock_response.created_at = 1000000000
+    mock_response.metadata = {}
+    mock_response.output_parsed = None
+    mock_response.output = []
+    mock_response.usage = None
+    mock_response.finish_reason = None
+    mock_response.conversation = None
+    mock_response.status = "completed"
+
+    raw = _as_raw(mock_response, headers={_SERVED_MODEL_HEADER: "gpt-4o-2024-08-06"})
+
+    with patch.object(client.client.responses, "create", return_value=raw):
+        response = await client.get_response(
+            messages=[Message(role="user", contents=["Test message"])],
+        )
+
+    assert response.model == "gpt-4o-2024-08-06"
+
+
+async def test_served_model_header_absent_keeps_response_model() -> None:
+    """When the served-model header is missing ChatResponse.model should come from the response body."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_response = MagicMock()
+    mock_response.id = "response_123"
+    mock_response.model = "test-model"
+    mock_response.created_at = 1000000000
+    mock_response.metadata = {}
+    mock_response.output_parsed = None
+    mock_response.output = []
+    mock_response.usage = None
+    mock_response.finish_reason = None
+    mock_response.conversation = None
+    mock_response.status = "completed"
+
+    # _as_raw sets headers to {} by default — i.e. no x-ms-served-model.
+    with patch.object(client.client.responses, "create", return_value=_as_raw(mock_response)):
+        response = await client.get_response(
+            messages=[Message(role="user", contents=["Test message"])],
+        )
+
+    assert response.model == "test-model"
+
+
+async def test_served_model_header_empty_string_does_not_override() -> None:
+    """Empty/whitespace header values should not overwrite the response body's model name."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_response = MagicMock()
+    mock_response.id = "response_123"
+    mock_response.model = "test-model"
+    mock_response.created_at = 1000000000
+    mock_response.metadata = {}
+    mock_response.output_parsed = None
+    mock_response.output = []
+    mock_response.usage = None
+    mock_response.finish_reason = None
+    mock_response.conversation = None
+    mock_response.status = "completed"
+
+    raw = _as_raw(mock_response, headers={_SERVED_MODEL_HEADER: "   "})
+
+    with patch.object(client.client.responses, "create", return_value=raw):
+        response = await client.get_response(
+            messages=[Message(role="user", contents=["Test message"])],
+        )
+
+    assert response.model == "test-model"
+
+
+async def test_served_model_header_captured_on_parse_path() -> None:
+    """The served-model header should also be captured on the structured-output (parse) path."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    mock_parsed_response = MagicMock()
+    mock_parsed_response.id = "parsed_response_123"
+    mock_parsed_response.text = "Parsed response"
+    mock_parsed_response.model = "test-model"
+    mock_parsed_response.created_at = 1000000000
+    mock_parsed_response.metadata = {}
+    mock_parsed_response.output_parsed = None
+    mock_parsed_response.usage = None
+    mock_parsed_response.finish_reason = None
+    mock_parsed_response.conversation = None
+
+    raw = _as_raw(mock_parsed_response, headers={_SERVED_MODEL_HEADER: "gpt-4o-2024-08-06"})
+
+    with patch.object(client.client.responses, "parse", return_value=raw):
+        response = await client.get_response(
+            messages=[Message(role="user", contents=["Test message"])],
+            options={"response_format": OutputStruct, "store": True},
+        )
+
+    assert response.model == "gpt-4o-2024-08-06"
+
+
+async def test_served_model_header_propagated_to_streaming_updates() -> None:
+    """In streaming mode the served-model header should overwrite update.model on every chunk."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    events = [
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=1,
+            logprobs=[],
+            delta="Hello",
+        ),
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=2,
+            logprobs=[],
+            delta=" world",
+        ),
+    ]
+
+    fake_stream = _FakeAsyncEventStream(events, headers={_SERVED_MODEL_HEADER: "gpt-4o-2024-08-06"})
+
+    with (
+        patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))),
+        patch.object(client.client.responses, "create", new=AsyncMock(return_value=fake_stream)),
+        patch.object(client, "_get_metadata_from_response", return_value={}),
+    ):
+        stream = client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True)
+        updates = [update async for update in stream]
+
+    assert updates, "Expected at least one streaming update"
+    for update in updates:
+        assert update.model == "gpt-4o-2024-08-06"
+
+
+async def test_served_model_header_aggregates_into_final_streaming_response() -> None:
+    """Aggregating updates via to_chat_response() should preserve the served-model value."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    events = [
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=1,
+            logprobs=[],
+            delta="Hello",
+        ),
+    ]
+
+    fake_stream = _FakeAsyncEventStream(events, headers={_SERVED_MODEL_HEADER: "gpt-4o-2024-08-06"})
+
+    with (
+        patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))),
+        patch.object(client.client.responses, "create", new=AsyncMock(return_value=fake_stream)),
+        patch.object(client, "_get_metadata_from_response", return_value={}),
+    ):
+        stream = client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True)
+        updates = [update async for update in stream]
+
+    final = ChatResponse.from_updates(updates)
+    assert final.model == "gpt-4o-2024-08-06"
+
+
+async def test_served_model_header_absent_in_streaming_updates() -> None:
+    """When the header is missing in streaming mode update.model should fall back to the deployment alias."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    events = [
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=1,
+            logprobs=[],
+            delta="Hello",
+        ),
+    ]
+
+    fake_stream = _FakeAsyncEventStream(events)  # default empty headers
+
+    with (
+        patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))),
+        patch.object(client.client.responses, "create", new=AsyncMock(return_value=fake_stream)),
+        patch.object(client, "_get_metadata_from_response", return_value={}),
+    ):
+        stream = client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True)
+        updates = [update async for update in stream]
+
+    assert updates, "Expected at least one streaming update"
+    for update in updates:
+        # Without the header, _parse_chunk_from_openai's default is the client's model name.
+        assert update.model == "test-model"
+
+
+async def test_served_model_header_not_captured_for_streaming_text_format() -> None:
+    """The streaming structured-output path uses ``responses.stream(...)`` and therefore cannot
+    surface the served-model header. Pin this behavior so any future change is intentional."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    events = [
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=1,
+            logprobs=[],
+            delta="Hello",
+        ),
+    ]
+
+    # `responses.stream(...)` returns an async context manager. The headers attribute
+    # is irrelevant because this code path never asks for it.
+    fake_stream_ctx = _FakeAsyncEventStreamContext(events)
+
+    with (
+        patch.object(
+            client,
+            "_prepare_request",
+            new=AsyncMock(return_value=(client.client, {"text_format": OutputStruct}, {})),
+        ),
+        patch.object(client.client.responses, "stream", return_value=fake_stream_ctx),
+        patch.object(client, "_get_metadata_from_response", return_value={}),
+    ):
+        stream = client._inner_get_response(messages=[Message(role="user", contents=["Hi"])], options={}, stream=True)
+        updates = [update async for update in stream]
+
+    assert updates, "Expected at least one streaming update"
+    for update in updates:
+        # No header override; model stays the deployment alias.
+        assert update.model == "test-model"
+
+
+async def test_streaming_text_format_preserves_final_structured_output() -> None:
+    """Streaming structured output should still parse into the final ChatResponse value."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+
+    events = [
+        ResponseTextDeltaEvent(
+            type="response.output_text.delta",
+            content_index=0,
+            item_id="text_item",
+            output_index=0,
+            sequence_number=1,
+            logprobs=[],
+            delta='{"location":"Seattle","weather":"Sunny"}',
+        ),
+    ]
+
+    fake_stream_ctx = _FakeAsyncEventStreamContext(events)
+
+    with (
+        patch.object(
+            client,
+            "_prepare_request",
+            new=AsyncMock(
+                return_value=(
+                    client.client,
+                    {"text_format": OutputStruct},
+                    {"response_format": OutputStruct},
+                )
+            ),
+        ),
+        patch.object(client.client.responses, "stream", return_value=fake_stream_ctx),
+        patch.object(client, "_get_metadata_from_response", return_value={}),
+    ):
+        stream = client.get_response(
+            messages=[Message(role="user", contents=["Hi"])],
+            options={"response_format": OutputStruct},
+            stream=True,
+        )
+        response = await stream.get_final_response()
+
+    assert response.model == "test-model"
+    assert response.value == OutputStruct(location="Seattle", weather="Sunny")
+
+
 async def test_bad_request_error_non_content_filter() -> None:
    """Test get_response BadRequestError without content_filter."""
    client = OpenAIChatClient(model="test-model", api_key="test-key")
@@ -953,7 +1284,9 @@ async def test_local_shell_tool_is_invoked_in_function_loop() -> None:
    mock_text_item.content = [mock_text_content]
    mock_response2.output = [mock_text_item]

-    with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create:
+    with patch.object(
+        client.client.responses, "create", side_effect=[_as_raw(mock_response1), _as_raw(mock_response2)]
+    ) as mock_create:
        await client.get_response(
            messages=[Message(role="user", contents=["What Python version is available?"])],
            options={"tools": [local_shell_tool]},
@@ -1026,7 +1359,9 @@ async def test_shell_call_is_invoked_as_local_shell_function_loop() -> None:
    mock_text_item.content = [mock_text_content]
    mock_response2.output = [mock_text_item]

-    with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create:
+    with patch.object(
+        client.client.responses, "create", side_effect=[_as_raw(mock_response1), _as_raw(mock_response2)]
+    ) as mock_create:
        await client.get_response(
            messages=[Message(role="user", contents=["What Python version is available?"])],
            options={"tools": [local_shell_tool]},
@@ -1097,7 +1432,9 @@ async def test_tool_loop_store_false_omits_reasoning_items_from_second_request()
    mock_text_item.content = [mock_text_content]
    mock_response2.output = [mock_text_item]

-    with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create:
+    with patch.object(
+        client.client.responses, "create", side_effect=[_as_raw(mock_response1), _as_raw(mock_response2)]
+    ) as mock_create:
        response = await client.get_response(
            messages=[Message(role="user", contents=["What's the weather in Amsterdam?"])],
            options={
@@ -2810,7 +3147,9 @@ async def test_end_to_end_mcp_approval_flow(span_exporter) -> None:
    mock_response2.output = [mock_text_item]

    # Patch the create call to return the two mocked responses in sequence
-    with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create:
+    with patch.object(
+        client.client.responses, "create", side_effect=[_as_raw(mock_response1), _as_raw(mock_response2)]
+    ) as mock_create:
        # First call: get the approval request
        response = await client.get_response(messages=[Message(role="user", contents=["Trigger approval"])])
        assert response.messages[0].contents[0].type == "function_approval_request"
@@ -4120,9 +4459,7 @@ async def test_prepare_options_with_conversation_id_strips_server_items_for_mixe
    types = [item.get("type") for item in options["input"]]
    assert "reasoning" not in types
    assert "function_call" not in types
-    output_call_ids = {
-        item["call_id"] for item in options["input"] if item.get("type") == "function_call_output"
-    }
+    output_call_ids = {item["call_id"] for item in options["input"] if item.get("type") == "function_call_output"}
    assert output_call_ids == {"call_history", "call_live"}
    assert options["previous_response_id"] == "resp_prev123"