From 57012abfa8a56062185aa1ceea63eda2ed705310 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 26 Jul 2026 22:18:38 -0700 Subject: [PATCH 1/2] [integrations][python] Apply Azure OpenAI native structured output Translate a BaseModel output schema into the provider's native response_format json-schema parameter instead of rejecting it, so callers with a schema are no longer forced onto the prompt-engineering fallback. Capability resolves against model_of_azure_deployment rather than model, because on Azure the latter is a user-chosen deployment name that carries no information about the backing model. The native path additionally requires an api_version at or above Azure's documented 2024-08-01 floor; below it the schema is never sent, so an older api-version cannot produce an unconstrained response that the caller mistakes for a structured one. A caller-supplied response_format that the native path would overwrite now raises rather than being silently replaced or reaching the SDK as a duplicate keyword argument. --- .../azure/azure_openai_chat_model.py | 145 ++++++- ...t_azure_openai_native_structured_output.py | 410 ++++++++++++++++++ .../test_output_schema_param_declared.py | 24 +- 3 files changed, 572 insertions(+), 7 deletions(-) create mode 100644 python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py diff --git a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py index b653b4d67..10a07f56c 100644 --- a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py @@ -19,7 +19,14 @@ from typing import Any, Dict, List, Sequence from openai import NOT_GIVEN, AzureOpenAI -from pydantic import Field, PrivateAttr + +# Private SDK module (leading underscore): the openai client itself uses this helper to +# build the strict json_schema for response_format, and there is no public re-export. It +# has existed at this path since the structured-output support in openai 1.66.3 (the +# pinned minimum). A future openai bump that moves it will fail loudly on import here. +from openai.lib._pydantic import to_strict_json_schema +from pydantic import BaseModel, Field, PrivateAttr +from typing_extensions import override from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage @@ -40,6 +47,70 @@ {"model", "model_of_azure_deployment", "temperature", "max_tokens", "logprobs"} ) +# Models with documented json_schema strict Structured Outputs support. Source of truth: +# https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/structured-outputs +# +# Matching is exact, never by prefix: Azure exposes a deployment's model name and model +# version as separate properties, so a name carries no version to discriminate on. The +# documented list includes gpt-4o only at versions 2024-08-06 and 2024-11-20 while +# version 2024-05-13 is unsupported, so a bare "gpt-4o" is ambiguous and is deliberately +# absent from the set below. An unrecognized name reports not-capable and degrades to +# the prompt fallback rather than failing at the provider. +# +# The source list prints "gpt-5.1-codex mini" with a space; it is transcribed hyphenated +# here because Azure model identifiers do not contain spaces. +_NATIVE_STRUCTURED_OUTPUT_MODELS = frozenset( + { + "gpt-5.1-codex", + "gpt-5.1-codex-mini", + "gpt-5.1", + "gpt-5.1-chat", + "gpt-5-pro", + "gpt-5-codex", + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "codex-mini", + "o3-pro", + "o3-mini", + "o1", + "gpt-4o-mini", + "gpt-4.1", + "gpt-4.1-nano", + "gpt-4.1-mini", + "o4-mini", + "o3", + } +) + +# Date prefix of 2024-08-01-preview, the earliest api-version Azure documents as +# supporting structured outputs. +_MIN_STRUCTURED_OUTPUT_API_VERSION = "2024-08-01" + + +def _native_response_format(output_schema: Any) -> Dict[str, Any] | None: + """Build the ``response_format`` for a native structured-output request. + + Returns ``None`` (leaving behavior unchanged) unless the schema is a ``BaseModel`` + subclass. A ``RowTypeInfo`` schema is skipped so it keeps the prompt-engineering + fallback. + """ + if output_schema is None: + return None + model = ( + output_schema.output_schema if isinstance(output_schema, OutputSchema) else None + ) + if not (isinstance(model, type) and issubclass(model, BaseModel)): + return None + return { + "type": "json_schema", + "json_schema": { + "name": model.__name__, + "schema": to_strict_json_schema(model), + "strict": True, + }, + } + class AzureOpenAIChatModelConnection(BaseChatModelConnection): """The connection to the Azure OpenAI LLM. @@ -114,6 +185,41 @@ def client(self) -> AzureOpenAI: ) return self._client + @override + def supports_native_structured_output(self, effective_model: str | None) -> bool: + """Whether Azure documents json_schema strict support for ``effective_model``. + + ``effective_model`` is the model backing an Azure deployment, not the deployment + name. See the module-level allowlist for the source of truth and for why the + match is exact. An unrecognized model reports ``False`` so it degrades to the + prompt-engineering fallback rather than failing at the provider. + + Reads no instance state, so it stays answerable on an instance that was never + initialized, where any field access would raise. + """ + if not effective_model: + return False + return effective_model in _NATIVE_STRUCTURED_OUTPUT_MODELS + + def _api_version_supports_structured_output(self) -> bool: + """Whether the configured api-version reaches the structured-output floor. + + Azure documents ``2024-08-01-preview`` as the first api-version supporting + structured outputs, and whether an older version rejects ``response_format`` or + silently ignores it is not documented. The request therefore never carries + ``response_format`` below the floor, which is safe under either behavior. + + The comparison assumes the documented api-version form, a zero-padded + ``YYYY-MM-DD`` date optionally suffixed ``-preview``; over that form comparing + the leading date lexicographically is exact. The GA ``v1`` literal sorts above + the floor, which matches Azure documenting ``v1`` as supporting structured + outputs. This is not a validator: a value of any other shape is not classified + reliably, and the service rejects an api-version it does not recognize. + """ + if not self.api_version: + return False + return self.api_version[:10] >= _MIN_STRUCTURED_OUTPUT_API_VERSION + def chat( self, messages: Sequence[ChatMessage], @@ -130,10 +236,13 @@ def chat( tools : Optional[List] List of tools that can be called by the model output_schema : OutputSchema | None - Rejected when non-``None``: this connection has no native structured-output - translation, so callers stay on the prompt-engineering fallback. - Declaring the parameter keeps a caller-supplied schema out of - ``**kwargs``, which is forwarded to the provider SDK. + The schema the response should conform to, or ``None`` for an unconstrained + response. Native structured output is applied only for a ``BaseModel`` + schema, on a deployment whose backing model the provider documents as + capable, and with an api-version that supports it; a ``RowTypeInfo`` schema, + an incapable model, or an older api-version keeps the prompt-engineering + fallback. Where native output applies, a caller-supplied + ``response_format`` conflicts with it and raises ``ValueError``. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -143,7 +252,6 @@ def chat( ChatMessage Model response message """ - self._reject_unsupported_output_schema(output_schema) tool_specs = None if tools is not None: tool_specs = [to_openai_tool(metadata=tool.metadata) for tool in tools] @@ -165,6 +273,31 @@ def chat( ) raise ValueError(msg) + # Capability belongs to the model backing the deployment, so it is the input to + # the check. The deployment name is chosen by the user and carries none. + if ( + output_schema is not None + and self.supports_native_structured_output(model_of_azure_deployment) + and self._api_version_supports_structured_output() + ): + response_format = _native_response_format(output_schema) + if response_format is not None: + caller_response_format = ( + "response_format" in kwargs + or "response_format" in additional_kwargs + ) + if caller_response_format: + msg = ( + f"The {response_format['json_schema']['name']} output schema " + f"is sent as response_format on deployment " + f"'{azure_deployment}', so response_format must not also be " + f"passed as a kwarg or in additional_kwargs. Remove that " + f"value, or omit output_schema to set response_format " + f"directly." + ) + raise ValueError(msg) + kwargs["response_format"] = response_format + response = self.client.chat.completions.create( # Azure OpenAI APIs use Azure deployment name as the model parameter model=azure_deployment, diff --git a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py new file mode 100644 index 000000000..278e7cd07 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py @@ -0,0 +1,410 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +from typing import Any +from unittest.mock import MagicMock + +import pytest +from pydantic import BaseModel +from pyflink.common.typeinfo import Types + +from flink_agents.api.agents.types import OutputSchema +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.integrations.chat_models.azure.azure_openai_chat_model import ( + AzureOpenAIChatModelConnection, +) +from flink_agents.plan.function import PythonFunction +from flink_agents.plan.tools.function_tool import FunctionTool + +# A deployment name is chosen by the user and carries no capability information, so +# every chat() call here uses one that is not a model name. +DEPLOYMENT = "my-deployment" + +CAPABLE_API_VERSION = "2024-08-01-preview" + +BELOW_FLOOR_API_VERSION = "2024-02-01" + +CALLER_RESPONSE_FORMAT = {"type": "json_object"} + + +class Person(BaseModel): + """A representative BaseModel output schema.""" + + name: str + age: int + + +ROW_TYPE = Types.ROW_NAMED(["name"], [Types.STRING()]) + + +def _add(a: int, b: int) -> int: + """Add two integers. + + Parameters + ---------- + a : int + first + b : int + second + + Returns: + ------- + int + sum + """ + return a + b + + +def _connection( + api_version: str = CAPABLE_API_VERSION, +) -> AzureOpenAIChatModelConnection: + conn = AzureOpenAIChatModelConnection( + name="azure_openai", + api_key="test-key", + azure_endpoint="https://example.openai.azure.com", + api_version=api_version, + ) + mock_client = MagicMock() + mock_message = MagicMock() + mock_message.role = "assistant" + mock_message.content = "ok" + mock_message.tool_calls = None + mock_client.chat.completions.create.return_value.choices = [ + MagicMock(message=mock_message) + ] + mock_client.chat.completions.create.return_value.usage = None + conn._client = mock_client + return conn + + +def _create_call_kwargs(conn: AzureOpenAIChatModelConnection) -> dict[str, Any]: + return conn.client.chat.completions.create.call_args.kwargs + + +def _chat_with_caller_response_format( + conn: AzureOpenAIChatModelConnection, + *, + model_of_azure_deployment: str, + in_additional_kwargs: bool, + schema: Any = Person, +) -> None: + """Chat with a caller-supplied response_format, optionally with an output schema. + + The value travels either inside additional_kwargs or as a direct kwarg; both end + up in the same create() call. A ``schema`` of ``None`` sends no output schema at + all rather than an empty one. + """ + channel = ( + {"additional_kwargs": {"response_format": CALLER_RESPONSE_FORMAT}} + if in_additional_kwargs + else {"response_format": CALLER_RESPONSE_FORMAT} + ) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment=model_of_azure_deployment, + output_schema=None if schema is None else OutputSchema(output_schema=schema), + **channel, + ) + + +def test_native_applied_for_capable_deployment_model() -> None: + """response_format json_schema strict applied for a BaseModel on a capable model.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + response_format = _create_call_kwargs(conn)["response_format"] + assert response_format["type"] == "json_schema" + assert response_format["json_schema"]["name"] == "Person" + assert response_format["json_schema"]["strict"] is True + assert response_format["json_schema"]["schema"]["additionalProperties"] is False + + +def test_capable_native_request_still_targets_the_deployment() -> None: + """The native branch leaves `model` as the deployment name. + + Capability is keyed on the backing model, but the provider is still addressed by + deployment; substituting one for the other would route the call to a deployment + that may not exist on the resource. + """ + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + assert _create_call_kwargs(conn)["model"] == DEPLOYMENT + + +def test_native_not_applied_when_deployment_model_absent() -> None: + """Native NOT applied when the backing model of the deployment is unknown.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_for_unknown_deployment_model() -> None: + """Native NOT applied for a backing model outside the allowlist.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="some-unknown-model", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_for_bare_gpt_4o() -> None: + """Native NOT applied for a bare `gpt-4o` backing model. + + Azure carries model name and model version as separate properties, so a bare + `gpt-4o` may be the 2024-05-13 version, which predates structured output support. + """ + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +@pytest.mark.parametrize("api_version", ["2024-10-21", "v1"]) +def test_native_applied_for_other_api_versions_above_floor(api_version: str) -> None: + """Native applied for a later GA date and for the GA `v1` literal. + + `v1` pins the non-date half of the comparison: it is admitted because it sorts + above the floor, not because it parses as a later date. Rewriting the gate to + parse dates would drop it while every date case still passed. + """ + conn = _connection(api_version=api_version) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" in _create_call_kwargs(conn) + + +def test_native_not_applied_when_api_version_below_floor() -> None: + """Native NOT applied when the configured api-version predates the floor.""" + conn = _connection(api_version=BELOW_FLOOR_API_VERSION) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_when_api_version_empty() -> None: + """Native NOT applied when no api-version is configured. + + The empty string stands in for an absent api-version: the field is required at + construction, so `None` is rejected by validation before chat() is ever reached. + """ + conn = _connection(api_version="") + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_when_schema_none() -> None: + """Native NOT applied when no output schema is supplied.""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=None, + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_not_applied_for_row_type_info() -> None: + """Native NOT applied for a RowTypeInfo schema (BaseModel-only scope).""" + conn = _connection() + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=ROW_TYPE), + ) + assert "response_format" not in _create_call_kwargs(conn) + + +def test_native_applied_even_when_tools_bound() -> None: + """Native applied for a BaseModel even when tools are bound. + + Azure documents structured outputs as unsupported with parallel function calls, + which constrains strict tool schemas rather than the response_format this branch + sets, so binding tools does not gate it. + """ + conn = _connection() + tool = FunctionTool(func=PythonFunction.from_callable(_add)) + conn.chat( + [ChatMessage(role=MessageRole.USER, content="hi")], + tools=[tool], + model=DEPLOYMENT, + model_of_azure_deployment="gpt-4o-mini", + output_schema=OutputSchema(output_schema=Person), + ) + assert "response_format" in _create_call_kwargs(conn) + + +@pytest.mark.parametrize("in_additional_kwargs", [True, False]) +def test_caller_response_format_conflicts_with_native_schema( + in_additional_kwargs: bool, +) -> None: + """A caller-supplied response_format alongside a natively applied schema raises. + + Both values would otherwise reach the same create() call, where the direct kwarg + is silently overwritten and the additional_kwargs one becomes a duplicate keyword + argument reported by the SDK rather than by this connection. + """ + conn = _connection() + with pytest.raises(ValueError, match="response_format") as excinfo: + _chat_with_caller_response_format( + conn, + model_of_azure_deployment="gpt-4o-mini", + in_additional_kwargs=in_additional_kwargs, + ) + assert "Person" in str(excinfo.value) + + +@pytest.mark.parametrize("in_additional_kwargs", [True, False]) +@pytest.mark.parametrize( + ("api_version", "model_of_azure_deployment", "schema"), + [ + (CAPABLE_API_VERSION, "gpt-4o", Person), + (CAPABLE_API_VERSION, "gpt-4o-mini", ROW_TYPE), + (CAPABLE_API_VERSION, "gpt-4o-mini", None), + (BELOW_FLOOR_API_VERSION, "gpt-4o-mini", Person), + ], + ids=[ + "incapable_model", + "row_type_info_schema", + "no_output_schema", + "api_version_below_floor", + ], +) +def test_caller_response_format_survives_when_native_is_skipped( + api_version: str, + model_of_azure_deployment: str, + schema: Any, + in_additional_kwargs: bool, +) -> None: + """The same caller input passes through untouched wherever native output is skipped. + + Native output is skipped for an incapable backing model, for a schema kind outside + the natively translatable set, for no schema at all, and for an api-version below + the floor. Only the branch that actually sends a schema as response_format may + reject the caller's own value, so identical caller code has to keep working along + every one of those paths, including the no-schema path taken by any caller that + drives response_format itself. + """ + conn = _connection(api_version=api_version) + _chat_with_caller_response_format( + conn, + model_of_azure_deployment=model_of_azure_deployment, + in_additional_kwargs=in_additional_kwargs, + schema=schema, + ) + assert _create_call_kwargs(conn)["response_format"] is CALLER_RESPONSE_FORMAT + + +@pytest.mark.parametrize( + "model", + [ + "gpt-5.1-codex", + "gpt-5.1-codex-mini", + "gpt-5.1", + "gpt-5.1-chat", + "gpt-5-pro", + "gpt-5-codex", + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "codex-mini", + "o3-pro", + "o3-mini", + "o1", + "gpt-4o-mini", + "gpt-4.1", + "gpt-4.1-nano", + "gpt-4.1-mini", + "o4-mini", + "o3", + ], +) +def test_capability_predicate_accepts_capable_models(model: str) -> None: + """The capability predicate accepts every documented capable Azure model name. + + The list is the whole allowlist, so dropping an entry is caught rather than only + narrowing capability silently. + """ + assert _connection().supports_native_structured_output(model) is True + + +@pytest.mark.parametrize( + "model", + [ + "gpt-4o", + "gpt-35-turbo", + "gpt-4", + "gpt-4o-2024-08-06", + "some-unknown-model", + None, + "", + ], +) +def test_capability_predicate_rejects_incapable_models(model: str | None) -> None: + """The capability predicate rejects incapable, version-suffixed, and empty names. + + A version-suffixed value such as `gpt-4o-2024-08-06` is an OpenAI snapshot name, + not a name Azure reports as the model behind a deployment. + """ + assert _connection().supports_native_structured_output(model) is False + + +def test_capability_predicate_reads_no_instance_state() -> None: + """The capability predicate is a pure function of its argument. + + The subclass walk that checks connection capabilities calls this on an instance + built with `__new__`, where reading any field raises AttributeError. + """ + uninitialized = AzureOpenAIChatModelConnection.__new__( + AzureOpenAIChatModelConnection + ) + assert uninitialized.supports_native_structured_output("gpt-5") is True diff --git a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py index e31d37535..9b814c3cf 100644 --- a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py +++ b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py @@ -87,6 +87,21 @@ def _connections_implementing_chat() -> List[Type[BaseChatModelConnection]]: ] +def _translates_schema_natively(cls: Type[BaseChatModelConnection]) -> bool: + """Whether ``cls`` applies an output schema through a native provider parameter. + + Overriding ``supports_native_structured_output`` is how a connection reports that + capability, so the same override marks the connections that must accept a schema + instead of rejecting it. Such a connection owns the decision of what to do with a + schema it cannot apply natively for the effective model, and its own tests pin + that behavior. + """ + return ( + cls.supports_native_structured_output + is not BaseChatModelConnection.supports_native_structured_output + ) + + def test_every_connection_declares_output_schema_param() -> None: """Every connection in the tree must declare ``output_schema`` defaulting to None. @@ -147,13 +162,20 @@ def test_every_connection_rejects_an_output_schema_it_cannot_translate() -> None own it lets a connection silently return an unconstrained response that the caller would treat as schema-conforming. Rejecting turns that into an error at the call. + A connection that does translate a schema natively is exempt, since rejecting is + exactly what it must not do. + The tree is walked rather than hand-listed, so a connection added to a module that is already imported is held to the contract for free. Reach still stops at those imports: ``__subclasses__()`` only sees classes that have been imported, so a connection in a brand-new module needs that module added at the top of this file. """ schema = OutputSchema(output_schema=_Answer) - connections = _connections_implementing_chat() + connections = [ + cls + for cls in _connections_implementing_chat() + if not _translates_schema_natively(cls) + ] assert connections, "the connection walk found nothing to check" offenders = [ From 53bb169807ac9d287ad232cc3c282f9396504932 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Mon, 27 Jul 2026 10:42:18 -0700 Subject: [PATCH 2/2] [integrations][java] Apply Azure OpenAI native structured output Translate a POJO output schema into openai-java's native responseFormat instead of inheriting the rejecting default, mirroring the Python Azure OpenAI connection: the capable-model allowlist and the api-version floor are identical in both languages. Capability is keyed on model_of_azure_deployment rather than model, because the latter names the Azure deployment and says nothing about the model behind it. The native path also requires an api_version at or above 2024-08-01, so a schema is never sent to a service version that predates structured output. Building the request inline left no seam to observe, so chat() is split into a 4-arg override, a private doChat, and package-private buildRequest and toResponse. The split turns one read of model_of_azure_deployment into two that must agree; toResponse makes that pinnable, and reading the key after the call is equivalent to reading it before because the parameter map is built per call and never retained. A caller-supplied response_format that the native path would overwrite now raises rather than being silently replaced. --- .../AzureOpenAIChatModelConnection.java | 309 ++++++++++--- .../openai/AzureOpenAIChatModelSetup.java | 4 +- .../AzureOpenAIChatModelConnectionTest.java | 413 +++++++++++++++++- 3 files changed, 660 insertions(+), 66 deletions(-) diff --git a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java index 8626b1b80..0354c8fec 100644 --- a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java +++ b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java @@ -25,10 +25,12 @@ import com.openai.azure.credential.AzureApiKeyCredential; import com.openai.client.OpenAIClient; import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.core.JsonSchemaLocalValidation; import com.openai.core.JsonValue; import com.openai.models.ChatModel; import com.openai.models.FunctionDefinition; import com.openai.models.FunctionParameters; +import com.openai.models.ResponseFormatJsonSchema; import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionCreateParams; import com.openai.models.chat.completions.ChatCompletionFunctionTool; @@ -98,8 +100,48 @@ public class AzureOpenAIChatModelConnection extends BaseChatModelConnection { private static final Set RESERVED_KWARG_KEYS = Set.of("model", "model_of_azure_deployment", "temperature", "max_tokens", "logprobs"); + // Models with documented json_schema strict Structured Outputs support. Source of truth: + // https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/structured-outputs + // + // Matching is exact, never by prefix: Azure exposes a deployment's model name and model version + // as separate properties, so a name carries no version to discriminate on. The documented list + // includes gpt-4o only at versions 2024-08-06 and 2024-11-20 while version 2024-05-13 is + // unsupported, so a bare "gpt-4o" is ambiguous and is deliberately absent from the set below. + // An unrecognized name reports not-capable and degrades to the prompt fallback rather than + // failing at the provider. + // + // The source list prints "gpt-5.1-codex mini" with a space; it is transcribed hyphenated here + // because Azure model identifiers do not contain spaces. + private static final Set NATIVE_STRUCTURED_OUTPUT_MODELS = + Set.of( + "gpt-5.1-codex", + "gpt-5.1-codex-mini", + "gpt-5.1", + "gpt-5.1-chat", + "gpt-5-pro", + "gpt-5-codex", + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "codex-mini", + "o3-pro", + "o3-mini", + "o1", + "gpt-4o-mini", + "gpt-4.1", + "gpt-4.1-nano", + "gpt-4.1-mini", + "o4-mini", + "o3"); + + // Date prefix of 2024-08-01-preview, the earliest api-version Azure documents as supporting + // structured outputs. + private static final String MIN_STRUCTURED_OUTPUT_API_VERSION = "2024-08-01"; + private final OpenAIClient client; + private final String apiVersion; + public AzureOpenAIChatModelConnection( ResourceDescriptor descriptor, ResourceContext resourceContext) { super(descriptor, resourceContext); @@ -113,6 +155,7 @@ public AzureOpenAIChatModelConnection( if (apiVersion == null || apiVersion.isBlank()) { throw new IllegalArgumentException("api_version should not be null or empty."); } + this.apiVersion = apiVersion; String azureEndpoint = descriptor.getArgument("azure_endpoint"); if (azureEndpoint == null || azureEndpoint.isBlank()) { @@ -151,84 +194,226 @@ public AzureOpenAIChatModelConnection( this.client = clientBuilder.build(); } + /** + * Whether Azure documents json_schema strict support for {@code effectiveModel}. + * + *

