[pull] main from apache:main

.github/workflows/pr_labeler.js

@@ -66,7 +66,7 @@ module.exports = async ({ github, context, core, workflowRunId, triggerWorkflow
     const baseBranch = pr.base.ref;
     const fixVersionLabelMap = {
       'main': 'fixVersion/0.4.0',
-      'release-0.3': 'fixVersion/0.3.1',
+      'release-0.3': 'fixVersion/0.3.2',
       'release-0.2': 'fixVersion/0.2.2',
       'release-0.1': 'fixVersion/0.1.2',
     };

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/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelConnection.java

@@ -80,4 +80,13 @@ public ResourceType getResourceType() {
      *     embeddings. The length of each array is determined by the model itself.
      */
     public abstract List<float[]> embed(List<String> texts, Map<String, Object> parameters);
+
+    public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object> parameters) {
+        return new EmbeddingResult<>(embed(text, parameters), null);
+    }
+
+    public EmbeddingResult<List<float[]>> embedWithUsage(
+            List<String> texts, Map<String, Object> parameters) {
+        return new EmbeddingResult<>(embed(texts, parameters), null);
+    }
 }

api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java

@@ -109,6 +109,17 @@ public float[] embed(String text, Map<String, Object> parameters) {
         return getConnection().embed(text, params);
     }
 
+    public EmbeddingResult<float[]> embedWithUsage(String text) {
+        return embedWithUsage(text, Collections.emptyMap());
+    }
+
+    public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object> parameters) {
+        Map<String, Object> params = this.getParameters();
+        params.putAll(parameters);
+        BaseEmbeddingModelConnection currentConnection = getConnection();
+        return currentConnection.embedWithUsage(text, params);
+    }
+
     /**
      * Generate embeddings for multiple texts.
      *
@@ -125,4 +136,16 @@ public List<float[]> embed(List<String> texts, Map<String, Object> parameters) {
         params.putAll(parameters);
         return getConnection().embed(texts, params);
     }
+
+    public EmbeddingResult<List<float[]>> embedWithUsage(List<String> texts) {
+        return embedWithUsage(texts, Collections.emptyMap());
+    }
+
+    public EmbeddingResult<List<float[]>> embedWithUsage(
+            List<String> texts, Map<String, Object> parameters) {
+        Map<String, Object> params = this.getParameters();
+        params.putAll(parameters);
+        BaseEmbeddingModelConnection currentConnection = getConnection();
+        return currentConnection.embedWithUsage(texts, params);
+    }
 }

api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingModelUtils.java

@@ -17,7 +17,9 @@
  */
 package org.apache.flink.agents.api.embedding.model;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 public class EmbeddingModelUtils {
     public static float[] toFloatArray(List list) {
@@ -34,4 +36,66 @@ public static float[] toFloatArray(List list) {
         }
         return array;
     }
+
+    public static EmbeddingResult<float[]> toSingleEmbeddingResult(Object result) {
+        Map<?, ?> values = toResultMap(result);
+        return new EmbeddingResult<>(
+                toFloatArray(toEmbeddingList(values.get("embeddings"))),
+                toTokenUsage(values.get("token_usage")));
+    }
+
+    public static EmbeddingResult<List<float[]>> toBatchEmbeddingResult(Object result) {
+        Map<?, ?> values = toResultMap(result);
+        List<?> rawEmbeddings = toEmbeddingList(values.get("embeddings"));
+        List<float[]> embeddings = new ArrayList<>();
+        for (Object embedding : rawEmbeddings) {
+            embeddings.add(toFloatArray(toEmbeddingList(embedding)));
+        }
+        return new EmbeddingResult<>(embeddings, toTokenUsage(values.get("token_usage")));
+    }
+
+    private static Map<?, ?> toResultMap(Object result) {
+        if (result instanceof Map) {
+            return (Map<?, ?>) result;
+        }
+        throw new IllegalArgumentException(
+                "Expected Map from Python embed_with_usage method, but got: "
+                        + (result == null ? "null" : result.getClass().getName()));
+    }
+
+    private static List<?> toEmbeddingList(Object embeddings) {
+        if (embeddings instanceof List) {
+            return (List<?>) embeddings;
+        }
+        throw new IllegalArgumentException(
+                "Expected List value in Python embedding result, but got: "
+                        + (embeddings == null ? "null" : embeddings.getClass().getName()));
+    }
+
+    private static EmbeddingTokenUsage toTokenUsage(Object tokenUsage) {
+        if (tokenUsage == null) {
+            return null;
+        }
+        if (!(tokenUsage instanceof Map)) {
+            throw new IllegalArgumentException(
+                    "Expected Map token_usage in Python embedding result, but got: "
+                            + tokenUsage.getClass().getName());
+        }
+
+        Map<?, ?> usage = (Map<?, ?>) tokenUsage;
+        return new EmbeddingTokenUsage(
+                toLong(usage.get("prompt_tokens"), "prompt_tokens"),
+                toLong(usage.get("total_tokens"), "total_tokens"));
+    }
+
+    private static long toLong(Object value, String fieldName) {
+        if (value instanceof Number) {
+            return ((Number) value).longValue();
+        }
+        throw new IllegalArgumentException(
+                "Expected numeric "
+                        + fieldName
+                        + " in Python embedding token usage, but got: "
+                        + (value == null ? "null" : value.getClass().getName()));
+    }
 }