From 91a87e54bf0a4b4bbfc2c0ad86e9d005898461c9 Mon Sep 17 00:00:00 2001 From: abhnvgrg Date: Sat, 12 Sep 2026 18:33:05 +0530 Subject: [PATCH] feat(azure): add structured output support for AzureAIChatCompletionClient Closes #5957 --- .../models/azure/_azure_ai_client.py | 101 +++++- .../models/azure/config/__init__.py | 2 +- .../models/test_azure_ai_model_client.py | 325 +++++++++++++++++- 3 files changed, 415 insertions(+), 13 deletions(-) diff --git a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py index 16d56b57f956..3c9a53d648fb 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py @@ -38,6 +38,7 @@ ImageContentItem, ImageDetailLevel, ImageUrl, + JsonSchemaFormat, StreamingChatChoiceUpdate, StreamingChatCompletionsUpdate, TextContentItem, @@ -193,7 +194,7 @@ class AzureAIChatCompletionClient(ChatCompletionClient): temperature: (optional,float) top_p: (optional,float) max_tokens: (optional,int) - response_format: (optional, literal["text", "json_object"]) + response_format: (optional, Union[literal["text", "json_object"], JsonSchemaFormat]) stop: (optional,List[str]) tools: (optional,List[ChatCompletionsToolDefinition]) tool_choice: (optional,Union[str, ChatCompletionsToolChoicePreset, ChatCompletionsNamedToolChoice]]) @@ -284,6 +285,57 @@ async def main(): if __name__ == "__main__": asyncio.run(main()) + To use structured output, set the `json_output` parameter to a Pydantic model class. + The model must support structured output, which is indicated by the `structured_output` + field of the `model_info`. The response content is a JSON string that conforms to the + schema of the model class: + + .. code-block:: python + + import asyncio + import os + + from autogen_core.models import UserMessage + from autogen_ext.models.azure import AzureAIChatCompletionClient + from azure.core.credentials import AzureKeyCredential + from pydantic import BaseModel + + + class CityInfo(BaseModel): + name: str + country: str + population: int + + + async def main(): + client = AzureAIChatCompletionClient( + model="gpt-4o", + endpoint="https://models.github.ai/inference", + credential=AzureKeyCredential(os.environ["GITHUB_TOKEN"]), + model_info={ + "json_output": True, + "function_calling": True, + "vision": True, + "family": "unknown", + "structured_output": True, + }, + ) + + result = await client.create( + [UserMessage(content="Give me info about Paris", source="user")], + json_output=CityInfo, + ) + assert isinstance(result.content, str) + city = CityInfo.model_validate_json(result.content) + print(city.name) # "Paris" + + # Close the client. + await client.close() + + + if __name__ == "__main__": + asyncio.run(main()) + """ @@ -341,15 +393,26 @@ def _validate_model_info( raise ValueError("Model does not support vision and image was provided") if json_output is not None: - if self.model_info["json_output"] is False and json_output is True: - raise ValueError("Model does not support JSON output") - - if isinstance(json_output, type): - # TODO: we should support this in the future. - raise ValueError("Structured output is not currently supported for AzureAIChatCompletionClient") - - if json_output is True and "response_format" not in create_args: - create_args["response_format"] = "json_object" + if json_output is True: + # JSON mode. + if self.model_info["json_output"] is False: + raise ValueError("Model does not support JSON output") + if "response_format" not in create_args: + create_args["response_format"] = "json_object" + elif json_output is False: + # Text mode, nothing to do. + pass + elif isinstance(json_output, type) and issubclass(json_output, BaseModel): + # Structured output mode with a Pydantic model class. + if self.model_info.get("structured_output") is False: + raise ValueError("Model does not support structured output") + create_args["response_format"] = JsonSchemaFormat( + name=json_output.__name__, + schema=json_output.model_json_schema(), + strict=True, + ) + else: + raise ValueError(f"json_output must be a boolean or a Pydantic model class, got {type(json_output)}") if self.model_info["json_output"] is False and json_output is True: raise ValueError("Model does not support JSON output") @@ -446,6 +509,13 @@ async def create( if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1: thought, content = parse_r1_content(content) + if isinstance(json_output, type) and issubclass(json_output, BaseModel) and isinstance(content, str): + # Validate the structured output, the content is kept as the raw JSON string. + try: + json_output.model_validate_json(content) + except Exception as e: + raise ValueError(f"Failed to parse structured output: {e}. Raw content: {content}") from e + response = CreateResult( finish_reason=finish_reason, # type: ignore content=content, @@ -520,7 +590,9 @@ async def create_stream( ) assert isinstance(chunk, StreamingChatCompletionsUpdate) choice = chunk.choices[0] if len(chunk.choices) > 0 else None - if choice and choice.finish_reason is not None: + # The SDK types finish_reason as required, but the service omits it on + # intermediate chunks, so use a truthiness check instead of a None check. + if choice and choice.finish_reason: if isinstance(choice.finish_reason, CompletionsFinishReason): finish_reason = cast(FinishReasons, choice.finish_reason.value) # Handle special case for TOOL_CALLS finish reason @@ -579,6 +651,13 @@ async def create_stream( if isinstance(content, str) and self._model_info["family"] == ModelFamily.R1: thought, content = parse_r1_content(content) + if isinstance(json_output, type) and issubclass(json_output, BaseModel) and isinstance(content, str): + # Validate the structured output, the content is kept as the raw JSON string. + try: + json_output.model_validate_json(content) + except Exception as e: + raise ValueError(f"Failed to parse structured output: {e}. Raw content: {content}") from e + result = CreateResult( finish_reason=finish_reason, content=content, diff --git a/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py b/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py index 38cf34b5378a..ac1737678f76 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/azure/config/__init__.py @@ -33,7 +33,7 @@ class AzureAICreateArguments(TypedDict, total=False): temperature: Optional[float] top_p: Optional[float] max_tokens: Optional[int] - response_format: Optional[Literal["text", "json_object"]] + response_format: Optional[Union[Literal["text", "json_object"], JsonSchemaFormat]] stop: Optional[List[str]] tools: Optional[List[ChatCompletionsToolDefinition]] tool_choice: Optional[Union[str, ChatCompletionsToolChoicePreset, ChatCompletionsNamedToolChoice]] diff --git a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py index dfc7af07302f..c4a865606915 100644 --- a/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py +++ b/python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py @@ -2,7 +2,7 @@ import logging import os from datetime import datetime -from typing import Any, AsyncGenerator, List, Type, Union +from typing import Any, AsyncGenerator, Dict, List, Type, Union from unittest.mock import AsyncMock, MagicMock import pytest @@ -21,6 +21,7 @@ ChatResponseMessage, CompletionsFinishReason, CompletionsUsage, + JsonSchemaFormat, StreamingChatChoiceUpdate, StreamingChatCompletionsUpdate, StreamingChatResponseMessageUpdate, @@ -29,6 +30,7 @@ FunctionCall as AzureFunctionCall, ) from azure.core.credentials import AzureKeyCredential +from pydantic import BaseModel async def _mock_create_stream(*args: Any, **kwargs: Any) -> AsyncGenerator[StreamingChatCompletionsUpdate, None]: @@ -973,3 +975,324 @@ async def test_azure_ai_tool_choice_specific_tool_streaming( assert final_result.content[0].name == "process_text" assert final_result.content[0].arguments == '{"input": "hello"}' assert final_result.thought == "Let me process this for you." + + +class CityInfo(BaseModel): + """Structured output type used by the structured output tests.""" + + name: str + country: str + population: int + + +@pytest.fixture +def structured_output_client(monkeypatch: pytest.MonkeyPatch) -> AzureAIChatCompletionClient: + """ + Returns a client that simulates structured JSON output for structured output tests. + """ + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> ChatCompletions: + await asyncio.sleep(0.01) + return ChatCompletions( + id="id", + created=datetime.now(), + model="model", + choices=[ + ChatChoice( + index=0, + finish_reason="stop", + message=ChatResponseMessage( + role="assistant", + content='{"name": "Paris", "country": "France", "population": 2161000}', + ), + ) + ], + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + return AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": True, + "function_calling": False, + "vision": False, + "family": "unknown", + "structured_output": True, + }, + model="model", + ) + + +@pytest.mark.asyncio +async def test_azure_ai_structured_output_create( + structured_output_client: AzureAIChatCompletionClient, +) -> None: + """Test structured output with json_output set to a Pydantic model class.""" + result = await structured_output_client.create( + messages=[UserMessage(content="Tell me about Paris", source="user")], + json_output=CityInfo, + ) + + assert isinstance(result.content, str) + city = CityInfo.model_validate_json(result.content) + assert city.name == "Paris" + assert city.country == "France" + assert city.population == 2161000 + + +@pytest.mark.asyncio +async def test_azure_ai_structured_output_unsupported(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Ensures error is raised if we request structured output but the model_info doesn't support it. + """ + mock_client = MagicMock() + mock_client.close = AsyncMock() + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + client = AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": False, + "function_calling": False, + "vision": False, + "family": "unknown", + "structured_output": False, + }, + model="model", + ) + + with pytest.raises(ValueError) as exc: + await client.create( + messages=[UserMessage(content="Hello", source="user")], + json_output=CityInfo, + ) + assert "Model does not support structured output" in str(exc.value) + + +@pytest.mark.asyncio +async def test_azure_ai_structured_output_invalid_content(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Ensures error is raised if the model returns content that doesn't conform to the schema. + """ + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> ChatCompletions: + await asyncio.sleep(0.01) + return ChatCompletions( + id="id", + created=datetime.now(), + model="model", + choices=[ + ChatChoice( + index=0, + finish_reason="stop", + message=ChatResponseMessage(role="assistant", content="Paris is the capital of France."), + ) + ], + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + client = AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": True, + "function_calling": False, + "vision": False, + "family": "unknown", + "structured_output": True, + }, + model="model", + ) + + with pytest.raises(ValueError) as exc: + await client.create( + messages=[UserMessage(content="Tell me about Paris", source="user")], + json_output=CityInfo, + ) + assert "Failed to parse structured output" in str(exc.value) + + +@pytest.fixture +def structured_output_stream_client(monkeypatch: pytest.MonkeyPatch) -> AzureAIChatCompletionClient: + """ + Returns a client that streams structured JSON output. + """ + + json_parts = ['{"name": "Paris"', ', "country": "France"', ', "population": 2161000}'] + + async def _mock_structured_stream( + *args: Any, **kwargs: Any + ) -> AsyncGenerator[StreamingChatCompletionsUpdate, None]: + for part in json_parts: + await asyncio.sleep(0.01) + yield StreamingChatCompletionsUpdate( + id="id", + choices=[ + StreamingChatChoiceUpdate( + index=0, + finish_reason="stop", + delta=StreamingChatResponseMessageUpdate(role="assistant", content=part), + ) + ], + created=datetime.now(), + model="model", + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30) + if part == json_parts[-1] + else None, + ) + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> Any: + if kwargs.get("stream", False): + return _mock_structured_stream(*args, **kwargs) + return None + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + return AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": True, + "function_calling": False, + "vision": False, + "family": "unknown", + "structured_output": True, + }, + model="model", + ) + + +@pytest.mark.asyncio +async def test_azure_ai_structured_output_stream( + structured_output_stream_client: AzureAIChatCompletionClient, +) -> None: + """Test structured output with streaming.""" + chunks: List[Union[str, CreateResult]] = [] + async for chunk in structured_output_stream_client.create_stream( + messages=[UserMessage(content="Tell me about Paris", source="user")], + json_output=CityInfo, + ): + chunks.append(chunk) + + final_result = chunks[-1] + assert isinstance(final_result, CreateResult) + assert isinstance(final_result.content, str) + city = CityInfo.model_validate_json(final_result.content) + assert city.name == "Paris" + assert city.country == "France" + assert city.population == 2161000 + + +@pytest.mark.asyncio +async def test_azure_ai_structured_output_passes_response_format(monkeypatch: pytest.MonkeyPatch) -> None: + """Verify that a JsonSchemaFormat is passed to the underlying client.complete().""" + captured_kwargs: Dict[str, Any] = {} + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def mock_complete(*args: Any, **kwargs: Any) -> ChatCompletions: + captured_kwargs.update(kwargs) + await asyncio.sleep(0.01) + return ChatCompletions( + id="id", + created=datetime.now(), + model="model", + choices=[ + ChatChoice( + index=0, + finish_reason="stop", + message=ChatResponseMessage( + role="assistant", + content='{"name": "Paris", "country": "France", "population": 2161000}', + ), + ) + ], + usage=CompletionsUsage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + mock_client.complete = mock_complete + + def mock_new(cls: Type[ChatCompletionsClient], *args: Any, **kwargs: Any) -> MagicMock: + return mock_client + + monkeypatch.setattr(ChatCompletionsClient, "__new__", mock_new) + + client = AzureAIChatCompletionClient( + endpoint="endpoint", + credential=AzureKeyCredential("api_key"), + model_info={ + "json_output": True, + "function_calling": False, + "vision": False, + "family": "unknown", + "structured_output": True, + }, + model="model", + ) + + await client.create( + messages=[UserMessage(content="Tell me about Paris", source="user")], + json_output=CityInfo, + ) + + assert "response_format" in captured_kwargs + response_format = captured_kwargs["response_format"] + assert isinstance(response_format, JsonSchemaFormat) + assert response_format.name == "CityInfo" + assert response_format.strict is True + assert response_format.schema == CityInfo.model_json_schema() + + +@pytest.mark.asyncio +async def test_azure_ai_json_output_bool(structured_output_client: AzureAIChatCompletionClient) -> None: + """Verify that json_output=True still uses the json_object response format.""" + result = await structured_output_client.create( + messages=[UserMessage(content="Tell me about Paris", source="user")], + json_output=True, + ) + assert isinstance(result.content, str) + + +@pytest.mark.asyncio +async def test_azure_ai_json_output_invalid_type(structured_output_client: AzureAIChatCompletionClient) -> None: + """Ensures error is raised if json_output is neither a boolean nor a Pydantic model class.""" + with pytest.raises(ValueError) as exc: + await structured_output_client.create( + messages=[UserMessage(content="Hello", source="user")], + json_output="json", # type: ignore + ) + assert "json_output must be a boolean or a Pydantic model class" in str(exc.value)