{@code effectiveModel} is the model backing an Azure deployment, not the deployment name. + * See the allowlist above for the source of truth and for why the match is exact. An + * unrecognized model reports {@code false} so it degrades to the prompt-engineering fallback + * rather than failing at the provider. + * + *

Reads no instance state, so capability stays answerable independently of how the + * connection was configured. + */ + @Override + protected boolean supportsNativeStructuredOutput(String effectiveModel) { + if (effectiveModel == null || effectiveModel.isEmpty()) { + return false; + } + return NATIVE_STRUCTURED_OUTPUT_MODELS.contains(effectiveModel); + } + + /** + * Whether the configured api-version reaches the structured-output floor. + * + *

Azure documents {@code 2024-08-01-preview} as the first api-version supporting structured + * outputs, and whether an older version rejects {@code response_format} or silently ignores it + * is not documented. The request therefore never carries {@code response_format} below the + * floor, which is safe under either behavior. + * + *

The comparison assumes the documented api-version form, a zero-padded {@code YYYY-MM-DD} + * date optionally suffixed {@code -preview}; over that form comparing the leading date + * lexicographically is exact. The GA {@code v1} literal sorts above the floor, which matches + * Azure documenting {@code v1} as supporting structured outputs. This is not a validator: a + * value of any other shape is not classified reliably, and the service rejects an api-version + * it does not recognize. The constructor rejects a null or blank api-version, so no value of + * that shape reaches here. + */ + private boolean apiVersionSupportsStructuredOutput() { + String datePrefix = + apiVersion.length() > MIN_STRUCTURED_OUTPUT_API_VERSION.length() + ? apiVersion.substring(0, MIN_STRUCTURED_OUTPUT_API_VERSION.length()) + : apiVersion; + return datePrefix.compareTo(MIN_STRUCTURED_OUTPUT_API_VERSION) >= 0; + } + @Override public ChatMessage chat( List messages, List tools, Map modelParams) { + return doChat(messages, tools, modelParams, null); + } + + /** + * Translates {@code outputSchema} into Azure's native strict {@code response_format} + * json_schema when it is a POJO {@link Class}, the model backing the deployment is one Azure + * documents json_schema strict support for, and the configured api-version reaches {@code + * 2024-08-01-preview}. Any other combination leaves the request unconstrained so that the + * prompt-engineering fallback still governs the response, rather than failing at the provider. + * + *

Capability is keyed on the {@code model_of_azure_deployment} model parameter rather than + * on the deployment the request targets, because a deployment name is chosen by the user and + * carries no model information. Leaving that parameter unset therefore keeps even a capable + * deployment on the fallback. + * + * @throws IllegalArgumentException if the schema is applied natively while {@code + * additional_kwargs} also carries a {@code response_format}, since the two would otherwise + * compete on the same request + */ + @Override + public ChatMessage chat( + List messages, + List tools, + Map modelParams, + Object outputSchema) { + return doChat(messages, tools, modelParams, outputSchema); + } + + private ChatMessage doChat( + List messages, + List tools, + Map modelParams, + Object outputSchema) { try { - Map mutableArgs = - modelParams != null ? new HashMap<>(modelParams) : new HashMap<>(); + ChatCompletionCreateParams params = + buildRequest(messages, tools, modelParams, outputSchema); + return toResponse(client.chat().completions().create(params), modelParams); + } catch (IllegalArgumentException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException("Failed to call Azure OpenAI chat completions API.", e); + } + } - String azureDeployment = (String) mutableArgs.remove("model"); - if (azureDeployment == null || azureDeployment.isBlank()) { - throw new IllegalArgumentException("model is required for Azure OpenAI API calls"); - } - String modelOfAzureDeployment = - (String) mutableArgs.remove("model_of_azure_deployment"); + // Package-private so response handling can be asserted against a constructed completion without + // issuing a live API call through the final OpenAI client. + ChatMessage toResponse(ChatCompletion completion, Map modelParams) { + // Read from the caller's map rather than the copy buildRequest consumed, and read without + // consuming: a caller may reuse the same map across calls. The map is assembled fresh for + // each call and no one retains it, so reading it once the response has arrived yields the + // same value as reading it before the request was issued. Token metrics report the model + // backing the deployment, which buildRequest only uses to decide capability. + String modelOfAzureDeployment = + modelParams != null ? (String) modelParams.get("model_of_azure_deployment") : null; + + ChatMessage response = + OpenAIChatCompletionsUtils.convertFromOpenAIMessage( + completion.choices().get(0).message()); + + if (modelOfAzureDeployment != null + && !modelOfAzureDeployment.isBlank() + && completion.usage().isPresent()) { + response.getExtraArgs().put("model_name", modelOfAzureDeployment); + response.getExtraArgs().put("promptTokens", completion.usage().get().promptTokens()); + response.getExtraArgs() + .put("completionTokens", completion.usage().get().completionTokens()); + } - ChatCompletionCreateParams.Builder builder = - ChatCompletionCreateParams.builder() - .model(ChatModel.of(azureDeployment)) - .messages(OpenAIChatCompletionsUtils.convertToOpenAIMessages(messages)); + return response; + } - if (tools != null && !tools.isEmpty()) { - builder.tools(convertTools(tools)); - } + // Package-private so the request body (including the native response_format) can be asserted + // without issuing a live API call through the final OpenAI client. + ChatCompletionCreateParams buildRequest( + List messages, + List tools, + Map rawModelParams, + Object outputSchema) { + Map mutableArgs = + rawModelParams != null ? new HashMap<>(rawModelParams) : new HashMap<>(); + + String azureDeployment = (String) mutableArgs.remove("model"); + if (azureDeployment == null || azureDeployment.isBlank()) { + throw new IllegalArgumentException("model is required for Azure OpenAI API calls"); + } + String modelOfAzureDeployment = (String) mutableArgs.remove("model_of_azure_deployment"); - Object temperature = mutableArgs.remove("temperature"); - if (temperature instanceof Number) { - builder.temperature(((Number) temperature).doubleValue()); - } + ChatCompletionCreateParams.Builder builder = + ChatCompletionCreateParams.builder() + .model(ChatModel.of(azureDeployment)) + .messages(OpenAIChatCompletionsUtils.convertToOpenAIMessages(messages)); - Object maxTokens = mutableArgs.remove("max_tokens"); - if (maxTokens instanceof Number) { - builder.maxCompletionTokens(((Number) maxTokens).longValue()); - } + if (tools != null && !tools.isEmpty()) { + builder.tools(convertTools(tools)); + } - Object logprobs = mutableArgs.remove("logprobs"); - if (Boolean.TRUE.equals(logprobs)) { - builder.logprobs(true); - } + // Capability belongs to the model backing the deployment, so it is the input to the check; + // the deployment name is chosen by the user and carries none. Native structured output + // applies only for a POJO Class schema — a RowTypeInfo (wrapped in OutputSchema) keeps the + // prompt-engineering fallback, as do an incapable model and an api-version below the floor. + String nativeSchemaName = null; + if (outputSchema instanceof Class + && supportsNativeStructuredOutput(modelOfAzureDeployment) + && apiVersionSupportsStructuredOutput()) { + Class schemaClass = (Class) outputSchema; + builder.responseFormat(toNativeResponseFormat(schemaClass)); + nativeSchemaName = schemaClass.getSimpleName(); + } - @SuppressWarnings("unchecked") - Map additionalKwargs = - (Map) mutableArgs.remove("additional_kwargs"); - if (additionalKwargs != null) { - Set collisions = new HashSet<>(additionalKwargs.keySet()); - collisions.retainAll(RESERVED_KWARG_KEYS); - if (!collisions.isEmpty()) { - throw new IllegalArgumentException( - "additional_kwargs must not contain reserved typed fields: " - + collisions - + ". Set these via the corresponding Setup field instead."); - } - for (Map.Entry entry : additionalKwargs.entrySet()) { - builder.putAdditionalBodyProperty( - entry.getKey(), toJsonValue(entry.getValue())); - } - } + Object temperature = mutableArgs.remove("temperature"); + if (temperature instanceof Number) { + builder.temperature(((Number) temperature).doubleValue()); + } - ChatCompletion completion = client.chat().completions().create(builder.build()); + Object maxTokens = mutableArgs.remove("max_tokens"); + if (maxTokens instanceof Number) { + builder.maxCompletionTokens(((Number) maxTokens).longValue()); + } - ChatMessage response = - OpenAIChatCompletionsUtils.convertFromOpenAIMessage( - completion.choices().get(0).message()); + Object logprobs = mutableArgs.remove("logprobs"); + if (Boolean.TRUE.equals(logprobs)) { + builder.logprobs(true); + } - if (modelOfAzureDeployment != null - && !modelOfAzureDeployment.isBlank() - && completion.usage().isPresent()) { - response.getExtraArgs().put("model_name", modelOfAzureDeployment); - response.getExtraArgs() - .put("promptTokens", completion.usage().get().promptTokens()); - response.getExtraArgs() - .put("completionTokens", completion.usage().get().completionTokens()); + @SuppressWarnings("unchecked") + Map additionalKwargs = + (Map) mutableArgs.remove("additional_kwargs"); + if (additionalKwargs != null) { + Set collisions = new HashSet<>(additionalKwargs.keySet()); + collisions.retainAll(RESERVED_KWARG_KEYS); + if (!collisions.isEmpty()) { + throw new IllegalArgumentException( + "additional_kwargs must not contain reserved typed fields: " + + collisions + + ". Set these via the corresponding Setup field instead."); + } + // Only the branch that actually sent a schema may reject the caller's own + // response_format; every path that skipped it leaves the value untouched. + if (nativeSchemaName != null && additionalKwargs.containsKey("response_format")) { + throw new IllegalArgumentException( + "The " + + nativeSchemaName + + " output schema is sent as response_format on deployment '" + + azureDeployment + + "', so response_format must not also be set in additional_kwargs." + + " Remove that value, or omit the output schema to set" + + " response_format directly."); + } + for (Map.Entry entry : additionalKwargs.entrySet()) { + builder.putAdditionalBodyProperty(entry.getKey(), toJsonValue(entry.getValue())); } - - return response; - } catch (IllegalArgumentException e) { - throw e; - } catch (Exception e) { - throw new RuntimeException("Failed to call Azure OpenAI chat completions API.", e); } + + return builder.build(); + } + + // Derives the strict json_schema response format from a POJO class via the SDK's typed + // structured-output builder. The Kotlin-facade StructuredOutputsKt.responseFormatFromClass is + // not callable from Java, so the response format is extracted through the typed builder, which + // generates the same strict draft-2020-12 schema, and then reattached to the standard builder. + private static ResponseFormatJsonSchema toNativeResponseFormat(Class schemaClass) { + return ChatCompletionCreateParams.builder() + .model(ChatModel.of("")) + .addUserMessage("") + .responseFormat(schemaClass, JsonSchemaLocalValidation.NO) + .build() + .rawParams() + .responseFormat() + .orElseThrow( + () -> + new IllegalStateException( + "OpenAI SDK did not produce a response_format for schema " + + schemaClass.getName())) + .asJsonSchema(); } @Override diff --git a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelSetup.java b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelSetup.java index 44a7c8431..d88d01cc1 100644 --- a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelSetup.java +++ b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelSetup.java @@ -32,7 +32,9 @@ * *

