This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 5cc3b585 [integrations][java][python] Apply Azure OpenAI native
structured output (#930)
5cc3b585 is described below
commit 5cc3b585b2ae2f7163323b970ca75f13b95bff0b
Author: Weiqing Yang <[email protected]>
AuthorDate: Tue Aug 4 02:14:38 2026 -0700
[integrations][java][python] Apply Azure OpenAI native structured output
(#930)
---
.../openai/AzureOpenAIChatModelConnection.java | 325 +++++++++++++---
.../openai/AzureOpenAIChatModelSetup.java | 4 +-
.../openai/AzureOpenAIChatModelConnectionTest.java | 431 +++++++++++++++++++-
.../chat_models/azure/azure_openai_chat_model.py | 161 +++++++-
.../test_azure_openai_native_structured_output.py | 432 +++++++++++++++++++++
5 files changed, 1281 insertions(+), 72 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 8626b1b8..31b2c5ab 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.AzureUrlPathMode;
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;
@@ -47,6 +49,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.regex.Pattern;
/**
* Chat model integration for Azure OpenAI Service. Built on the openai-java
SDK using its built-in
@@ -98,8 +101,49 @@ public class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
private static final Set<String> RESERVED_KWARG_KEYS =
Set.of("model", "model_of_azure_deployment", "temperature",
"max_tokens", "logprobs");
+ // Models that both have documented json_schema strict Structured Outputs
support and are served
+ // on the Chat Completions API, which is the API this connection calls.
The set is that
+ // intersection, taken from two sources:
+ //
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/structured-outputs
lists the
+ // models supporting Structured Outputs on any API, and
+ //
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
carries the
+ // per-model feature table whose "Chat Completions API" row excludes the
models Azure serves
+ // only on the Responses API.
+ //
+ // 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.
+ private static final Set<String> NATIVE_STRUCTURED_OUTPUT_MODELS =
+ Set.of(
+ "gpt-5.1",
+ "gpt-5.1-chat",
+ "gpt-5",
+ "gpt-5-mini",
+ "gpt-5-nano",
+ "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";
+
+ // Leading zero-padded YYYY-MM-DD date of the api-version form Azure
documents, which is a date
+ // optionally carrying a suffix such as -preview.
+ private static final Pattern API_VERSION_DATE_PREFIX =
Pattern.compile("^\\d{4}-\\d{2}-\\d{2}");
+
private final OpenAIClient client;
+ private final String apiVersion;
+
public AzureOpenAIChatModelConnection(
ResourceDescriptor descriptor, ResourceContext resourceContext) {
super(descriptor, resourceContext);
@@ -113,6 +157,7 @@ public class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
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 +196,240 @@ public class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
this.client = clientBuilder.build();
}
+ /**
+ * Whether Azure documents json_schema strict support for {@code
effectiveModel}.
+ *
+ * <p>{@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.
+ *
+ * <p>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.
+ *
+ * <p>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.
+ *
+ * <p>Only the documented api-version form is classified, a zero-padded
{@code YYYY-MM-DD} date
+ * optionally suffixed {@code -preview}; over that form comparing the
leading date
+ * lexicographically is exact. A value of any other shape, including the
GA {@code v1} literal,
+ * reports {@code false} and keeps the prompt fallback. That is the right
answer for {@code v1}
+ * under the default {@code AUTO} path mode against a resource endpoint,
where the request is
+ * built on the deployment-scoped path {@code
/openai/deployments/{deployment}/chat/completions}
+ * with the api-version carried as a query parameter, so the literal is
sent as {@code
+ * ?api-version=v1} rather than selecting Azure's {@code /openai/v1}
endpoint. Under {@code
+ * UNIFIED}, or an endpoint already ending in {@code /openai/v1}, the
request does reach the
+ * unified endpoint, and reporting not-capable there costs only the prompt
fallback. The
+ * constructor rejects a null or blank api-version, so no value of that
shape reaches here.
+ */
+ private boolean apiVersionSupportsStructuredOutput() {
+ if (!API_VERSION_DATE_PREFIX.matcher(apiVersion).lookingAt()) {
+ return false;
+ }
+ return apiVersion
+ .substring(0,
MIN_STRUCTURED_OUTPUT_API_VERSION.length())
+ .compareTo(MIN_STRUCTURED_OUTPUT_API_VERSION)
+ >= 0;
+ }
+
@Override
public ChatMessage chat(
List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
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.
+ *
+ * <p>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<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
+ return doChat(messages, tools, modelParams, outputSchema);
+ }
+
+ private ChatMessage doChat(
+ List<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> modelParams,
+ Object outputSchema) {
try {
- Map<String, Object> 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<String, Object>
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<ChatMessage> messages,
+ List<Tool> tools,
+ Map<String, Object> rawModelParams,
+ Object outputSchema) {
+ Map<String, Object> 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.
+ //
+ // TODO(#912): the requested strategy is not visible here, so this
re-check cannot tell an
+ // explicit NATIVE request apart from one that merely resolved to
native. A caller asking
+ // for NATIVE therefore gets an unconstrained response instead of an
error whenever this
+ // branch is skipped, which on Azure also happens when the api-version
is below the floor
+ // or when model_of_azure_deployment is unset and capability cannot be
resolved at all.
+ // Once strategy resolution is wired up, NATIVE must either bypass
this re-check or fail
+ // explicitly.
+ String nativeSchemaName = null;
+ if (outputSchema instanceof Class
+ && supportsNativeStructuredOutput(modelOfAzureDeployment)
+ && apiVersionSupportsStructuredOutput()) {
+ Class<?> schemaClass = (Class<?>) outputSchema;
+ builder.responseFormat(toNativeResponseFormat(schemaClass));
+ nativeSchemaName = schemaClass.getSimpleName();
+ }
- @SuppressWarnings("unchecked")
- Map<String, Object> additionalKwargs =
- (Map<String, Object>)
mutableArgs.remove("additional_kwargs");
- if (additionalKwargs != null) {
- Set<String> 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<String, Object> 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<String, Object> additionalKwargs =
+ (Map<String, Object>) mutableArgs.remove("additional_kwargs");
+ if (additionalKwargs != null) {
+ Set<String> 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<String, Object> 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 <T> ResponseFormatJsonSchema
toNativeResponseFormat(Class<T> 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 44a7c843..d88d01cc 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 @@ import java.util.Optional;
*
* <p>{@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.
*
* <p>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 60a29729..97b9a2cb 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<String, Object> 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<String, Object> params(String modelOfAzureDeployment) {
+ Map<String, Object> params = new HashMap<>();
+ params.put("model", DEPLOYMENT);
+ if (modelOfAzureDeployment != null) {
+ params.put("model_of_azure_deployment", modelOfAzureDeployment);
+ }
+ return params;
+ }
+
+ private static Map<String, Object> paramsWithCallerResponseFormat(
+ String modelOfAzureDeployment) {
+ Map<String, Object> params = params(modelOfAzureDeployment);
+ params.put("additional_kwargs", Map.of("response_format",
CALLER_RESPONSE_FORMAT));
+ return params;
+ }
+
+ private static List<ChatMessage> userMessage() {
+ return List.of(new ChatMessage(MessageRole.USER, "hi"));
+ }
+
@Test
@DisplayName("Constructor throws when api_key is missing")
void testConstructorMissingApiKey() {
@@ -128,4 +201,356 @@ class AzureOpenAIChatModelConnectionTest {
.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-08-01", "2024-10-21"})
+ @DisplayName("Native applied for a bare GA date at or above the floor")
+ void testNativeAppliedForGaDateAtOrAboveFloor(String apiVersion) {
+ // The documented floor is the preview form 2024-08-01-preview, so
these pin that a bare GA
+ // date carrying no -preview suffix is admitted, and that 2024-08-01
is the inclusive
+ // boundary.
+ ChatCompletionCreateParams request =
+ connection(apiVersion)
+ .buildRequest(
+ userMessage(), List.of(),
params("gpt-4o-mini"), Person.class);
+
+ assertThat(request.responseFormat()).isPresent();
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"v1", "latest"})
+ @DisplayName("Native NOT applied for an api-version outside the documented
dated form")
+ void testNativeNotAppliedForNonDateApiVersion(String apiVersion) {
+ // Every one of these sorts above the floor as a string, so only
classifying the dated form
+ // keeps them out. The v1 literal in particular does not select
Azure's v1 endpoint on the
+ // default path mode: it is sent as a query parameter on the
deployment-scoped
+ // chat/completions path.
+ ChatCompletionCreateParams request =
+ connection(apiVersion)
+ .buildRequest(
+ userMessage(), List.of(),
params("gpt-4o-mini"), Person.class);
+
+ assertThat(request.responseFormat()).isEmpty();
+ }
+
+ @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<name STRING>";
+
+ 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",
+ "gpt-5.1-chat",
+ "gpt-5",
+ "gpt-5-mini",
+ "gpt-5-nano",
+ "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",
+ "gpt-5.1-codex",
+ "gpt-5.1-codex-mini",
+ "gpt-5-pro",
+ "gpt-5-codex",
+ "codex-mini",
+ "o3-pro"
+ })
+ @DisplayName("Capability predicate rejects incapable, Responses-only, 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. The codex,
gpt-5-pro and o3-pro names do
+ // support structured outputs but are served only on the Responses
API, so they are
+ // incapable on the chat completions API this connection calls.
+
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<String, Object> args =
paramsWithCallerResponseFormat("gpt-4o-mini");
+
+ assertThatThrownBy(() -> conn.chat(userMessage(), null, args,
Person.class))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("response_format")
+ .hasMessageContaining("Person");
+ }
+
+ private static Stream<Arguments> 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<name STRING>"),
+ 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<String,
Object>>() {}))
+ .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<String, Object> callerParams = params("gpt-4o-mini");
+
+ connection().toResponse(completionWithUsage(11L, 7L), callerParams);
+
+ assertThat(callerParams).containsEntry("model_of_azure_deployment",
"gpt-4o-mini");
+ }
+
+ private static Stream<Arguments> 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<String, Object>
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);
+ }
+ }
}
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 b653b4d6..422d9d09 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
@@ -16,10 +16,18 @@
# limitations under the License.
#################################################################################
import logging
+import re
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 +48,72 @@ _RESERVED_KWARG_KEYS = frozenset(
{"model", "model_of_azure_deployment", "temperature", "max_tokens",
"logprobs"}
)
+# Models that both have documented json_schema strict Structured Outputs
support and are
+# served on the Chat Completions API, which is the API this connection calls.
The set is
+# that intersection, taken from two sources:
+#
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/structured-outputs
+# lists the models supporting Structured Outputs on any API, and
+# https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
carries the
+# per-model feature table whose "Chat Completions API" row excludes the models
Azure
+# serves only on the Responses API.
+#
+# 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.
+_NATIVE_STRUCTURED_OUTPUT_MODELS = frozenset(
+ {
+ "gpt-5.1",
+ "gpt-5.1-chat",
+ "gpt-5",
+ "gpt-5-mini",
+ "gpt-5-nano",
+ "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"
+
+# Leading zero-padded YYYY-MM-DD date of the api-version form Azure documents,
which
+# is a date optionally carrying a suffix such as -preview. The ASCII flag
restricts \d
+# to 0-9, so a value written with other Unicode decimal digits is not read as
a date.
+_API_VERSION_DATE_PREFIX = re.compile(r"^\d{4}-\d{2}-\d{2}", re.ASCII)
+
+
+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 +188,46 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
)
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.
+
+ Only the documented api-version form is classified, a zero-padded
+ ``YYYY-MM-DD`` date optionally suffixed ``-preview``; over that form
comparing
+ the leading date lexicographically is exact. A value of any other
shape,
+ including the GA ``v1`` literal, reports ``False`` and keeps the prompt
+ fallback. That is the accurate answer for ``v1``: ``AzureOpenAI``
reaches the
+ service through the deployment-scoped path
+ ``/openai/deployments/{deployment}/chat/completions`` with the
api-version
+ carried as a query parameter, so the ``v1`` literal is sent as
+ ``?api-version=v1`` rather than selecting Azure's ``/openai/v1``
endpoint.
+ """
+ if not self.api_version:
+ return False
+ if not _API_VERSION_DATE_PREFIX.match(self.api_version):
+ return False
+ return self.api_version[:10] >= _MIN_STRUCTURED_OUTPUT_API_VERSION
+
def chat(
self,
messages: Sequence[ChatMessage],
@@ -130,10 +244,13 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
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 +260,6 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
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 +281,39 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
)
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.
+ #
+ # TODO(#912): the requested strategy is not visible here, so this
check cannot
+ # tell an explicit NATIVE request apart from one that merely resolved
to native.
+ # A caller asking for NATIVE therefore gets an unconstrained response
instead of
+ # an error whenever this branch is skipped, which on Azure also
happens when the
+ # api-version is below the floor or when model_of_azure_deployment is
unset and
+ # capability cannot be resolved at all. Once strategy resolution is
wired up,
+ # NATIVE must either bypass this check or fail explicitly.
+ 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 00000000..063f3dc2
--- /dev/null
+++
b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_native_structured_output.py
@@ -0,0 +1,432 @@
+################################################################################
+# 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)
+
+
[email protected]("api_version", ["2024-08-01", "2024-10-21"])
+def test_native_applied_for_ga_date_at_or_above_floor(api_version: str) ->
None:
+ """Native applied for a bare GA date at or above the floor.
+
+ The documented floor is the preview form `2024-08-01-preview`, so these
pin that a
+ bare GA date carrying no `-preview` suffix is admitted, and that
`2024-08-01` is the
+ inclusive boundary.
+ """
+ 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)
+
+
[email protected]("api_version", ["v1", "latest"])
+def test_native_not_applied_for_non_date_api_version(api_version: str) -> None:
+ """Native NOT applied for an api-version outside the documented dated form.
+
+ Every one of these sorts above the floor as a string, so only classifying
the
+ dated form keeps them out. The `v1` literal in particular does not reach
Azure's
+ v1 endpoint from here: `AzureOpenAI` sends it as a query parameter on the
+ deployment-scoped chat/completions path.
+ """
+ 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" not 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)
+
+
[email protected]("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)
+
+
[email protected]("in_additional_kwargs", [True, False])
[email protected](
+ ("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
+
+
[email protected](
+ "model",
+ [
+ "gpt-5.1",
+ "gpt-5.1-chat",
+ "gpt-5",
+ "gpt-5-mini",
+ "gpt-5-nano",
+ "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
+
+
[email protected](
+ "model",
+ [
+ "gpt-4o",
+ "gpt-35-turbo",
+ "gpt-4",
+ "gpt-4o-2024-08-06",
+ "some-unknown-model",
+ "gpt-5.1-codex",
+ "gpt-5.1-codex-mini",
+ "gpt-5-pro",
+ "gpt-5-codex",
+ "codex-mini",
+ "o3-pro",
+ None,
+ "",
+ ],
+)
+def test_capability_predicate_rejects_incapable_models(model: str | None) ->
None:
+ """The capability predicate rejects incapable, Responses-only, 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. The codex,
`gpt-5-pro`
+ and `o3-pro` names do support structured outputs but are served only on the
+ Responses API, so they are incapable on the chat completions API this
connection
+ calls.
+ """
+ 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