[integrations][openai] Apply OpenAI native structured output

integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java

@@ -22,11 +22,13 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 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.ReasoningEffort;
+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;
@@ -43,6 +45,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * A chat model integration for the OpenAI Chat Completions service using the official Java SDK.
@@ -119,11 +122,53 @@ public OpenAICompletionsConnection(
         this.client = builder.build();
     }
 
+    // Models for which OpenAI documents json_schema strict Structured Outputs support.
+    // Source of truth: https://platform.openai.com/docs/guides/structured-outputs — json_schema is
+    // supported on the gpt-4o-mini and gpt-4o-2024-08-06 snapshots "and later"; gpt-4-turbo,
+    // earlier models, and gpt-3.5-turbo get JSON mode only.
+    //
+    // "and later" is temporal, not a name prefix: gpt-4o-2024-05-13 predates the cutoff and does
+    // NOT support Structured Outputs, so matching a bare "gpt-4o" prefix would misclassify it as
+    // capable and fail silently at the provider. Prefix matching is therefore used only for the
+    // gpt-4o-mini family, whose entire lifetime post-dates the cutoff; every other capable model is
+    // matched exactly. An unrecognized model reports not-capable and degrades to the prompt
+    // fallback rather than failing at the provider.
+    private static final String NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX = "gpt-4o-mini";
+    private static final Set<String> NATIVE_STRUCTURED_OUTPUT_MODELS =
+            Set.of("gpt-4o", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20");
+
+    @Override
+    protected boolean supportsNativeStructuredOutput(String effectiveModel) {
+        if (effectiveModel == null) {
+            return false;
+        }
+        return effectiveModel.startsWith(NATIVE_STRUCTURED_OUTPUT_FAMILY_PREFIX)
+                || NATIVE_STRUCTURED_OUTPUT_MODELS.contains(effectiveModel);
+    }
+
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {
+        return doChat(messages, tools, modelParams, null);
+    }
+
+    @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 {
-            ChatCompletionCreateParams params = buildRequest(messages, tools, modelParams);
+            ChatCompletionCreateParams params =
+                    buildRequest(messages, tools, modelParams, outputSchema);
             ChatCompletion completion = client.chat().completions().create(params);
             ChatMessage response =
                     OpenAIChatCompletionsUtils.convertFromOpenAIMessage(
@@ -150,8 +195,13 @@ public ChatMessage chat(
         }
     }
 
-    private ChatCompletionCreateParams buildRequest(
-            List<ChatMessage> messages, List<Tool> tools, Map<String, Object> rawModelParams) {
+    // 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> modelParams =
                 rawModelParams != null ? new HashMap<>(rawModelParams) : new HashMap<>();
 
@@ -170,6 +220,13 @@ private ChatCompletionCreateParams buildRequest(
             builder.tools(convertTools(tools, strictMode));
         }
 
+        // Native structured output applies only for a POJO Class schema on a model the provider
+        // documents as capable; a RowTypeInfo (wrapped in OutputSchema) or an incapable model keeps
+        // the prompt-engineering fallback.
+        if (outputSchema instanceof Class && supportsNativeStructuredOutput(modelName)) {
+            builder.responseFormat(toNativeResponseFormat((Class<?>) outputSchema));
+        }
+
         Object temperature = modelParams.remove("temperature");
         if (temperature instanceof Number) {
             builder.temperature(((Number) temperature).doubleValue());
@@ -208,6 +265,26 @@ private ChatCompletionCreateParams buildRequest(
         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();
+    }
+
     private List<ChatCompletionTool> convertTools(List<Tool> tools, boolean strictMode) {
         List<ChatCompletionTool> openaiTools = new ArrayList<>(tools.size());
         for (Tool tool : tools) {

integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java

@@ -0,0 +1,191 @@
+/*
+ * 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.
+ */
+
+package org.apache.flink.agents.integrations.chatmodels.openai;
+
+import com.openai.models.ResponseFormatJsonSchema;
+import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+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 java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests for {@link OpenAICompletionsConnection}'s native structured-output behavior. These
+ * assert the built request body without a live API call by inspecting {@code buildRequest}, and
+ * exercise the model-dependent capability predicate directly.
+ */
+class OpenAICompletionsConnectionTest {
+
+    private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null);
+
+    /** A representative POJO output schema. */
+    public static class Person {
+        public String name;
+        public int age;
+    }
+
+    private static OpenAICompletionsConnection connection() {
+        ResourceDescriptor desc =
+                ResourceDescriptor.Builder.newBuilder(OpenAICompletionsConnection.class.getName())
+                        .addInitialArgument("api_key", "test-key")
+                        .addInitialArgument("model", "gpt-4o")
+                        .build();
+        return new OpenAICompletionsConnection(desc, NOOP);
+    }
+
+    private static Map<String, Object> params(String model) {
+        Map<String, Object> params = new HashMap<>();
+        params.put("model", model);
+        return params;
+    }
+
+    private static List<ChatMessage> userMessage() {
+        return List.of(new ChatMessage(MessageRole.USER, "hi"));
+    }
+
+    @Test
+    @DisplayName("Native response_format json_schema strict applied for a POJO on a capable model")
+    void testNativeAppliedForPojoCapableModel() {
+        ChatCompletionCreateParams params =
+                connection().buildRequest(userMessage(), List.of(), params("gpt-4o"), Person.class);
+
+        assertThat(params.responseFormat()).isPresent();
+        ResponseFormatJsonSchema jsonSchema = params.responseFormat().get().asJsonSchema();
+        assertThat(jsonSchema.jsonSchema().strict()).contains(true);
+    }
+
+    @Test
+    @DisplayName("Native NOT applied for a POJO on an incapable model (prompt fallback)")
+    void testNativeNotAppliedForIncapableModel() {
+        ChatCompletionCreateParams params =
+                connection()
+                        .buildRequest(
+                                userMessage(), List.of(), params("gpt-3.5-turbo"), Person.class);
+
+        assertThat(params.responseFormat()).isEmpty();
+    }
+
+    @Test
+    @DisplayName("Native NOT applied for a pre-cutoff same-family gpt-4o snapshot")
+    void testNativeNotAppliedForPreCutoffSnapshot() {
+        // gpt-4o-2024-05-13 predates the Structured Outputs cutoff even though it shares the gpt-4o
+        // prefix; treating it as capable would fail silently at the provider.
+        ChatCompletionCreateParams params =
+                connection()
+                        .buildRequest(
+                                userMessage(),
+                                List.of(),
+                                params("gpt-4o-2024-05-13"),
+                                Person.class);
+
+        assertThat(params.responseFormat()).isEmpty();
+    }
+
+    @Test
+    @DisplayName("Native NOT applied when no output schema is supplied")
+    void testNativeNotAppliedWhenSchemaNull() {
+        ChatCompletionCreateParams params =
+                connection().buildRequest(userMessage(), List.of(), params("gpt-4o"), null);
+
+        assertThat(params.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 (not a bare POJO Class), so it must
+        // not activate native structured output; any non-Class schema object exercises the same
+        // instanceof gate.
+        Object nonClassSchema = "row<name STRING>";
+
+        ChatCompletionCreateParams params =
+                connection()
+                        .buildRequest(userMessage(), List.of(), params("gpt-4o"), nonClassSchema);
+
+        assertThat(params.responseFormat()).isEmpty();
+    }
+
+    @Test
+    @DisplayName("Native applied for a POJO even when tools are bound (no empty-tools gate)")
+    void testNativeAppliedEvenWhenToolsBound() {
+        ChatCompletionCreateParams params =
+                connection()
+                        .buildRequest(
+                                userMessage(),
+                                List.of(new StubTool()),
+                                params("gpt-4o"),
+                                Person.class);
+
+        assertThat(params.responseFormat()).isPresent();
+    }
+
+    @Test
+    @DisplayName("Capability predicate accepts the documented capable models")
+    void testCapabilityPredicateAcceptsCapableModels() {
+        OpenAICompletionsConnection connection = connection();
+
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o")).isTrue();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-08-06")).isTrue();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-11-20")).isTrue();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o-mini")).isTrue();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o-mini-2024-07-18")).isTrue();
+    }
+
+    @Test
+    @DisplayName("Capability predicate rejects incapable, pre-cutoff, unknown, and null models")
+    void testCapabilityPredicateRejectsIncapableModels() {
+        OpenAICompletionsConnection connection = connection();
+
+        assertThat(connection.supportsNativeStructuredOutput("gpt-3.5-turbo")).isFalse();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4-turbo")).isFalse();
+        assertThat(connection.supportsNativeStructuredOutput("gpt-4o-2024-05-13")).isFalse();
+        assertThat(connection.supportsNativeStructuredOutput("some-unknown-model")).isFalse();
+        assertThat(connection.supportsNativeStructuredOutput(null)).isFalse();
+    }
+
+    /** 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);
+        }
+    }
+}