[api][python] Add explicit output schema parameter to the chat path (structured-output foundation)

api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java

@@ -25,6 +25,8 @@
 import org.apache.flink.agents.api.resource.ResourceType;
 import org.apache.flink.agents.api.tools.Tool;
 
+import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Map;
 
@@ -45,6 +47,26 @@ public ResourceType getResourceType() {
         return ResourceType.CHAT_MODEL_CONNECTION;
     }
 
+    /**
+     * Whether this connection can apply the provider's native structured-output API for the given
+     * model.
+     *
+     * <p>Capability is <b>model-dependent</b>, not connection-wide: a single provider connection
+     * commonly serves both models that accept a native schema parameter and models that do not. It
+     * is therefore evaluated against the <i>effective</i> model at request-build time — the model
+     * actually being called, which per-request parameters may override.
+     *
+     * <p>The default {@code false} keeps a connection on the prompt-engineering fallback. An
+     * unrecognized model must report {@code false} so that it degrades to the fallback rather than
+     * failing at the provider.
+     *
+     * @param effectiveModel the model the request will be issued against, may be null
+     * @return true if a schema can be applied natively for {@code effectiveModel}
+     */
+    protected boolean supportsNativeStructuredOutput(String effectiveModel) {
+        return false;
+    }
+
     /**
      * Process a chat request and return a chat response.
      *
@@ -55,4 +77,44 @@ public ResourceType getResourceType() {
      */
     public abstract ChatMessage chat(
             List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams);
+
+    /**
+     * Process a chat request that carries an output schema, and return a chat response.
+     *
+     * <p>{@code outputSchema} is framework-level execution metadata, kept off {@code modelParams}
+     * so that it can never reach a provider SDK request as a generation parameter. It is either a
+     * POJO {@link Class} or an {@link org.apache.flink.agents.api.agents.OutputSchema} (a {@code
+     * RowTypeInfo} wrapper); the two cases are distinguished by the connection that consumes it.
+     *
+     * <p>A schema must not be handed to a connection that has no native translation for it: this
+     * default implementation rejects a non-null {@code outputSchema} rather than dropping it, so an
+     * unconstrained response can never be mistaken for a schema-conforming one. A null {@code
+     * outputSchema} delegates to {@link #chat(List, List, Map)}. A connection that does translate a
+     * schema into a native provider parameter overrides this overload, and reports its capability
+     * via {@link #supportsNativeStructuredOutput(String)}.
+     *
+     * @param messages the input chat messages
+     * @param tools the tools can be called by the model
+     * @param modelParams the additional arguments passed to the model
+     * @param outputSchema the schema the response should conform to, or null for an unconstrained
+     *     response
+     * @return the chat response containing model outputs
+     * @throws UnsupportedOperationException if {@code outputSchema} is non-null and this connection
+     *     has no native structured-output translation
+     */
+    public ChatMessage chat(
+            List<ChatMessage> messages,
+            List<Tool> tools,
+            Map<String, Object> modelParams,
+            @Nullable Object outputSchema) {
+        if (outputSchema != null) {
+            throw new UnsupportedOperationException(
+                    getClass().getName()
+                            + " has no native structured-output translation, so it cannot honor"
+                            + " the given output schema. Override chat(List, List, Map, Object) to"
+                            + " translate the schema natively, or pass no schema so the caller"
+                            + " applies the prompt-engineering fallback.");
+        }
+        return chat(messages, tools, modelParams);
+    }
 }

api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java

@@ -48,6 +48,7 @@ public abstract class BaseChatModelSetup extends Resource {
     @Nullable protected String skillDiscoveryPrompt;
     protected List<String> allowedCommands;
     protected List<String> allowedScriptDirs;
+    protected StructuredOutputStrategy structuredOutputStrategy;
 
     @Nullable protected BaseChatModelConnection connection;
     protected final List<Tool> tools = new ArrayList<>();
@@ -67,6 +68,10 @@ public BaseChatModelSetup(ResourceDescriptor descriptor, ResourceContext resourc
                 declaredScriptDirs == null
                         ? new ArrayList<>()
                         : new ArrayList<>(declaredScriptDirs);
+        this.structuredOutputStrategy =
+                StructuredOutputStrategy.fromArgument(
+                        descriptor.getArgument("structured_output_strategy"),
+                        StructuredOutputStrategy.AUTO);
     }
 
     /**
@@ -219,4 +224,15 @@ public List<String> getAllowedCommands() {
     public List<String> getAllowedScriptDirs() {
         return allowedScriptDirs;
     }
+
+    /**
+     * The configured intent about how an output schema should be applied, defaulting to {@link
+     * StructuredOutputStrategy#AUTO}. Whether native structured output is actually applied combines
+     * this policy with the connection's model-dependent capability.
+     *
+     * @return the structured output strategy
+     */
+    public StructuredOutputStrategy getStructuredOutputStrategy() {
+        return structuredOutputStrategy;
+    }
 }