{@code model} (inherited from {@link BaseChatModelSetup}) is the Azure deployment name, not * the underlying OpenAI model name. The underlying model name can be supplied via {@code - * model_of_azure_deployment} and is used solely for token-metrics tracking. + * model_of_azure_deployment}. It labels token-usage metrics, and it is also the name that decides + * whether a request can carry a native structured-output schema, since the deployment name carries + * no model information. Leaving it unset keeps requests on the prompt-engineering fallback. * *

Example usage: * diff --git a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java index 60a29729b..04581e6fd 100644 --- a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java +++ b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java @@ -18,34 +18,107 @@ package org.apache.flink.agents.integrations.chatmodels.openai; +import com.fasterxml.jackson.core.type.TypeReference; +import com.openai.models.ChatModel; +import com.openai.models.ResponseFormatJsonSchema; +import com.openai.models.chat.completions.ChatCompletion; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +import com.openai.models.chat.completions.ChatCompletionMessage; +import com.openai.models.completions.CompletionUsage; import org.apache.flink.agents.api.chat.messages.ChatMessage; import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.agents.api.tools.ToolMetadata; +import org.apache.flink.agents.api.tools.ToolParameters; +import org.apache.flink.agents.api.tools.ToolResponse; +import org.apache.flink.agents.api.tools.ToolType; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.NullAndEmptySource; +import org.junit.jupiter.params.provider.ValueSource; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** - * Unit tests for {@link AzureOpenAIChatModelConnection} — constructor validation only, no network - * access. End-to-end tests against a real Azure OpenAI deployment live in {@link - * AzureOpenAIChatModelIT}. + * Unit tests for {@link AzureOpenAIChatModelConnection} — constructor validation and request + * building, with no network access. End-to-end tests against a real Azure OpenAI deployment live in + * {@link AzureOpenAIChatModelIT}. */ class AzureOpenAIChatModelConnectionTest { private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + /** A deployment name is chosen by the user and carries no capability information. */ + private static final String DEPLOYMENT = "my-deployment"; + + private static final String CAPABLE_API_VERSION = "2024-08-01-preview"; + + private static final String BELOW_FLOOR_API_VERSION = "2024-02-01"; + + private static final Map CALLER_RESPONSE_FORMAT = Map.of("type", "json_object"); + private static ResourceDescriptor.Builder connectionDescriptor() { return ResourceDescriptor.Builder.newBuilder( AzureOpenAIChatModelConnection.class.getName()); } + /** A representative POJO output schema. */ + public static class Person { + public String name; + public int age; + } + + private static AzureOpenAIChatModelConnection connection(String apiVersion) { + ResourceDescriptor desc = + connectionDescriptor() + .addInitialArgument("api_key", "test-key") + .addInitialArgument("api_version", apiVersion) + .addInitialArgument("azure_endpoint", "https://example.openai.azure.com") + .build(); + return new AzureOpenAIChatModelConnection(desc, NOOP); + } + + private static AzureOpenAIChatModelConnection connection() { + return connection(CAPABLE_API_VERSION); + } + + /** + * Model params addressing {@link #DEPLOYMENT}. A null {@code modelOfAzureDeployment} omits the + * key entirely, which is how the setup emits an unset backing model. + */ + private static Map params(String modelOfAzureDeployment) { + Map params = new HashMap<>(); + params.put("model", DEPLOYMENT); + if (modelOfAzureDeployment != null) { + params.put("model_of_azure_deployment", modelOfAzureDeployment); + } + return params; + } + + private static Map paramsWithCallerResponseFormat( + String modelOfAzureDeployment) { + Map params = params(modelOfAzureDeployment); + params.put("additional_kwargs", Map.of("response_format", CALLER_RESPONSE_FORMAT)); + return params; + } + + private static List userMessage() { + return List.of(new ChatMessage(MessageRole.USER, "hi")); + } + @Test @DisplayName("Constructor throws when api_key is missing") void testConstructorMissingApiKey() { @@ -128,4 +201,338 @@ void testChatRejectsReservedKeyInAdditionalKwargs() { .hasMessageContaining("additional_kwargs") .hasMessageContaining("temperature"); } + + @Test + @DisplayName("Native response_format json_schema strict applied for a POJO on a capable model") + void testNativeAppliedForCapableDeploymentModel() { + ChatCompletionCreateParams request = + connection() + .buildRequest( + userMessage(), List.of(), params("gpt-4o-mini"), Person.class); + + assertThat(request.responseFormat()).isPresent(); + ResponseFormatJsonSchema jsonSchema = request.responseFormat().get().asJsonSchema(); + // The SDK derives the wire name from the class, so it identifies the schema without being + // equal to the class name. + assertThat(jsonSchema.jsonSchema().name()).contains("Person"); + assertThat(jsonSchema.jsonSchema().strict()).contains(true); + } + + @Test + @DisplayName("A native request still targets the deployment, not the backing model") + void testCapableNativeRequestStillTargetsTheDeployment() { + // Capability is keyed on the backing model, but the provider is still addressed by + // deployment; substituting one for the other would route the call to a deployment that may + // not exist on the resource. + ChatCompletionCreateParams request = + connection() + .buildRequest( + userMessage(), List.of(), params("gpt-4o-mini"), Person.class); + + assertThat(request.model()).isEqualTo(ChatModel.of(DEPLOYMENT)); + } + + @Test + @DisplayName("Native NOT applied when the backing model of the deployment is absent") + void testNativeNotAppliedWhenDeploymentModelAbsent() { + ChatCompletionCreateParams request = + connection().buildRequest(userMessage(), List.of(), params(null), Person.class); + + assertThat(request.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied for a backing model outside the allowlist") + void testNativeNotAppliedForUnknownDeploymentModel() { + ChatCompletionCreateParams request = + connection() + .buildRequest( + userMessage(), + List.of(), + params("some-unknown-model"), + Person.class); + + assertThat(request.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied for a bare gpt-4o backing model") + void testNativeNotAppliedForBareGpt4o() { + // Azure carries model name and model version as separate properties, so a bare gpt-4o may + // be the 2024-05-13 version, which predates structured-output support. + ChatCompletionCreateParams request = + connection().buildRequest(userMessage(), List.of(), params("gpt-4o"), Person.class); + + assertThat(request.responseFormat()).isEmpty(); + } + + @ParameterizedTest + @ValueSource(strings = {"2024-10-21", "v1"}) + @DisplayName("Native applied for a later GA date and for the GA v1 literal") + void testNativeAppliedForOtherApiVersionsAboveFloor(String apiVersion) { + // v1 pins the non-date half of the comparison: it is admitted because it sorts above the + // floor, not because it parses as a later date. Rewriting the gate to parse dates would + // drop it while every date case still passed. + ChatCompletionCreateParams request = + connection(apiVersion) + .buildRequest( + userMessage(), List.of(), params("gpt-4o-mini"), Person.class); + + assertThat(request.responseFormat()).isPresent(); + } + + @Test + @DisplayName("Native NOT applied when the configured api-version predates the floor") + void testNativeNotAppliedWhenApiVersionBelowFloor() { + // An absent api-version cannot reach this gate at all: the constructor rejects it, which + // testConstructorMissingApiVersion pins. + ChatCompletionCreateParams request = + connection(BELOW_FLOOR_API_VERSION) + .buildRequest( + userMessage(), List.of(), params("gpt-4o-mini"), Person.class); + + assertThat(request.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied when no output schema is supplied") + void testNativeNotAppliedWhenSchemaNull() { + ChatCompletionCreateParams request = + connection().buildRequest(userMessage(), List.of(), params("gpt-4o-mini"), null); + + assertThat(request.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native NOT applied for a non-POJO schema form (POJO-only scope)") + void testNativeNotAppliedForNonPojoSchema() { + // A RowTypeInfo schema arrives wrapped in OutputSchema rather than as a bare POJO Class, so + // it must not activate native structured output. OutputSchema cannot be instantiated here + // because RowTypeInfo is not on this module's classpath; any non-Class schema object + // exercises the same gate. + Object nonClassSchema = "row"; + + ChatCompletionCreateParams request = + connection() + .buildRequest( + userMessage(), List.of(), params("gpt-4o-mini"), nonClassSchema); + + assertThat(request.responseFormat()).isEmpty(); + } + + @Test + @DisplayName("Native applied for a POJO even when tools are bound") + void testNativeAppliedEvenWhenToolsBound() { + // Azure documents structured outputs as unsupported with parallel function calls, which + // constrains strict tool schemas rather than the response_format this branch sets, so + // binding tools does not gate it. + ChatCompletionCreateParams request = + connection() + .buildRequest( + userMessage(), + List.of(new StubTool()), + params("gpt-4o-mini"), + Person.class); + + assertThat(request.responseFormat()).isPresent(); + } + + @ParameterizedTest + @ValueSource( + strings = { + "gpt-5.1-codex", + "gpt-5.1-codex-mini", + "gpt-5.1", + "gpt-5.1-chat", + "gpt-5-pro", + "gpt-5-codex", + "gpt-5", + "gpt-5-mini", + "gpt-5-nano", + "codex-mini", + "o3-pro", + "o3-mini", + "o1", + "gpt-4o-mini", + "gpt-4.1", + "gpt-4.1-nano", + "gpt-4.1-mini", + "o4-mini", + "o3" + }) + @DisplayName("Capability predicate accepts every documented capable Azure model name") + void testCapabilityPredicateAcceptsCapableModels(String model) { + // The list is the whole allowlist, so dropping an entry is caught rather than only + // narrowing capability silently. + assertThat(connection().supportsNativeStructuredOutput(model)).isTrue(); + } + + @ParameterizedTest + @NullAndEmptySource + @ValueSource( + strings = { + "gpt-4o", + "gpt-35-turbo", + "gpt-4", + "gpt-4o-2024-08-06", + "some-unknown-model" + }) + @DisplayName("Capability predicate rejects incapable, version-suffixed, and empty names") + void testCapabilityPredicateRejectsIncapableModels(String model) { + // A version-suffixed value such as gpt-4o-2024-08-06 is an OpenAI snapshot name, not a name + // Azure reports as the model behind a deployment. + assertThat(connection().supportsNativeStructuredOutput(model)).isFalse(); + } + + @Test + @DisplayName( + "A caller-supplied response_format alongside a natively applied schema is rejected") + void testCallerResponseFormatConflictsWithNativeSchema() { + // Both values would otherwise reach the same request, where the additional body property + // silently competes with the typed response_format the schema produced. + AzureOpenAIChatModelConnection conn = connection(); + Map args = paramsWithCallerResponseFormat("gpt-4o-mini"); + + assertThatThrownBy(() -> conn.chat(userMessage(), null, args, Person.class)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("response_format") + .hasMessageContaining("Person"); + } + + private static Stream nonNativePaths() { + return Stream.of( + Arguments.of("incapable_model", CAPABLE_API_VERSION, "gpt-4o", Person.class), + Arguments.of( + "non_pojo_schema", CAPABLE_API_VERSION, "gpt-4o-mini", "row"), + Arguments.of("no_output_schema", CAPABLE_API_VERSION, "gpt-4o-mini", null), + Arguments.of( + "api_version_below_floor", + BELOW_FLOOR_API_VERSION, + "gpt-4o-mini", + Person.class)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("nonNativePaths") + @DisplayName("A caller-supplied response_format survives wherever native output is skipped") + void testCallerResponseFormatSurvivesWhenNativeIsSkipped( + String label, String apiVersion, String modelOfAzureDeployment, Object outputSchema) { + // Only the branch that actually sends a schema as response_format may reject the caller's + // own value, so identical caller code has to keep working along every path that skips it, + // including the no-schema path taken by any caller that drives response_format itself. + ChatCompletionCreateParams request = + connection(apiVersion) + .buildRequest( + userMessage(), + List.of(), + paramsWithCallerResponseFormat(modelOfAzureDeployment), + outputSchema); + + assertThat(request.responseFormat()).isEmpty(); + assertThat(request._additionalBodyProperties()) + .hasEntrySatisfying( + "response_format", + value -> + assertThat( + value.convert( + new TypeReference< + Map>() {})) + .isEqualTo(CALLER_RESPONSE_FORMAT)); + } + + @Test + @DisplayName("Token metrics label the response with the model backing the deployment") + void testResponseCarriesBackingModelTokenMetrics() { + // The deployment name is what the request targets, but usage has to be attributed to the + // model behind it, which is the only name that identifies what actually ran. + ChatMessage response = + connection().toResponse(completionWithUsage(11L, 7L), params("gpt-4o-mini")); + + assertThat(response.getExtraArgs()) + .containsEntry("model_name", "gpt-4o-mini") + .containsEntry("promptTokens", 11L) + .containsEntry("completionTokens", 7L); + } + + @Test + @DisplayName("Response handling leaves the caller's model params intact") + void testResponseHandlingDoesNotConsumeCallerModelParams() { + // The backing model is read from the map the caller owns. Consuming the entry would strip + // token metrics from every later call that reuses the same map. + Map callerParams = params("gpt-4o-mini"); + + connection().toResponse(completionWithUsage(11L, 7L), callerParams); + + assertThat(callerParams).containsEntry("model_of_azure_deployment", "gpt-4o-mini"); + } + + private static Stream incompleteMetricsInputs() { + return Stream.of( + Arguments.of("backing_model_unset", completionWithUsage(11L, 7L), params(null)), + Arguments.of( + "completion_without_usage", + completionWithoutUsage(), + params("gpt-4o-mini"))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("incompleteMetricsInputs") + @DisplayName("Token metrics are omitted when the backing model or the usage report is absent") + void testNoTokenMetricsWhenMetricsInputsAreIncomplete( + String label, ChatCompletion completion, Map modelParams) { + // Leaving the backing model unset is the documented default, and a completion may arrive + // without a usage report; both drop the metrics rather than costing the caller the reply. + ChatMessage response = connection().toResponse(completion, modelParams); + + assertThat(response.getExtraArgs()) + .doesNotContainKeys("model_name", "promptTokens", "completionTokens"); + } + + private static ChatCompletion completionWithUsage(long promptTokens, long completionTokens) { + return completionBuilder() + .usage( + CompletionUsage.builder() + .promptTokens(promptTokens) + .completionTokens(completionTokens) + .totalTokens(promptTokens + completionTokens) + .build()) + .build(); + } + + private static ChatCompletion completionWithoutUsage() { + return completionBuilder().build(); + } + + private static ChatCompletion.Builder completionBuilder() { + ChatCompletionMessage message = + ChatCompletionMessage.builder().content("hi").refusal(Optional.empty()).build(); + return ChatCompletion.builder() + .id("completion-1") + .created(0L) + .model(DEPLOYMENT) + .addChoice( + ChatCompletion.Choice.builder() + .finishReason(ChatCompletion.Choice.FinishReason.STOP) + .index(0L) + .logprobs(Optional.empty()) + .message(message) + .build()); + } + + /** Minimal tool stub; only its presence in the tools list matters. */ + private static class StubTool extends Tool { + StubTool() { + super(new ToolMetadata("add", "adds", "{\"type\":\"object\"}")); + } + + @Override + public ToolType getToolType() { + return ToolType.FUNCTION; + } + + @Override + public ToolResponse call(ToolParameters parameters) { + return ToolResponse.success(null); + } + } }