Skip to content

Commit e9d97ce

Browse files
authored
Python: fix(azure-ai): Fix response_format handling for structured outputs (#3114)
* fix(azure-ai): read response_format from chat_options instead of run_options * refactor: use explicit None checks for response_format * Fix mypy error * Mypy fix
1 parent f4ab586 commit e9d97ce

4 files changed

Lines changed: 234 additions & 24 deletions

File tree

python/packages/azure-ai/agent_framework_azure_ai/_client.py

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -300,13 +300,26 @@ def _create_text_format_config(
300300
raise ServiceInvalidRequestError("response_format must be a Pydantic model or mapping.")
301301

302302
async def _get_agent_reference_or_create(
303-
self, run_options: dict[str, Any], messages_instructions: str | None
303+
self,
304+
run_options: dict[str, Any],
305+
messages_instructions: str | None,
306+
chat_options: ChatOptions | None = None,
304307
) -> dict[str, str]:
305308
"""Determine which agent to use and create if needed.
306309
310+
Args:
311+
run_options: The prepared options for the API call.
312+
messages_instructions: Instructions extracted from messages.
313+
chat_options: The chat options containing response_format and other settings.
314+
307315
Returns:
308316
dict[str, str]: The agent reference to use.
309317
"""
318+
# chat_options is needed separately because the base class excludes response_format
319+
# from run_options (transforming it to text/text_format for OpenAI). Azure's agent
320+
# creation API requires the original response_format to build its own config format.
321+
if chat_options is None:
322+
chat_options = ChatOptions()
310323
# Agent name must be explicitly provided by the user.
311324
if self.agent_name is None:
312325
raise ServiceInitializationError(
@@ -341,8 +354,14 @@ async def _get_agent_reference_or_create(
341354
if "top_p" in run_options:
342355
args["top_p"] = run_options["top_p"]
343356

344-
if "response_format" in run_options:
345-
response_format = run_options["response_format"]
357+
# response_format is accessed from chat_options or additional_properties
358+
# since the base class excludes it from run_options
359+
response_format: Any = (
360+
chat_options.response_format
361+
if chat_options.response_format is not None
362+
else chat_options.additional_properties.get("response_format")
363+
)
364+
if response_format:
346365
args["text"] = PromptAgentDefinitionText(format=self._create_text_format_config(response_format))
347366

348367
# Combine instructions from messages and options
@@ -390,12 +409,12 @@ async def _prepare_options(
390409

391410
if not self._is_application_endpoint:
392411
# Application-scoped response APIs do not support "agent" property.
393-
agent_reference = await self._get_agent_reference_or_create(run_options, instructions)
412+
agent_reference = await self._get_agent_reference_or_create(run_options, instructions, chat_options)
394413
run_options["extra_body"] = {"agent": agent_reference}
395414

396415
# Remove properties that are not supported on request level
397416
# but were configured on agent level
398-
exclude = ["model", "tools", "response_format", "temperature", "top_p"]
417+
exclude = ["model", "tools", "response_format", "temperature", "top_p", "text", "text_format"]
399418

400419
for property in exclude:
401420
run_options.pop(property, None)

python/packages/azure-ai/tests/test_azure_ai_client.py

Lines changed: 111 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -723,9 +723,10 @@ async def test_azure_ai_client_agent_creation_with_response_format(
723723
mock_agent.version = "1.0"
724724
mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent)
725725

726-
run_options = {"model": "test-model", "response_format": ResponseFormatModel}
726+
run_options = {"model": "test-model"}
727+
chat_options = ChatOptions(response_format=ResponseFormatModel)
727728

728-
await client._get_agent_reference_or_create(run_options, None) # type: ignore
729+
await client._get_agent_reference_or_create(run_options, None, chat_options) # type: ignore
729730

730731
# Verify agent was created with response format configuration
731732
call_args = mock_project_client.agents.create_version.call_args
@@ -776,19 +777,18 @@ async def test_azure_ai_client_agent_creation_with_mapping_response_format(
776777
"additionalProperties": False,
777778
}
778779

779-
run_options = {
780-
"model": "test-model",
781-
"response_format": {
782-
"type": "json_schema",
783-
"json_schema": {
784-
"name": runtime_schema["title"],
785-
"strict": True,
786-
"schema": runtime_schema,
787-
},
780+
run_options = {"model": "test-model"}
781+
response_format_mapping = {
782+
"type": "json_schema",
783+
"json_schema": {
784+
"name": runtime_schema["title"],
785+
"strict": True,
786+
"schema": runtime_schema,
788787
},
789788
}
789+
chat_options = ChatOptions(response_format=response_format_mapping) # type: ignore
790790

791-
await client._get_agent_reference_or_create(run_options, None) # type: ignore
791+
await client._get_agent_reference_or_create(run_options, None, chat_options) # type: ignore
792792

793793
call_args = mock_project_client.agents.create_version.call_args
794794
created_definition = call_args[1]["definition"]
@@ -805,7 +805,7 @@ async def test_azure_ai_client_agent_creation_with_mapping_response_format(
805805
async def test_azure_ai_client_prepare_options_excludes_response_format(
806806
mock_project_client: MagicMock,
807807
) -> None:
808-
"""Test that prepare_options excludes response_format from final run options."""
808+
"""Test that prepare_options excludes response_format, text, and text_format from final run options."""
809809
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
810810

811811
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
@@ -815,7 +815,12 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
815815
patch.object(
816816
client.__class__.__bases__[0],
817817
"_prepare_options",
818-
return_value={"model": "test-model", "response_format": ResponseFormatModel},
818+
return_value={
819+
"model": "test-model",
820+
"response_format": ResponseFormatModel,
821+
"text": {"format": {"type": "json_schema", "name": "test"}},
822+
"text_format": ResponseFormatModel,
823+
},
819824
),
820825
patch.object(
821826
client,
@@ -825,8 +830,11 @@ async def test_azure_ai_client_prepare_options_excludes_response_format(
825830
):
826831
run_options = await client._prepare_options(messages, chat_options)
827832

828-
# response_format should be excluded from final run options
833+
# response_format, text, and text_format should be excluded from final run options
834+
# because they are configured at agent level, not request level
829835
assert "response_format" not in run_options
836+
assert "text" not in run_options
837+
assert "text_format" not in run_options
830838
# But extra_body should contain agent reference
831839
assert "extra_body" in run_options
832840
assert run_options["extra_body"]["agent"]["name"] == "test-agent"
@@ -1009,3 +1017,91 @@ async def test_azure_ai_chat_client_agent_with_tools() -> None:
10091017
assert response.text is not None
10101018
assert len(response.text) > 0
10111019
assert any(word in response.text.lower() for word in ["sunny", "25"])
1020+
1021+
1022+
class ReleaseBrief(BaseModel):
1023+
"""Structured output model for release brief."""
1024+
1025+
title: str = Field(description="A short title for the release.")
1026+
summary: str = Field(description="A brief summary of what was released.")
1027+
highlights: list[str] = Field(description="Key highlights from the release.")
1028+
model_config = ConfigDict(extra="forbid")
1029+
1030+
1031+
@pytest.mark.flaky
1032+
@skip_if_azure_ai_integration_tests_disabled
1033+
async def test_azure_ai_chat_client_agent_with_response_format() -> None:
1034+
"""Test ChatAgent with response_format (structured output) using AzureAIClient."""
1035+
async with (
1036+
temporary_chat_client(agent_name="ResponseFormatAgent") as chat_client,
1037+
ChatAgent(chat_client=chat_client) as agent,
1038+
):
1039+
response = await agent.run(
1040+
"Summarize the following release notes into a ReleaseBrief:\n\n"
1041+
"Version 2.0 Release Notes:\n"
1042+
"- Added new streaming API for real-time responses\n"
1043+
"- Improved error handling with detailed messages\n"
1044+
"- Performance boost of 50% in batch processing\n"
1045+
"- Fixed memory leak in connection pooling",
1046+
response_format=ReleaseBrief,
1047+
)
1048+
1049+
# Validate response
1050+
assert isinstance(response, AgentRunResponse)
1051+
assert response.value is not None
1052+
assert isinstance(response.value, ReleaseBrief)
1053+
1054+
# Validate structured output fields
1055+
brief = response.value
1056+
assert len(brief.title) > 0
1057+
assert len(brief.summary) > 0
1058+
assert len(brief.highlights) > 0
1059+
1060+
1061+
@pytest.mark.flaky
1062+
@skip_if_azure_ai_integration_tests_disabled
1063+
async def test_azure_ai_chat_client_agent_with_runtime_json_schema() -> None:
1064+
"""Test ChatAgent with runtime JSON schema (structured output) using AzureAIClient."""
1065+
runtime_schema = {
1066+
"title": "WeatherDigest",
1067+
"type": "object",
1068+
"properties": {
1069+
"location": {"type": "string"},
1070+
"conditions": {"type": "string"},
1071+
"temperature_c": {"type": "number"},
1072+
"advisory": {"type": "string"},
1073+
},
1074+
"required": ["location", "conditions", "temperature_c", "advisory"],
1075+
"additionalProperties": False,
1076+
}
1077+
1078+
async with (
1079+
temporary_chat_client(agent_name="RuntimeSchemaAgent") as chat_client,
1080+
ChatAgent(chat_client=chat_client) as agent,
1081+
):
1082+
response = await agent.run(
1083+
"Give a brief weather digest for Seattle.",
1084+
additional_chat_options={
1085+
"response_format": {
1086+
"type": "json_schema",
1087+
"json_schema": {
1088+
"name": runtime_schema["title"],
1089+
"strict": True,
1090+
"schema": runtime_schema,
1091+
},
1092+
},
1093+
},
1094+
)
1095+
1096+
# Validate response
1097+
assert isinstance(response, AgentRunResponse)
1098+
assert response.text is not None
1099+
1100+
# Parse JSON and validate structure
1101+
import json
1102+
1103+
parsed = json.loads(response.text)
1104+
assert "location" in parsed
1105+
assert "conditions" in parsed
1106+
assert "temperature_c" in parsed
1107+
assert "advisory" in parsed

python/packages/core/agent_framework/openai/_responses_client.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -439,16 +439,23 @@ async def _prepare_options(
439439
if (tool_choice := run_options.get("tool_choice")) and isinstance(tool_choice, dict) and "mode" in tool_choice:
440440
run_options["tool_choice"] = tool_choice["mode"]
441441

442-
# additional properties
442+
# additional properties (excluding response_format which is handled separately)
443443
additional_options = {
444-
key: value for key, value in chat_options.additional_properties.items() if value is not None
444+
key: value
445+
for key, value in chat_options.additional_properties.items()
446+
if value is not None and key != "response_format"
445447
}
446448
if additional_options:
447449
run_options.update(additional_options)
448450

449451
# response format and text config (after additional_properties so user can pass text via additional_properties)
450-
response_format = chat_options.response_format
451-
text_config = run_options.pop("text", None)
452+
# Check both chat_options.response_format and additional_properties for response_format
453+
response_format: Any = (
454+
chat_options.response_format
455+
if chat_options.response_format is not None
456+
else chat_options.additional_properties.get("response_format")
457+
)
458+
text_config: Any = run_options.pop("text", None)
452459
response_format, text_config = self._prepare_response_and_text_format(
453460
response_format=response_format, text_config=text_config
454461
)

python/packages/core/tests/openai/test_openai_responses_client.py

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2355,3 +2355,91 @@ async def test_openai_responses_client_agent_local_mcp_tool() -> None:
23552355
assert len(response.text) > 0
23562356
# Should contain Azure-related content since it's asking about Azure CLI
23572357
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
2358+
2359+
2360+
class ReleaseBrief(BaseModel):
2361+
"""Structured output model for release brief testing."""
2362+
2363+
title: str
2364+
summary: str
2365+
highlights: list[str]
2366+
model_config = {"extra": "forbid"}
2367+
2368+
2369+
@pytest.mark.flaky
2370+
@skip_if_openai_integration_tests_disabled
2371+
async def test_openai_responses_client_agent_with_response_format_pydantic() -> None:
2372+
"""Integration test for response_format with Pydantic model using OpenAI Responses Client."""
2373+
async with ChatAgent(
2374+
chat_client=OpenAIResponsesClient(),
2375+
instructions="You are a helpful assistant that returns structured JSON responses.",
2376+
) as agent:
2377+
response = await agent.run(
2378+
"Summarize the following release notes into a ReleaseBrief:\n\n"
2379+
"Version 2.0 Release Notes:\n"
2380+
"- Added new streaming API for real-time responses\n"
2381+
"- Improved error handling with detailed messages\n"
2382+
"- Performance boost of 50% in batch processing\n"
2383+
"- Fixed memory leak in connection pooling",
2384+
response_format=ReleaseBrief,
2385+
)
2386+
2387+
# Validate response
2388+
assert isinstance(response, AgentRunResponse)
2389+
assert response.value is not None
2390+
assert isinstance(response.value, ReleaseBrief)
2391+
2392+
# Validate structured output fields
2393+
brief = response.value
2394+
assert len(brief.title) > 0
2395+
assert len(brief.summary) > 0
2396+
assert len(brief.highlights) > 0
2397+
2398+
2399+
@pytest.mark.flaky
2400+
@skip_if_openai_integration_tests_disabled
2401+
async def test_openai_responses_client_agent_with_runtime_json_schema() -> None:
2402+
"""Integration test for response_format with runtime JSON schema using OpenAI Responses Client."""
2403+
runtime_schema = {
2404+
"title": "WeatherDigest",
2405+
"type": "object",
2406+
"properties": {
2407+
"location": {"type": "string"},
2408+
"conditions": {"type": "string"},
2409+
"temperature_c": {"type": "number"},
2410+
"advisory": {"type": "string"},
2411+
},
2412+
"required": ["location", "conditions", "temperature_c", "advisory"],
2413+
"additionalProperties": False,
2414+
}
2415+
2416+
async with ChatAgent(
2417+
chat_client=OpenAIResponsesClient(),
2418+
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
2419+
) as agent:
2420+
response = await agent.run(
2421+
"Give a brief weather digest for Seattle.",
2422+
additional_chat_options={
2423+
"response_format": {
2424+
"type": "json_schema",
2425+
"json_schema": {
2426+
"name": runtime_schema["title"],
2427+
"strict": True,
2428+
"schema": runtime_schema,
2429+
},
2430+
},
2431+
},
2432+
)
2433+
2434+
# Validate response
2435+
assert isinstance(response, AgentRunResponse)
2436+
assert response.text is not None
2437+
2438+
# Parse JSON and validate structure
2439+
import json
2440+
2441+
parsed = json.loads(response.text)
2442+
assert "location" in parsed
2443+
assert "conditions" in parsed
2444+
assert "temperature_c" in parsed
2445+
assert "advisory" in parsed

0 commit comments

Comments
 (0)