api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java

@@ -0,0 +1,112 @@
+/*
+ * 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.api.chat.model;
+
+import java.util.Locale;
+
+/**
+ * User intent about how an output schema should be applied to a chat request.
+ *
+ * <p>This expresses <b>policy</b> only. Whether a connection <i>can</i> apply the provider's native
+ * structured-output API is a separate, model-dependent <b>capability</b> question answered by
+ * {@link BaseChatModelConnection#supportsNativeStructuredOutput(String)}. Policy and capability are
+ * combined at request-build time.
+ */
+public enum StructuredOutputStrategy {
+    /**
+     * Use the provider's native structured-output API when the effective model is capable of it,
+     * and fall back to prompt engineering otherwise. This is the default.
+     */
+    AUTO,
+
+    /**
+     * Always use the provider's native structured-output API, without consulting the capability
+     * predicate.
+     */
+    NATIVE,
+
+    /**
+     * Never use the provider's native structured-output API; rely on prompt engineering alone. This
+     * matches the behavior of connections that have no native translation.
+     */
+    PROMPT;
+
+    /**
+     * Resolves this policy against a connection's model-dependent capability into whether the
+     * provider's native structured-output API should be used.
+     *
+     * <ul>
+     *   <li>{@code AUTO} defers to {@code modelCapable}: native when the effective model can, else
+     *       the prompt-engineering fallback.
+     *   <li>{@code NATIVE} always resolves to native, ignoring {@code modelCapable}, so an explicit
+     *       user intent surfaces a provider error rather than silently degrading.
+     *   <li>{@code PROMPT} never resolves to native.
+     * </ul>
+     *
+     * @param modelCapable whether the connection reports the effective model as natively capable
+     * @return true if native structured output should be applied
+     */
+    public boolean resolvesToNative(boolean modelCapable) {
+        switch (this) {
+            case NATIVE:
+                return true;
+            case PROMPT:
+                return false;
+            case AUTO:
+            default:
+                return modelCapable;
+        }
+    }
+
+    /**
+     * Resolves a strategy from a descriptor argument, which may arrive either as a {@code
+     * StructuredOutputStrategy} or — across the Python bridge, where arguments are carried as JSON
+     * — as its case-insensitive name.
+     *
+     * @param value the raw descriptor argument, may be null
+     * @param defaultValue the strategy to use when {@code value} is null
+     * @return the resolved strategy
+     * @throws IllegalArgumentException if {@code value} is neither null, a {@code
+     *     StructuredOutputStrategy}, nor the name of one
+     */
+    public static StructuredOutputStrategy fromArgument(
+            Object value, StructuredOutputStrategy defaultValue) {
+        if (value == null) {
+            return defaultValue;
+        }
+        if (value instanceof StructuredOutputStrategy) {
+            return (StructuredOutputStrategy) value;
+        }
+        if (value instanceof String) {
+            try {
+                return valueOf(((String) value).toUpperCase(Locale.ROOT));
+            } catch (IllegalArgumentException e) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "Unknown structured output strategy '%s'. Expected one of: AUTO, NATIVE, PROMPT.",
+                                value),
+                        e);
+            }
+        }
+        throw new IllegalArgumentException(
+                String.format(
+                        "Unsupported structured output strategy type '%s'. Expected a StructuredOutputStrategy or its name.",
+                        value.getClass().getName()));
+    }
+}

api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java

