[api][python][java] Track embedding token usage metrics

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()));
+    }
 }

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

@@ -0,0 +1,42 @@
+/*
+ * 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.embedding.model;
+
+import javax.annotation.Nullable;
+
+/** Embedding provider result with optional token usage metadata. */
+public class EmbeddingResult<T> {
+    private final T embeddings;
+
+    @Nullable private final EmbeddingTokenUsage tokenUsage;
+
+    public EmbeddingResult(T embeddings, @Nullable EmbeddingTokenUsage tokenUsage) {
+        this.embeddings = embeddings;
+        this.tokenUsage = tokenUsage;
+    }
+
+    public T getEmbeddings() {
+        return embeddings;
+    }
+
+    @Nullable
+    public EmbeddingTokenUsage getTokenUsage() {
+        return tokenUsage;
+    }
+}

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

@@ -0,0 +1,38 @@
+/*
+ * 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.embedding.model;
+
+/** Token usage reported by an embedding provider. */
+public class EmbeddingTokenUsage {
+    private final long promptTokens;
+    private final long totalTokens;
+
+    public EmbeddingTokenUsage(long promptTokens, long totalTokens) {
+        this.promptTokens = promptTokens;
+        this.totalTokens = totalTokens;
+    }
+
+    public long getPromptTokens() {
+        return promptTokens;
+    }
+
+    public long getTotalTokens() {
+        return totalTokens;
+    }
+}

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

@@ -19,6 +19,7 @@
 
 import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection;
 import org.apache.flink.agents.api.embedding.model.EmbeddingModelUtils;
+import org.apache.flink.agents.api.embedding.model.EmbeddingResult;
 import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
 import org.apache.flink.agents.api.resource.ResourceContext;
 import org.apache.flink.agents.api.resource.ResourceDescriptor;
@@ -42,6 +43,9 @@
 public class PythonEmbeddingModelConnection extends BaseEmbeddingModelConnection
         implements PythonResourceWrapper {
 
+    private static final String CALL_EMBED_WITH_USAGE =
+            "python_java_utils.call_embedding_with_usage";
+
     private final PyObject embeddingModel;
     private final PythonResourceAdapter adapter;
 
@@ -119,6 +123,31 @@ public List<float[]> embed(List<String> texts, Map<String, Object> parameters) {
                         + (results == null ? "null" : results.getClass().getName()));
     }
 
+    @Override
+    public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object> parameters) {
+        checkState(
+                embeddingModel != null,
+                "EmbeddingModelSetup is not initialized. Cannot perform embed operation.");
+
+        Map<String, Object> kwargs = new HashMap<>(parameters);
+        kwargs.put("text", text);
+        Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModel, kwargs);
+        return EmbeddingModelUtils.toSingleEmbeddingResult(result);
+    }
+
+    @Override
+    public EmbeddingResult<List<float[]>> embedWithUsage(
+            List<String> texts, Map<String, Object> parameters) {
+        checkState(
+                embeddingModel != null,
+                "EmbeddingModelSetup is not initialized. Cannot perform embed operation.");
+
+        Map<String, Object> kwargs = new HashMap<>(parameters);
+        kwargs.put("text", texts);
+        Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModel, kwargs);
+        return EmbeddingModelUtils.toBatchEmbeddingResult(result);
+    }
+
     @Override
     public Object getPythonResource() {
         return embeddingModel;

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

@@ -19,6 +19,7 @@
 
 import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup;
 import org.apache.flink.agents.api.embedding.model.EmbeddingModelUtils;
+import org.apache.flink.agents.api.embedding.model.EmbeddingResult;
 import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
 import org.apache.flink.agents.api.resource.ResourceContext;
 import org.apache.flink.agents.api.resource.ResourceDescriptor;
@@ -42,6 +43,9 @@
  */
 public class PythonEmbeddingModelSetup extends BaseEmbeddingModelSetup
         implements PythonResourceWrapper {
+    private static final String CALL_EMBED_WITH_USAGE =
+            "python_java_utils.call_embedding_with_usage";
+
     private final PyObject embeddingModelSetup;
     private final PythonResourceAdapter adapter;
 
@@ -124,6 +128,31 @@ public List<float[]> embed(List<String> texts, Map<String, Object> parameters) {
                         + (results == null ? "null" : results.getClass().getName()));
     }
 
+    @Override
+    public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object> parameters) {
+        checkState(
+                embeddingModelSetup != null,
+                "EmbeddingModelSetup is not initialized. Cannot perform embed operation.");
+
+        Map<String, Object> kwargs = new HashMap<>(parameters);
+        kwargs.put("text", text);
+        Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModelSetup, kwargs);
+        return EmbeddingModelUtils.toSingleEmbeddingResult(result);
+    }
+
+    @Override
+    public EmbeddingResult<List<float[]>> embedWithUsage(
+            List<String> texts, Map<String, Object> parameters) {
+        checkState(
+                embeddingModelSetup != null,
+                "EmbeddingModelSetup is not initialized. Cannot perform embed operation.");
+
+        Map<String, Object> kwargs = new HashMap<>(parameters);
+        kwargs.put("text", texts);
+        Object result = adapter.invoke(CALL_EMBED_WITH_USAGE, embeddingModelSetup, kwargs);
+        return EmbeddingModelUtils.toBatchEmbeddingResult(result);
+    }
+
     @Override
     public Map<String, Object> getParameters() {
         return Map.of();