feat(ext): support structured output in AzureAIChatCompletionClient

python/packages/autogen-ext/src/autogen_ext/models/azure/_azure_ai_client.py

@@ -38,6 +38,7 @@
     ImageContentItem,
     ImageDetailLevel,
     ImageUrl,
+    JsonSchemaFormat,
     StreamingChatChoiceUpdate,
     StreamingChatCompletionsUpdate,
     TextContentItem,
@@ -345,8 +346,12 @@ def _validate_model_info(
                 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 self.model_info["structured_output"] is False:
+                    raise ValueError("Model does not support structured output")
+                schema_name = re.sub(r"[^a-zA-Z0-9_-]", "_", json_output.__name__)[:64]
+                create_args["response_format"] = JsonSchemaFormat(
+                    name=schema_name, schema=json_output.model_json_schema(), strict=False
+                )
 
             if json_output is True and "response_format" not in create_args:
                 create_args["response_format"] = "json_object"

python/packages/autogen-ext/tests/models/test_azure_ai_model_client.py

@@ -28,7 +28,9 @@
 from azure.ai.inference.models import (
     FunctionCall as AzureFunctionCall,
 )
+from azure.ai.inference.models import JsonSchemaFormat
 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,33 @@ 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 _Answer(BaseModel):
+    value: int
+
+
+def _make_azure_client(monkeypatch: pytest.MonkeyPatch, structured_output: bool) -> AzureAIChatCompletionClient:
+    mock_client = MagicMock()
+    monkeypatch.setattr(ChatCompletionsClient, "__new__", lambda cls, *a, **kw: mock_client)
+    return AzureAIChatCompletionClient(
+        endpoint="endpoint",
+        credential=AzureKeyCredential("api_key"),
+        model="model",
+        model_info={"json_output": True, "function_calling": False, "vision": False, "family": "unknown", "structured_output": structured_output},
+    )
+
+
+def test_structured_output_raises_when_not_supported(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = _make_azure_client(monkeypatch, structured_output=False)
+    create_args: dict[str, Any] = {}
+    with pytest.raises(ValueError, match="does not support structured output"):
+        client._validate_model_info([], [], _Answer, create_args)
+
+
+def test_structured_output_sets_json_schema_format(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = _make_azure_client(monkeypatch, structured_output=True)
+    create_args: dict[str, Any] = {}
+    client._validate_model_info([], [], _Answer, create_args)
+    assert isinstance(create_args.get("response_format"), JsonSchemaFormat)
+    assert create_args["response_format"].name == "_Answer"