@@ -237,6 +237,7 @@ void testChatResponseFormat() {
     /** Connection that captures the messages passed to it for assertions. */
     private static class RecordingConnection extends BaseChatModelConnection {
         List<ChatMessage> capturedMessages;
+        Map<String, Object> capturedModelParams;
 
         RecordingConnection() {
             super(
@@ -249,6 +250,7 @@ private static class RecordingConnection extends BaseChatModelConnection {
         public ChatMessage chat(
                 List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {
             this.capturedMessages = new ArrayList<>(messages);
+            this.capturedModelParams = new HashMap<>(modelParams);
             return new ChatMessage(MessageRole.ASSISTANT, "ok");
         }
     }
@@ -320,6 +322,124 @@ void testChatRefillsTemplateOnSubsequentInvocations() {
         assertEquals("tool result", connection.capturedMessages.get(1).getContent());
     }
 
+    @Test
+    @DisplayName("Default chat() overload rejects an outputSchema it cannot translate")
+    void testDefaultChatOverloadRejectsOutputSchema() {
+        RecordingConnection connection = new RecordingConnection();
+
+        // Dropping the schema instead would return an unconstrained response that the
+        // caller has no way to tell apart from a schema-conforming one.
+        assertThrows(
+                UnsupportedOperationException.class,
+                () ->
+                        connection.chat(
+                                List.of(new ChatMessage(MessageRole.USER, "hi")),
+                                List.of(),
+                                new HashMap<>(),
+                                new Object()));
+
+        // The rejection has to precede the delegation: a delegate-then-throw ordering
+        // would still issue a real provider request before failing.
+        assertNull(connection.capturedMessages);
+    }
+
+    @Test
+    @DisplayName("Default chat() overload delegates to the 3-arg chat() for a null outputSchema")
+    void testDefaultChatOverloadDelegatesForNullOutputSchema() {
+        RecordingConnection connection = new RecordingConnection();
+        Map<String, Object> modelParams = new HashMap<>();
+        modelParams.put("temperature", 0.5);
+
+        ChatMessage response =
+                connection.chat(
+                        List.of(new ChatMessage(MessageRole.USER, "hi")),
+                        List.of(),
+                        modelParams,
+                        null);
+
+        // The 3-arg chat() ran (it is what produces "ok") and the overload added nothing
+        // to modelParams that could travel on to a provider SDK request.
+        assertEquals("ok", response.getContent());
+        assertEquals(Map.of("temperature", 0.5), connection.capturedModelParams);
+    }
+
+    @Test
+    @DisplayName("Default capability predicate reports no native structured output for any model")
+    void testDefaultCapabilityPredicateIsFalse() {
+        RecordingConnection connection = new RecordingConnection();
+
+        assertFalse(connection.supportsNativeStructuredOutput("gpt-4o"));
+        assertFalse(connection.supportsNativeStructuredOutput("gpt-3.5-turbo"));
+        assertFalse(connection.supportsNativeStructuredOutput(null));
+    }
+
+    @Test
+    @DisplayName("Structured-output strategy defaults to AUTO when the descriptor omits it")
+    void testStructuredOutputStrategyDefaultsToAuto() {
+        RecordingChatModelSetup setup =
+                new RecordingChatModelSetup(new RecordingConnection(), null);
+
+        assertEquals(StructuredOutputStrategy.AUTO, setup.getStructuredOutputStrategy());
+    }
+
+    @Test
+    @DisplayName("Structured-output strategy defaults to AUTO when the descriptor argument is null")
+    void testStructuredOutputStrategyDefaultsToAutoForNullArgument() {
+        // A descriptor argument present with a null value is indistinguishable from an
+        // absent one here, so it resolves to the same default rather than failing.
+        TestChatModel model =
+                new TestChatModel(
+                        new ResourceDescriptor(
+                                TestChatModel.class.getName(),
+                                Collections.singletonMap("structured_output_strategy", null)),
+                        null);
+
+        assertEquals(StructuredOutputStrategy.AUTO, model.getStructuredOutputStrategy());
+    }
+
+    @Test
+    @DisplayName("Structured-output strategy is read from the descriptor argument")
+    void testStructuredOutputStrategyReadFromDescriptor() {
+        TestChatModel model =
+                new TestChatModel(
+                        new ResourceDescriptor(
+                                TestChatModel.class.getName(),
+                                Map.of("structured_output_strategy", "native")),
+                        null);
+
+        assertEquals(StructuredOutputStrategy.NATIVE, model.getStructuredOutputStrategy());
+    }
+
+    @Test
+    @DisplayName("An unrecognized structured-output strategy is rejected instead of defaulting")
+    void testUnknownStructuredOutputStrategyRejected() {
+        ResourceDescriptor descriptor =
+                new ResourceDescriptor(
+                        TestChatModel.class.getName(),
+                        Map.of("structured_output_strategy", "bogus"));
+
+        assertThrows(IllegalArgumentException.class, () -> new TestChatModel(descriptor, null));
+    }
+
+    @Test
+    @DisplayName("AUTO resolves to native only when the effective model is capable")
+    void testAutoStrategyResolvesToNativeOnlyWhenCapable() {
+        assertTrue(StructuredOutputStrategy.AUTO.resolvesToNative(true));
+        assertFalse(StructuredOutputStrategy.AUTO.resolvesToNative(false));
+    }
+
+    @Test
+    @DisplayName("NATIVE forces native even when the model is not capable")
+    void testNativeStrategyForcesNativeRegardlessOfCapability() {
+        assertTrue(StructuredOutputStrategy.NATIVE.resolvesToNative(false));
+    }
+
+    @Test
+    @DisplayName("PROMPT never resolves to native even when the model is capable")
+    void testPromptStrategyNeverResolvesToNative() {
+        assertFalse(StructuredOutputStrategy.PROMPT.resolvesToNative(true));
+    }
+
     @Test
     @DisplayName("Test chat with long input")
     void testChatWithLongInput() {