diff --git a/.github/workflows/pr_labeler.js b/.github/workflows/pr_labeler.js index 9228c0781..6f5ac9309 100644 --- a/.github/workflows/pr_labeler.js +++ b/.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', }; diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java index 5181073a7..b4200958a 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelConnection.java +++ b/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. + * + *

Capability is model-dependent, 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 effective model at request-build time — the model + * actually being called, which per-request parameters may override. + * + *

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 messages, List tools, Map modelParams); + + /** + * Process a chat request that carries an output schema, and return a chat response. + * + *

{@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. + * + *

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 messages, + List tools, + Map 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); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java index 34b2b2994..6c59f6ef7 100644 --- a/api/src/main/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetup.java +++ b/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 allowedCommands; protected List allowedScriptDirs; + protected StructuredOutputStrategy structuredOutputStrategy; @Nullable protected BaseChatModelConnection connection; protected final List 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 getAllowedCommands() { public List 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; + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java b/api/src/main/java/org/apache/flink/agents/api/chat/model/StructuredOutputStrategy.java new file mode 100644 index 000000000..4026d1992 --- /dev/null +++ b/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. + * + *

This expresses policy only. Whether a connection can apply the provider's native + * structured-output API is a separate, model-dependent capability 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. + * + *

+ * + * @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())); + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelConnection.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelConnection.java index 4d46e3eed..01210e8f5 100644 --- a/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelConnection.java +++ b/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 embed(List texts, Map parameters); + + public EmbeddingResult embedWithUsage(String text, Map parameters) { + return new EmbeddingResult<>(embed(text, parameters), null); + } + + public EmbeddingResult> embedWithUsage( + List texts, Map parameters) { + return new EmbeddingResult<>(embed(texts, parameters), null); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java index e7c19893d..605c189d5 100644 --- a/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java +++ b/api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java @@ -109,6 +109,17 @@ public float[] embed(String text, Map parameters) { return getConnection().embed(text, params); } + public EmbeddingResult embedWithUsage(String text) { + return embedWithUsage(text, Collections.emptyMap()); + } + + public EmbeddingResult embedWithUsage(String text, Map parameters) { + Map 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 embed(List texts, Map parameters) { params.putAll(parameters); return getConnection().embed(texts, params); } + + public EmbeddingResult> embedWithUsage(List texts) { + return embedWithUsage(texts, Collections.emptyMap()); + } + + public EmbeddingResult> embedWithUsage( + List texts, Map parameters) { + Map params = this.getParameters(); + params.putAll(parameters); + BaseEmbeddingModelConnection currentConnection = getConnection(); + return currentConnection.embedWithUsage(texts, params); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingModelUtils.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingModelUtils.java index c1c71e589..08128eb45 100644 --- a/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingModelUtils.java +++ b/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 toSingleEmbeddingResult(Object result) { + Map values = toResultMap(result); + return new EmbeddingResult<>( + toFloatArray(toEmbeddingList(values.get("embeddings"))), + toTokenUsage(values.get("token_usage"))); + } + + public static EmbeddingResult> toBatchEmbeddingResult(Object result) { + Map values = toResultMap(result); + List rawEmbeddings = toEmbeddingList(values.get("embeddings")); + List 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())); + } } diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingResult.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingResult.java new file mode 100644 index 000000000..8c23d14d7 --- /dev/null +++ b/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 { + 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; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingTokenUsage.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/EmbeddingTokenUsage.java new file mode 100644 index 000000000..be1abdd98 --- /dev/null +++ b/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; + } +} diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java index 785ed03a8..d9bdb8d41 100644 --- a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java +++ b/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 embed(List texts, Map parameters) { + (results == null ? "null" : results.getClass().getName())); } + @Override + public EmbeddingResult embedWithUsage(String text, Map parameters) { + checkState( + embeddingModel != null, + "EmbeddingModelSetup is not initialized. Cannot perform embed operation."); + + Map 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> embedWithUsage( + List texts, Map parameters) { + checkState( + embeddingModel != null, + "EmbeddingModelSetup is not initialized. Cannot perform embed operation."); + + Map 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; diff --git a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java index c0efd5c35..b6febbc58 100644 --- a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java +++ b/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 embed(List texts, Map parameters) { + (results == null ? "null" : results.getClass().getName())); } + @Override + public EmbeddingResult embedWithUsage(String text, Map parameters) { + checkState( + embeddingModelSetup != null, + "EmbeddingModelSetup is not initialized. Cannot perform embed operation."); + + Map 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> embedWithUsage( + List texts, Map parameters) { + checkState( + embeddingModelSetup != null, + "EmbeddingModelSetup is not initialized. Cannot perform embed operation."); + + Map 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 getParameters() { return Map.of(); diff --git a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java index 43f8c8b05..d27d79cb8 100644 --- a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelTest.java +++ b/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 capturedMessages; + Map capturedModelParams; RecordingConnection() { super( @@ -249,6 +250,7 @@ private static class RecordingConnection extends BaseChatModelConnection { public ChatMessage chat( List messages, List tools, Map 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 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() { diff --git a/api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupEmbeddingResultTest.java b/api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupEmbeddingResultTest.java new file mode 100644 index 000000000..fbbf738da --- /dev/null +++ b/api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupEmbeddingResultTest.java @@ -0,0 +1,114 @@ +/* + * 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 org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; + +/** Test cases for embedding results returned by model setups. */ +class BaseEmbeddingModelSetupEmbeddingResultTest { + + @Test + void testEmbedWithUsageDelegatesProviderUsage() { + BaseEmbeddingModelSetup setup = + new BaseEmbeddingModelSetup( + new ResourceDescriptor( + "test", Map.of("connection", "connection", "model", "model")), + mock(ResourceContext.class)) { + @Override + public Map getParameters() { + return Collections.emptyMap(); + } + }; + setup.connection = + new BaseEmbeddingModelConnection( + new ResourceDescriptor("connection", Collections.emptyMap()), + mock(ResourceContext.class)) { + @Override + public float[] embed(String text, Map parameters) { + return new float[] {0.1f, 0.2f}; + } + + @Override + public java.util.List embed( + java.util.List texts, Map parameters) { + throw new UnsupportedOperationException(); + } + + @Override + public EmbeddingResult embedWithUsage( + String text, Map parameters) { + return new EmbeddingResult<>( + embed(text, parameters), new EmbeddingTokenUsage(7L, 9L)); + } + }; + + EmbeddingResult result = setup.embedWithUsage("hello"); + + assertArrayEquals(new float[] {0.1f, 0.2f}, result.getEmbeddings()); + assertEquals(7L, result.getTokenUsage().getPromptTokens()); + assertEquals(9L, result.getTokenUsage().getTotalTokens()); + } + + @Test + void testEmbedDelegatesToExistingConnectionMethod() { + BaseEmbeddingModelSetup setup = + new BaseEmbeddingModelSetup( + new ResourceDescriptor( + "test", Map.of("connection", "connection", "model", "model")), + mock(ResourceContext.class)) { + @Override + public Map getParameters() { + return Collections.emptyMap(); + } + }; + setup.connection = + new BaseEmbeddingModelConnection( + new ResourceDescriptor("connection", Collections.emptyMap()), + mock(ResourceContext.class)) { + @Override + public float[] embed(String text, Map parameters) { + return new float[] {0.1f, 0.2f}; + } + + @Override + public java.util.List embed( + java.util.List texts, Map parameters) { + throw new UnsupportedOperationException(); + } + + @Override + public EmbeddingResult embedWithUsage( + String text, Map parameters) { + return new EmbeddingResult<>( + new float[] {0.3f, 0.4f}, new EmbeddingTokenUsage(7L, 9L)); + } + }; + + assertArrayEquals(new float[] {0.1f, 0.2f}, setup.embed("hello")); + } +} diff --git a/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnectionTest.java b/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnectionTest.java index 470d51270..a89bf3313 100644 --- a/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnectionTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnectionTest.java @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.api.embedding.model.python; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.python.PythonResourceAdapter; @@ -136,6 +137,42 @@ void testEmbedSingleTextWithEmptyParameters() { assertThat(result).hasSize(2); } + @Test + void testEmbedWithUsageMultipleTexts() { + List texts = List.of("first", "second"); + when(mockAdapter.invoke( + eq("python_java_utils.call_embedding_with_usage"), + eq(mockEmbeddingModel), + any(Map.class))) + .thenReturn( + Map.of( + "embeddings", + List.of(List.of(0.1, 0.2), List.of(0.3, 0.4)), + "token_usage", + Map.of("prompt_tokens", 7, "total_tokens", 9))); + + EmbeddingResult> result = + pythonEmbeddingModelConnection.embedWithUsage(texts, Map.of("batch_size", 2)); + + assertThat(result.getEmbeddings()).hasSize(2); + assertThat(result.getEmbeddings().get(0)).containsExactly(0.1f, 0.2f); + assertThat(result.getEmbeddings().get(1)).containsExactly(0.3f, 0.4f); + assertThat(result.getTokenUsage()).isNotNull(); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(7L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(9L); + verify(mockAdapter) + .invoke( + eq("python_java_utils.call_embedding_with_usage"), + eq(mockEmbeddingModel), + argThat( + kwargs -> { + Map values = (Map) kwargs; + assertThat(values).containsEntry("text", texts); + assertThat(values).containsEntry("batch_size", 2); + return true; + })); + } + @Test void testEmbedSingleTextWithNullEmbeddingModelThrowsException() { PythonEmbeddingModelConnection connectionWithNullModel = diff --git a/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetupTest.java b/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetupTest.java index d8071d7cd..430ddc4c3 100644 --- a/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetupTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetupTest.java @@ -17,6 +17,7 @@ */ package org.apache.flink.agents.api.embedding.model.python; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.python.PythonResourceAdapter; @@ -143,6 +144,41 @@ void testEmbedSingleTextWithEmptyParameters() { assertThat(result).hasSize(2); } + @Test + void testEmbedWithUsageSingleText() { + String text = "test text"; + Map parameters = Map.of("model", "test-model"); + when(mockAdapter.invoke( + eq("python_java_utils.call_embedding_with_usage"), + eq(mockEmbeddingModelSetup), + any(Map.class))) + .thenReturn( + Map.of( + "embeddings", + List.of(0.1, 0.2), + "token_usage", + Map.of("prompt_tokens", 7, "total_tokens", 9))); + + EmbeddingResult result = + pythonEmbeddingModelSetup.embedWithUsage(text, parameters); + + assertThat(result.getEmbeddings()).containsExactly(0.1f, 0.2f); + assertThat(result.getTokenUsage()).isNotNull(); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(7L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(9L); + verify(mockAdapter) + .invoke( + eq("python_java_utils.call_embedding_with_usage"), + eq(mockEmbeddingModelSetup), + argThat( + kwargs -> { + Map values = (Map) kwargs; + assertThat(values).containsEntry("text", text); + assertThat(values).containsEntry("model", "test-model"); + return true; + })); + } + @Test void testEmbedSingleTextWithNullEmbeddingModelSetupThrowsException() { PythonEmbeddingModelSetup setupWithNullModel = diff --git a/code_review.md b/code_review.md index b99d41d17..f78e7bd6b 100644 --- a/code_review.md +++ b/code_review.md @@ -30,6 +30,21 @@ Keep these findings separate. A design-focused comment does not need to prove an immediate runtime bug, but it must explain the concrete complexity, ambiguity, or maintenance risk introduced by the change. +## Change-Type Review Guides + +For common change types, start from the focused guide below as an on-ramp, then +apply the full passes in this document. Each guide narrows the checklist to the +concerns that matter most for that area and links examples from past reviews. +Guides load on demand, so the general passes here stay short. + +| Change type | Focus | Guide | +|---|---|---| +| `runtime/` state and recovery | serde and replay type fidelity, real failure-path tests | [review-guides/runtime-state-recovery.md](review-guides/runtime-state-recovery.md) | +| Python-Java bridge | cross-language parity, type mapping across Pemja | planned | +| `api/` contract | API shape, compatibility policy, deprecation | planned | +| `dist` and dependency | shading, LICENSE and NOTICE, dist registration | planned | +| docs-only | facts match their source of truth | planned | + ## Required Review Passes 1. Understand the issue or feature goal. diff --git a/docs/content/docs/development/chat_models.md b/docs/content/docs/development/chat_models.md index 51213b73b..1fd4f9fa1 100644 --- a/docs/content/docs/development/chat_models.md +++ b/docs/content/docs/development/chat_models.md @@ -643,6 +643,102 @@ Some popular options include: Model availability depends on your Azure region and subscription. Always check the official Azure OpenAI documentation for regional availability before implementing in production. {{< /hint >}} +### Gemini + +Google Gemini provides cloud-based chat models through the Gemini Developer API and Vertex AI. The Flink Agents Gemini integration uses the official Google Gen AI SDK and supports text conversations, system instructions, and tool calling. + +{{< hint warning >}} +Vertex AI support is experimental. The connection path has been smoke-tested during construction but has not yet been verified end to end. +{{< /hint >}} + +{{< hint info >}} +Gemini is only supported in Java currently. To use Gemini from Python agents, see [Using Cross-Language Providers](#using-cross-language-providers). +{{< /hint >}} + +#### Prerequisites + +1. For the Gemini Developer API, create an API key in [Google AI Studio](https://aistudio.google.com/app/apikey) +2. For Vertex AI, enable Vertex AI in your Google Cloud project and configure Google Cloud credentials + +#### GeminiChatModelConnection Parameters + +{{< tabs "GeminiChatModelConnection Parameters" >}} + +{{< tab "Java" >}} + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `api_key` | String | Required unless `base_url` is set or `vertex_ai` is `true` | Gemini Developer API key | +| `base_url` | String | None | Custom endpoint, such as a proxy that injects credentials | +| `model` | String | None | Default model name, used when no model is supplied per setup | +| `timeout` | int | None | API request timeout in seconds | +| `vertex_ai` | boolean | `false` | Use the experimental Vertex AI backend; not yet verified end to end | +| `project` | String | None | Vertex AI project id | +| `location` | String | None | Vertex AI location | + +{{< /tab >}} + +{{< /tabs >}} + +#### GeminiChatModelSetup Parameters + +{{< tabs "GeminiChatModelSetup Parameters" >}} + +{{< tab "Java" >}} + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `connection` | String | Required | Reference to connection method name | +| `model` | String | `"gemini-3.1-pro-preview"` | Name of the chat model to use | +| `prompt` | Prompt \| String | None | Prompt template or reference to prompt resource | +| `tools` | List | None | List of tool names available to the model | +| `temperature` | double | `0.1` | Sampling temperature (0.0 to 2.0) | +| `max_output_tokens` | long | `1024` | Maximum number of tokens to generate | +| `additional_kwargs` | Map | `{}` | Additional Gemini parameters (`top_k`, `top_p`, `stop_sequences`) | + +{{< /tab >}} + +{{< /tabs >}} + +#### Usage Example + +{{< tabs "Gemini Usage Example" >}} + +{{< tab "Java" >}} +```java +public class MyAgent extends Agent { + @ChatModelConnection + public static ResourceDescriptor geminiConnection() { + return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.GEMINI_CONNECTION) + .addInitialArgument("api_key", System.getenv("GEMINI_API_KEY")) + .build(); + } + + @ChatModelSetup + public static ResourceDescriptor geminiChatModel() { + return ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.GEMINI_SETUP) + .addInitialArgument("connection", "geminiConnection") + .addInitialArgument("model", "gemini-3.1-pro-preview") + .addInitialArgument("temperature", 0.1d) + .addInitialArgument("max_output_tokens", 1024) + .build(); + } + + ... +} +``` +{{< /tab >}} + +{{< /tabs >}} + +#### Available Models + +Visit the [Gemini models documentation](https://ai.google.dev/gemini-api/docs/models) for the complete and up-to-date list of available models. + +{{< hint warning >}} +Model availability and names may differ between the Gemini Developer API and Vertex AI. Always check the official Gemini documentation before implementing in production. +{{< /hint >}} + ### Ollama Ollama provides local chat models that run on your machine, offering privacy, control, and no API costs. @@ -1370,4 +1466,4 @@ public class MyChatModelSetup extends BaseChatModelSetup { The built-in `chat_model_action` listens to `ChatRequestEvent` and `ToolResponseEvent`. To request a chat completion, send a `ChatRequestEvent`. If the model returns a final answer, the action sends a `ChatResponseEvent`. -If the model asks to call tools, `chat_model_action` sends a `ToolRequestEvent` instead of a final `ChatResponseEvent`. After the tools finish, it receives the matching `ToolResponseEvent`, appends the tool results to the chat history, and calls the model again. This loop continues until the model returns a final response. For details on how tools are executed, see [Built-in Events and Actions in Tool Use]({{< ref "docs/development/tool_use#built-in-events-and-actions" >}}). \ No newline at end of file +If the model asks to call tools, `chat_model_action` sends a `ToolRequestEvent` instead of a final `ChatResponseEvent`. After the tools finish, it receives the matching `ToolResponseEvent`, appends the tool results to the chat history, and calls the model again. This loop continues until the model returns a final response. For details on how tools are executed, see [Built-in Events and Actions in Tool Use]({{< ref "docs/development/tool_use#built-in-events-and-actions" >}}). diff --git a/docs/content/docs/faq/faq.md b/docs/content/docs/faq/faq.md index bc0cb279c..4b9fd7e7e 100644 --- a/docs/content/docs/faq/faq.md +++ b/docs/content/docs/faq/faq.md @@ -101,6 +101,7 @@ Flink Agents provides built-in integrations for many ecosystem providers. Some i | [Anthropic]({{< ref "docs/development/chat_models#anthropic" >}}) | ✅ | ✅ | | [Azure AI]({{< ref "docs/development/chat_models#azure-ai" >}}) | ❌ | ✅ | | [Azure OpenAI]({{< ref "docs/development/chat_models#azure-openai" >}}) | ✅ | ✅ | +| [Gemini]({{< ref "docs/development/chat_models#gemini" >}}) | ❌ | ✅ | | [Ollama]({{< ref "docs/development/chat_models#ollama" >}}) | ✅ | ✅ | | [OpenAI]({{< ref "docs/development/chat_models#openai" >}}) | ✅ | ✅ | | [Tongyi (DashScope)]({{< ref "docs/development/chat_models#tongyi-dashscope" >}}) | ✅ | ❌ | diff --git a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/EmbeddingCrossLanguageAgent.java b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/EmbeddingCrossLanguageAgent.java index be5f8d494..14bfbbb4d 100644 --- a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/EmbeddingCrossLanguageAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/EmbeddingCrossLanguageAgent.java @@ -29,6 +29,7 @@ import org.apache.flink.agents.api.annotation.EmbeddingModelSetup; import org.apache.flink.agents.api.context.RunnerContext; import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceName; @@ -105,11 +106,14 @@ public static void testEmbeddingGeneration(Event event, RunnerContext ctx) throw org.apache.flink.agents.api.resource.ResourceType .EMBEDDING_MODEL); - float[] embedding = embeddingModel.embed(text); + EmbeddingResult embeddingResult = embeddingModel.embedWithUsage(text); + float[] embedding = embeddingResult.getEmbeddings(); System.out.printf("[TEST] Generated embedding with dimension: %d%n", embedding.length); validateEmbeddingResult(id, text, embedding); - List embeddings = embeddingModel.embed(List.of(text)); + EmbeddingResult> embeddingsResult = + embeddingModel.embedWithUsage(List.of(text)); + List embeddings = embeddingsResult.getEmbeddings(); validateEmbeddingResults(id, List.of(text), embeddings); // Create a minimal test result to avoid serialization issues diff --git a/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java index a7dc2926c..b39ec6b20 100644 --- a/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java +++ b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java @@ -23,6 +23,8 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import org.apache.flink.agents.api.RetryExecutor; import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; +import org.apache.flink.agents.api.embedding.model.EmbeddingTokenUsage; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; @@ -80,36 +82,67 @@ public class BedrockEmbeddingModelConnection extends BaseEmbeddingModelConnectio public BedrockEmbeddingModelConnection( ResourceDescriptor descriptor, ResourceContext resourceContext) { super(descriptor, resourceContext); + this.client = createClient(resolveRegion(descriptor)); + this.defaultModel = resolveDefaultModel(descriptor); + this.embedPool = Executors.newFixedThreadPool(resolveEmbedConcurrency(descriptor)); + this.retryExecutor = createRetryExecutor(descriptor); + } + + BedrockEmbeddingModelConnection( + ResourceDescriptor descriptor, + ResourceContext resourceContext, + BedrockRuntimeClient client, + ExecutorService embedPool, + RetryExecutor retryExecutor, + String defaultModel) { + super(descriptor, resourceContext); + this.client = client; + this.defaultModel = defaultModel; + this.embedPool = embedPool; + this.retryExecutor = retryExecutor; + } + + private static BedrockRuntimeClient createClient(String region) { + return BedrockRuntimeClient.builder() + .region(Region.of(region)) + .credentialsProvider(DefaultCredentialsProvider.create()) + .build(); + } + private static String resolveRegion(ResourceDescriptor descriptor) { String region = descriptor.getArgument("region"); if (region == null || region.isBlank()) { region = "us-east-1"; } + return region; + } - this.client = - BedrockRuntimeClient.builder() - .region(Region.of(region)) - .credentialsProvider(DefaultCredentialsProvider.create()) - .build(); - + private static String resolveDefaultModel(ResourceDescriptor descriptor) { String model = descriptor.getArgument("model"); - this.defaultModel = (model != null && !model.isBlank()) ? model : DEFAULT_MODEL; + return (model != null && !model.isBlank()) ? model : DEFAULT_MODEL; + } + private static int resolveEmbedConcurrency(ResourceDescriptor descriptor) { Integer concurrency = descriptor.getArgument("embed_concurrency"); - int threads = concurrency != null ? concurrency : 4; - this.embedPool = Executors.newFixedThreadPool(threads); + return concurrency != null ? concurrency : 4; + } + private static RetryExecutor createRetryExecutor(ResourceDescriptor descriptor) { Integer retries = descriptor.getArgument("max_retries"); - this.retryExecutor = - RetryExecutor.builder() - .maxRetries(retries != null ? retries : 5) - .initialBackoffMs(200) - .retryablePredicate(BedrockEmbeddingModelConnection::isRetryable) - .build(); + return RetryExecutor.builder() + .maxRetries(retries != null ? retries : 5) + .initialBackoffMs(200) + .retryablePredicate(BedrockEmbeddingModelConnection::isRetryable) + .build(); } @Override public float[] embed(String text, Map parameters) { + return embedWithUsage(text, parameters).getEmbeddings(); + } + + @Override + public EmbeddingResult embedWithUsage(String text, Map parameters) { String model = (String) parameters.getOrDefault("model", defaultModel); Integer dimensions = (Integer) parameters.get("dimensions"); @@ -138,12 +171,21 @@ public float[] embed(String text, Map parameters) { for (int i = 0; i < embeddingNode.size(); i++) { embedding[i] = (float) embeddingNode.get(i).asDouble(); } - return embedding; + return new EmbeddingResult<>(embedding, extractTokenUsage(result)); } catch (Exception e) { throw new RuntimeException("Failed to parse Bedrock embedding response.", e); } } + private static EmbeddingTokenUsage extractTokenUsage(JsonNode result) { + JsonNode inputTokenCount = result.get("inputTextTokenCount"); + if (inputTokenCount == null || !inputTokenCount.isNumber()) { + return null; + } + long tokens = inputTokenCount.asLong(); + return new EmbeddingTokenUsage(tokens, tokens); + } + private static boolean isRetryable(Exception e) { String msg = e.toString(); return msg.contains("ThrottlingException") @@ -155,27 +197,52 @@ private static boolean isRetryable(Exception e) { @Override public List embed(List texts, Map parameters) { + return embedWithUsage(texts, parameters).getEmbeddings(); + } + + @Override + public EmbeddingResult> embedWithUsage( + List texts, Map parameters) { if (texts.size() <= 1) { List results = new ArrayList<>(texts.size()); + EmbeddingTokenUsage totalUsage = null; for (String text : texts) { - results.add(embed(text, parameters)); + EmbeddingResult result = embedWithUsage(text, parameters); + results.add(result.getEmbeddings()); + totalUsage = mergeUsage(totalUsage, result.getTokenUsage()); } - return results; + return new EmbeddingResult<>(results, totalUsage); } @SuppressWarnings("unchecked") - CompletableFuture[] futures = + CompletableFuture>[] futures = texts.stream() .map( text -> CompletableFuture.supplyAsync( - () -> embed(text, parameters), embedPool)) + () -> embedWithUsage(text, parameters), embedPool)) .toArray(CompletableFuture[]::new); CompletableFuture.allOf(futures).join(); List results = new ArrayList<>(texts.size()); - for (CompletableFuture f : futures) { - results.add(f.join()); + EmbeddingTokenUsage totalUsage = null; + for (CompletableFuture> f : futures) { + EmbeddingResult result = f.join(); + results.add(result.getEmbeddings()); + totalUsage = mergeUsage(totalUsage, result.getTokenUsage()); + } + return new EmbeddingResult<>(results, totalUsage); + } + + private static EmbeddingTokenUsage mergeUsage( + EmbeddingTokenUsage left, EmbeddingTokenUsage right) { + if (left == null) { + return right; + } + if (right == null) { + return left; } - return results; + return new EmbeddingTokenUsage( + left.getPromptTokens() + right.getPromptTokens(), + left.getTotalTokens() + right.getTotalTokens()); } @Override diff --git a/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java b/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java index 0891d7067..99c8e9b81 100644 --- a/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java +++ b/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java @@ -18,17 +18,31 @@ package org.apache.flink.agents.integrations.embeddingmodels.bedrock; +import org.apache.flink.agents.api.RetryExecutor; import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection; import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; import org.apache.flink.agents.api.resource.ResourceContext; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient; +import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelRequest; +import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse; +import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; /** Tests for {@link BedrockEmbeddingModelConnection} and {@link BedrockEmbeddingModelSetup}. */ class BedrockEmbeddingModelTest { @@ -94,4 +108,42 @@ void testSetupParametersNoDimensions() { assertThat(setup.getParameters()).doesNotContainKey("dimensions"); } + + @Test + @DisplayName("Batch embedding aggregates token usage from worker results") + void testBatchEmbeddingAggregatesTokenUsage() throws Exception { + BedrockRuntimeClient client = mock(BedrockRuntimeClient.class); + when(client.invokeModel(any(InvokeModelRequest.class))) + .thenReturn( + InvokeModelResponse.builder() + .body( + SdkBytes.fromUtf8String( + "{\"embedding\":[0.1,0.2],\"inputTextTokenCount\":5}")) + .build()); + + ExecutorService embedPool = Executors.newFixedThreadPool(2); + BedrockEmbeddingModelConnection conn = + new BedrockEmbeddingModelConnection( + connDescriptor(null), + NOOP, + client, + embedPool, + RetryExecutor.builder().maxRetries(0).build(), + "mock-model"); + + try { + EmbeddingResult> result = + conn.embedWithUsage(List.of("first", "second"), Map.of("model", "mock-model")); + + assertThat(result.getEmbeddings()).hasSize(2); + assertThat(result.getEmbeddings().get(0)).containsExactly(0.1f, 0.2f); + assertThat(result.getEmbeddings().get(1)).containsExactly(0.1f, 0.2f); + assertThat(result.getTokenUsage()).isNotNull(); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(10L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(10L); + verify(client, times(2)).invokeModel(any(InvokeModelRequest.class)); + } finally { + conn.close(); + } + } } diff --git a/python/flink_agents/api/chat_models/chat_model.py b/python/flink_agents/api/chat_models/chat_model.py index ac4a814aa..d40f238ee 100644 --- a/python/flink_agents/api/chat_models/chat_model.py +++ b/python/flink_agents/api/chat_models/chat_model.py @@ -17,11 +17,13 @@ ################################################################################# import re from abc import ABC, abstractmethod +from enum import Enum from typing import Any, ClassVar, Dict, List, Mapping, Sequence, Tuple, cast -from pydantic import Field, PrivateAttr +from pydantic import Field, PrivateAttr, field_validator from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ( ChatMessage, MessageRole, @@ -33,6 +35,70 @@ from flink_agents.api.tools.tool import Tool +class StructuredOutputStrategy(str, Enum): + """User intent about how an output schema should be applied to a chat request. + + This expresses *policy* only. Whether a connection *can* apply the provider's + native structured-output API is a separate, model-dependent *capability* + question. Policy and capability are combined at request-build time. + + Inherits from ``str`` so the value survives the JSON-carried bridge to Java. + Java serializes this enum as its *name* ("NATIVE") while the value here is + lowercase, so ``_missing_`` accepts either form in any case — matching the + case-insensitive resolver on the Java side. + + Attributes: + ---------- + AUTO : str + Use the provider's native structured-output API when the effective model is + capable of it, and fall back to prompt engineering otherwise. The default. + NATIVE : str + Always use the provider's native structured-output API, without consulting + the capability predicate. + PROMPT : str + Never use the provider's native structured-output API; rely on prompt + engineering alone. Matches the behavior of connections that have no native + translation. + """ + + AUTO = "auto" + NATIVE = "native" + PROMPT = "prompt" + + def resolves_to_native(self, model_capable: bool) -> bool: # noqa: FBT001 + """Resolve this policy against a model's capability into whether to go native. + + ``AUTO`` defers to ``model_capable`` (native when the effective model can, else + the prompt-engineering fallback); ``NATIVE`` always resolves to native, ignoring + ``model_capable``, so an explicit user intent surfaces a provider error rather + than silently degrading; ``PROMPT`` never resolves to native. + + Parameters + ---------- + model_capable : bool + Whether the connection reports the effective model as natively capable. + + Returns: + ------- + bool + ``True`` if native structured output should be applied. + """ + if self is StructuredOutputStrategy.NATIVE: + return True + if self is StructuredOutputStrategy.PROMPT: + return False + return model_capable + + @classmethod + def _missing_(cls, value: object) -> "StructuredOutputStrategy | None": + if isinstance(value, str): + normalized = value.lower() + for member in cls: + if normalized in (member.value, member.name.lower()): + return member + return None + + class BaseChatModelConnection(Resource, ABC): """Base abstract class for chat model connection. @@ -54,6 +120,61 @@ def resource_type(cls) -> ResourceType: """Return resource type of class.""" return ResourceType.CHAT_MODEL_CONNECTION + def supports_native_structured_output(self, effective_model: str | None) -> bool: + """Whether this connection can natively structure output for a given model. + + Capability is *model-dependent*, not connection-wide: a single provider + connection commonly serves both models that accept a native schema parameter and + models that do not, so it is evaluated against the *effective* model at + request-build time — the model actually being called, which per-request + parameters may override. + + The default ``False`` keeps a connection on the prompt-engineering fallback. A + connection that translates a schema into a native provider parameter overrides + this; an unrecognized model must report ``False`` so it degrades to the fallback + rather than failing at the provider. + + Parameters + ---------- + effective_model : str | None + The model the request will be issued against, may be ``None``. + + Returns: + ------- + bool + ``True`` if a schema can be applied natively for ``effective_model``. + """ + return False + + def _reject_unsupported_output_schema( + self, output_schema: OutputSchema | None + ) -> None: + """Refuse an output schema this connection cannot translate natively. + + ``chat`` is abstract here, so there is no inherited body that could absorb a + schema loudly. A connection without a native structured-output translation + calls this as the first statement of its ``chat`` instead, which turns a + schema it could only drop into an error rather than an unconstrained response + that the caller would mistake for a schema-conforming one. + + Args: + output_schema: The schema the response should conform to. ``None`` returns + without effect. + + Raises: + NotImplementedError: If ``output_schema`` is not ``None``. + """ + if output_schema is None: + return + cls = type(self) + msg = ( + f"{cls.__module__}.{cls.__qualname__} has no native structured-output" + " translation, so it cannot honor the given output schema. Override chat()" + " to translate the schema natively, or pass no schema so the caller applies" + " the prompt-engineering fallback." + ) + raise NotImplementedError(msg) + DEFAULT_REASONING_PATTERNS: ClassVar[Tuple[re.Pattern[str], ...]] = ( re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), re.compile(r"(.*?)", re.DOTALL | re.IGNORECASE), @@ -106,6 +227,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -116,6 +238,15 @@ def chat( Input message sequence tools : Optional[List] List of tools that can be called by the model + output_schema : OutputSchema | None + The schema the response should conform to, or ``None`` for an + unconstrained response. This is framework-level execution metadata, and + every implementation must declare it as a named parameter rather than let + it fall into ``**kwargs``: ``**kwargs`` is forwarded to the provider SDK, + so a schema landing there would reach the request body. Implementations + without a native structured-output translation reject a non-``None`` value + via ``_reject_unsupported_output_schema``, so a caller that wants the + prompt-engineering fallback must pass ``None``. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -153,6 +284,30 @@ class BaseChatModelSetup(Resource): skill_discovery_prompt: str | None = None allowed_commands: List[str] = Field(default_factory=list) allowed_script_dirs: List[str] = Field(default_factory=list) + structured_output_strategy: StructuredOutputStrategy = Field( + default=StructuredOutputStrategy.AUTO, + description=( + "Intent about how an output schema should be applied. Whether native " + "structured output is actually used combines this policy with the " + "connection's model-dependent capability. An explicitly null value is " + "normalized to AUTO, so a validated setup always carries a real strategy." + ), + ) + + @field_validator("structured_output_strategy", mode="before") + @classmethod + def _normalize_null_strategy(cls, value: Any) -> Any: + """Normalize an explicitly null strategy to the ``AUTO`` default. + + A configuration source can carry the key with a null value instead of + omitting it. Java cannot tell those two apart — its descriptor argument + lookup returns null in both cases and resolves them to ``AUTO`` — so an + explicit null resolves to ``AUTO`` here too rather than being rejected. + Unknown non-null values still fail validation. + """ + if value is None: + return StructuredOutputStrategy.AUTO + return value @property @abstractmethod @@ -256,9 +411,8 @@ def chat( # Call chat model connection to execute chat merged_kwargs = self.model_kwargs.copy() merged_kwargs.update(kwargs) - return self._get_connection().chat( - messages, tools=self._get_tools(), **merged_kwargs - ) + connection = self._get_connection() + return connection.chat(messages, tools=self._get_tools(), **merged_kwargs) def _record_token_metrics( self, model_name: str, prompt_tokens: int, completion_tokens: int diff --git a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py index 921662360..462674bd7 100644 --- a/python/flink_agents/api/chat_models/tests/test_chat_model_base.py +++ b/python/flink_agents/api/chat_models/tests/test_chat_model_base.py @@ -18,12 +18,14 @@ from typing import Any, Dict, List, Sequence import pytest -from pydantic import Field, ValidationError +from pydantic import BaseModel, Field, ValidationError +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, BaseChatModelSetup, + StructuredOutputStrategy, ) from flink_agents.api.prompts.prompt import Prompt from flink_agents.api.tools.tool import Tool @@ -41,18 +43,29 @@ def model_kwargs(self) -> Dict[str, Any]: return {"model": self.model} +class _Answer(BaseModel): + """A representative BaseModel output schema.""" + + text: str + + class _RecordingConnection(BaseChatModelConnection): - """Connection that captures the messages it receives for inspection.""" + """Connection that captures the messages and kwargs it receives for inspection.""" captured_messages: List[ChatMessage] = Field(default_factory=list) + captured_kwargs: Dict[str, Any] = Field(default_factory=dict) + captured_output_schema: OutputSchema | None = None def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: self.captured_messages = list(messages) + self.captured_kwargs = dict(kwargs) + self.captured_output_schema = output_schema return ChatMessage(role=MessageRole.ASSISTANT, content="ok") @@ -125,3 +138,120 @@ def test_chat_refills_template_on_subsequent_invocations() -> None: assert len(connection.captured_messages) == 2 assert connection.captured_messages[0].content == "Task: v1" assert connection.captured_messages[1].content == "tool result" + + +def test_default_capability_predicate_is_false() -> None: + """A connection reports no native structured output for any model by default.""" + connection = _RecordingConnection() + + assert connection.supports_native_structured_output("gpt-4o") is False + assert connection.supports_native_structured_output("gpt-3.5-turbo") is False + assert connection.supports_native_structured_output(None) is False + + +def test_output_schema_guard_rejects_a_schema() -> None: + """The guard refuses a schema a connection cannot translate natively. + + Dropping it instead would return an unconstrained response that the caller has no + way to tell apart from a schema-conforming one. + """ + connection = _RecordingConnection() + schema = OutputSchema(output_schema=_Answer) + + with pytest.raises(NotImplementedError, match="_RecordingConnection"): + connection._reject_unsupported_output_schema(schema) + + +def test_output_schema_guard_passes_through_none() -> None: + """A caller on the prompt-engineering fallback passes None and is let through.""" + connection = _RecordingConnection() + + assert connection._reject_unsupported_output_schema(None) is None + + +def test_setup_routes_output_schema_through_to_connection() -> None: + """chat() forwards a caller's ``output_schema`` on to the connection intact. + + The setup filters what reaches the connection, so a schema it dropped or consumed + would leave the connection unable to apply one at all. That the schema cannot land + in ``**kwargs`` is a separate, tree-wide invariant covered by the connection + signature guard. + """ + setup = _RecordingChatModelSetup(connection="c", model="m") + connection = _RecordingConnection() + setup._resolved_connection = connection + + schema = OutputSchema(output_schema=_Answer) + setup.chat([], output_schema=schema) + + assert connection.captured_output_schema is schema + + +def test_structured_output_strategy_defaults_to_auto() -> None: + """The setup policy defaults to AUTO when unset.""" + setup = _RecordingChatModelSetup(connection="c", model="m") + + assert setup.structured_output_strategy is StructuredOutputStrategy.AUTO + + +@pytest.mark.parametrize("raw", ["NATIVE", "native", "Native"]) +def test_structured_output_strategy_coerces_name_and_value_case_insensitively( + raw: str, +) -> None: + """The policy coerces from either its name or its value, in any case. + + Java serializes this enum as its name ("NATIVE") and its own resolver accepts any + case, so a Python side that only accepted the lowercase value would reject what + Java sends. + """ + setup = _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=raw + ) + + assert setup.structured_output_strategy is StructuredOutputStrategy.NATIVE + + +def test_structured_output_strategy_normalizes_explicit_none_to_auto() -> None: + """An explicitly null policy resolves to AUTO instead of being rejected. + + Java cannot distinguish a configuration that carries the key as null from one + that omits it, and resolves both to AUTO, so a null arriving on Python must not + fail validation or leave the attribute None. + """ + setup = _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=None + ) + + assert setup.structured_output_strategy is StructuredOutputStrategy.AUTO + + +@pytest.mark.parametrize("raw", ["bogus", ""]) +def test_structured_output_strategy_rejects_unrecognized_value(raw: str) -> None: + """Only null normalizes to AUTO; every other unrecognized value still raises. + + An empty string reaches this field in practice from an empty YAML scalar or an + unset environment substitution, and it must not be mistaken for an omitted value: + normalizing on falsiness rather than on null would accept it as AUTO. A non-empty + unrecognized name survives that same falsiness check, so it takes `"bogus"` to + catch a resolver that coerces any unknown string to AUTO. + """ + with pytest.raises(ValidationError): + _RecordingChatModelSetup( + connection="c", model="m", structured_output_strategy=raw + ) + + +def test_auto_strategy_resolves_to_native_only_when_capable() -> None: + """AUTO defers to the model's capability.""" + assert StructuredOutputStrategy.AUTO.resolves_to_native(True) is True + assert StructuredOutputStrategy.AUTO.resolves_to_native(False) is False + + +def test_native_strategy_forces_native_regardless_of_capability() -> None: + """NATIVE resolves to native even when the model is not capable.""" + assert StructuredOutputStrategy.NATIVE.resolves_to_native(False) is True + + +def test_prompt_strategy_never_resolves_to_native() -> None: + """PROMPT never resolves to native even when the model is capable.""" + assert StructuredOutputStrategy.PROMPT.resolves_to_native(True) is False diff --git a/python/flink_agents/api/embedding_models/embedding_model.py b/python/flink_agents/api/embedding_models/embedding_model.py index 5529c817e..c31ab5324 100644 --- a/python/flink_agents/api/embedding_models/embedding_model.py +++ b/python/flink_agents/api/embedding_models/embedding_model.py @@ -16,13 +16,32 @@ # limitations under the License. ################################################################################# from abc import ABC, abstractmethod -from typing import Any, Dict, Sequence, cast +from dataclasses import dataclass +from typing import Any, Dict, Generic, Sequence, TypeVar, cast from pydantic import Field from typing_extensions import override from flink_agents.api.resource import Resource, ResourceType +EmbeddingValue = TypeVar("EmbeddingValue", list[float], list[list[float]]) + + +@dataclass(frozen=True) +class EmbeddingTokenUsage: + """Token usage reported by an embedding provider.""" + + prompt_tokens: int = 0 + total_tokens: int = 0 + + +@dataclass(frozen=True) +class EmbeddingResult(Generic[EmbeddingValue]): + """Embedding provider result with optional token usage metadata.""" + + embeddings: EmbeddingValue + token_usage: EmbeddingTokenUsage | None = None + class BaseEmbeddingModelConnection(Resource, ABC): """Base abstract class for text embedding model connection. @@ -62,6 +81,12 @@ def embed( The dimension of the vector depends on the specific embedding model used. """ + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings and return provider token usage when available.""" + return EmbeddingResult(embeddings=self.embed(text, **kwargs)) + class BaseEmbeddingModelSetup(Resource, ABC): """Base abstract class for text embedding model setup. @@ -123,3 +148,11 @@ def embed( merged_kwargs = self.model_kwargs.copy() merged_kwargs.update(kwargs) return self._get_connection().embed(text, **merged_kwargs) + + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings and return provider token usage when available.""" + merged_kwargs = self.model_kwargs.copy() + merged_kwargs.update(kwargs) + return self._get_connection().embed_with_usage(text, **merged_kwargs) diff --git a/python/flink_agents/api/embedding_models/tests/__init__.py b/python/flink_agents/api/embedding_models/tests/__init__.py new file mode 100644 index 000000000..1b9e7e00a --- /dev/null +++ b/python/flink_agents/api/embedding_models/tests/__init__.py @@ -0,0 +1,18 @@ +################################################################################ +# 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. +################################################################################# + diff --git a/python/flink_agents/api/embedding_models/tests/test_embedding_result.py b/python/flink_agents/api/embedding_models/tests/test_embedding_result.py new file mode 100644 index 000000000..e42193520 --- /dev/null +++ b/python/flink_agents/api/embedding_models/tests/test_embedding_result.py @@ -0,0 +1,127 @@ +################################################################################ +# 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, Dict, Sequence +from unittest.mock import MagicMock + +from flink_agents.api.embedding_models.embedding_model import ( + BaseEmbeddingModelConnection, + BaseEmbeddingModelSetup, + EmbeddingResult, + EmbeddingTokenUsage, +) +from flink_agents.api.resource import Resource, ResourceType +from flink_agents.api.resource_context import ResourceContext + + +class FakeEmbeddingModelConnection(BaseEmbeddingModelConnection): + def embed( + self, text: str | Sequence[str], **kwargs: Any + ) -> list[float] | list[list[float]]: + if isinstance(text, str): + return [0.1, 0.2] + return [[0.1, 0.2] for _ in text] + + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + return EmbeddingResult( + embeddings=self.embed(text, **kwargs), + token_usage=EmbeddingTokenUsage(prompt_tokens=7, total_tokens=9), + ) + + +class FakeEmbeddingModelConnectionWithoutUsage(BaseEmbeddingModelConnection): + def embed( + self, text: str | Sequence[str], **kwargs: Any + ) -> list[float] | list[list[float]]: + if isinstance(text, str): + return [0.1, 0.2] + return [[0.1, 0.2] for _ in text] + + +class DispatchAwareEmbeddingModelConnection(BaseEmbeddingModelConnection): + def embed( + self, text: str | Sequence[str], **kwargs: Any + ) -> list[float] | list[list[float]]: + if isinstance(text, str): + return [0.1, 0.2] + return [[0.1, 0.2] for _ in text] + + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + return EmbeddingResult( + embeddings=[0.3, 0.4] + if isinstance(text, str) + else [[0.3, 0.4] for _ in text], + token_usage=EmbeddingTokenUsage(prompt_tokens=7, total_tokens=9), + ) + + +class FakeEmbeddingModelSetup(BaseEmbeddingModelSetup): + @property + def model_kwargs(self) -> Dict[str, Any]: + return {} + + +def _make_setup(connection: BaseEmbeddingModelConnection) -> FakeEmbeddingModelSetup: + def get_resource(name: str, resource_type: ResourceType) -> Resource: + assert name == "mock-connection" + assert resource_type == ResourceType.EMBEDDING_MODEL_CONNECTION + return connection + + ctx = MagicMock(spec=ResourceContext) + ctx.get_resource = get_resource + setup = FakeEmbeddingModelSetup( + name="embedding", + connection="mock-connection", + model="mock-model", + resource_context=ctx, + ) + setup.open() + return setup + + +def test_embedding_result_returns_provider_usage() -> None: + setup = _make_setup(FakeEmbeddingModelConnection(name="connection")) + + result = setup.embed_with_usage("hello") + + assert result.embeddings == [0.1, 0.2] + assert result.token_usage == EmbeddingTokenUsage(prompt_tokens=7, total_tokens=9) + + +def test_embedding_result_defaults_to_no_usage() -> None: + setup = _make_setup(FakeEmbeddingModelConnectionWithoutUsage(name="connection")) + + result = setup.embed_with_usage("hello") + + assert result.embeddings == [0.1, 0.2] + assert result.token_usage is None + + +def test_embed_preserves_existing_embedding_only_api() -> None: + setup = _make_setup(FakeEmbeddingModelConnection(name="connection")) + + assert setup.embed("hello") == [0.1, 0.2] + + +def test_embed_preserves_connection_embed_dispatch() -> None: + setup = _make_setup(DispatchAwareEmbeddingModelConnection(name="connection")) + + assert setup.embed("hello") == [0.1, 0.2] diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py index 625bb69f5..2aa03c264 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/mock_chat_model_agent.py @@ -29,6 +29,7 @@ from pyflink.datastream import KeySelector from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -95,9 +96,16 @@ def chat( self, messages: Sequence[ChatMessage], tools: List | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Generate a tool call or a response according to input.""" + """Generate a tool call or a response according to input. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + structured-output translation. Declaring the parameter keeps a caller-supplied + schema out of ``**kwargs``. + """ + self._reject_unsupported_output_schema(output_schema) if "sum" in messages[-1].content: input = messages[-1].content # Validate the tool was bound before the model was invoked. diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py index 0b30f9afb..9358fec9d 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_integration/tool_parameter_injection_agent.py @@ -18,6 +18,7 @@ from pyflink.datastream import KeySelector from flink_agents.api.agents.agent import Agent +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -113,9 +114,16 @@ def chat( self, messages: list[ChatMessage], tools: list[BaseTool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: object, ) -> ChatMessage: - """Return a tool call for user input, or echo the tool response.""" + """Return a tool call for user input, or echo the tool response. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + structured-output translation. Declaring the parameter keeps a caller-supplied + schema out of ``**kwargs``. + """ + self._reject_unsupported_output_schema(output_schema) last_message = messages[-1] if last_message.role == MessageRole.TOOL: return ChatMessage(role=MessageRole.ASSISTANT, content=last_message.content) diff --git a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_agent.py b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_agent.py index dbce4891b..e57fbc07f 100644 --- a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_agent.py +++ b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/embedding_model_cross_language_agent.py @@ -77,7 +77,8 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ) # Test single text embedding - embedding = embeddingModel.embed(input_text) + embedding_result = embeddingModel.embed_with_usage(input_text) + embedding = embedding_result.embeddings print(f"[TEST] Generated embedding with dimension: {len(embedding)}") # Validate single embedding result @@ -98,7 +99,8 @@ def process_input(event: Event, ctx: RunnerContext) -> None: ) # Test batch embedding - embeddings = embeddingModel.embed([input_text]) + embeddings_result = embeddingModel.embed_with_usage([input_text]) + embeddings = embeddings_result.embeddings print(f"[TEST] Generated batch embeddings: count={len(embeddings)}") # Validate batch embedding results diff --git a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py index af5483dc1..afcd3f124 100644 --- a/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py +++ b/python/flink_agents/integrations/chat_models/anthropic/anthropic_chat_model.py @@ -24,6 +24,7 @@ from pydantic import Field, PrivateAttr from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -168,9 +169,17 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Direct communication with Anthropic model service for chat conversation.""" + """Direct communication with Anthropic model service for chat conversation. + + A non-``None`` ``output_schema`` is rejected: 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. + """ + self._reject_unsupported_output_schema(output_schema) anthropic_tools = None if tools is not None: anthropic_tools = [ 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 18a092128..b653b4d67 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 @@ -21,6 +21,7 @@ from openai import NOT_GIVEN, AzureOpenAI from pydantic import Field, PrivateAttr +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -117,6 +118,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -127,6 +129,11 @@ def chat( Input message sequence 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -136,6 +143,7 @@ def chat( 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] diff --git a/python/flink_agents/integrations/chat_models/ollama_chat_model.py b/python/flink_agents/integrations/chat_models/ollama_chat_model.py index 7c36ec38e..b0061fb8d 100644 --- a/python/flink_agents/integrations/chat_models/ollama_chat_model.py +++ b/python/flink_agents/integrations/chat_models/ollama_chat_model.py @@ -21,6 +21,7 @@ from ollama import Client, Message from pydantic import Field +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -85,9 +86,17 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Process a sequence of messages, and return a response.""" + """Process a sequence of messages, and return a response. + + A non-``None`` ``output_schema`` is rejected: 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. + """ + self._reject_unsupported_output_schema(output_schema) ollama_messages = self.__convert_to_ollama_messages(messages) # Convert tool format diff --git a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py index 0d487772b..68c2414ba 100644 --- a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py +++ b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py @@ -22,6 +22,7 @@ from pydantic import Field, PrivateAttr from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -138,6 +139,7 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: """Direct communication with model service for chat conversation. @@ -148,6 +150,11 @@ def chat( Input message sequence 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. **kwargs : Any Additional parameters passed to the model service (e.g., temperature, max_tokens, etc.) @@ -157,6 +164,7 @@ def chat( 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] diff --git a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py new file mode 100644 index 000000000..e31d37535 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py @@ -0,0 +1,168 @@ +################################################################################ +# 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. +################################################################################# +"""Guards the ``output_schema`` parameter contract across every chat connection.""" + +import inspect +from typing import Iterator, List, Type + +from pydantic import BaseModel + +from flink_agents.api.agents.types import OutputSchema +from flink_agents.api.chat_models import java_chat_model as api_java_chat_model +from flink_agents.api.chat_models.chat_model import BaseChatModelConnection +from flink_agents.e2e_tests.e2e_tests_integration import ( + mock_chat_model_agent, + tool_parameter_injection_agent, +) +from flink_agents.integrations.chat_models import ollama_chat_model, tongyi_chat_model +from flink_agents.integrations.chat_models.anthropic import anthropic_chat_model +from flink_agents.integrations.chat_models.azure import azure_openai_chat_model +from flink_agents.integrations.chat_models.openai import openai_chat_model +from flink_agents.runtime.java import java_chat_model as runtime_java_chat_model + +# A class is only discoverable through __subclasses__() once it has been imported. +# Importing every module that defines a connection is what gives the walk below its +# reach — including the cross-language bridge and the e2e test doubles. +_MODULES_DEFINING_CONNECTIONS = ( + anthropic_chat_model, + api_java_chat_model, + azure_openai_chat_model, + mock_chat_model_agent, + ollama_chat_model, + openai_chat_model, + runtime_java_chat_model, + tongyi_chat_model, + tool_parameter_injection_agent, +) + + +class _Answer(BaseModel): + """A representative output schema to hand to a connection.""" + + text: str + + +def _subclasses_recursive(cls: Type[object]) -> Iterator[Type[object]]: + for subclass in cls.__subclasses__(): + yield subclass + yield from _subclasses_recursive(subclass) + + +def _is_test_local(cls: Type[object]) -> bool: + """Whether ``cls`` is defined inside a ``tests`` package. + + A whole-tree pytest run imports every test module into one process, so doubles + declared inside a test file also turn up in the walk. Such a double may + deliberately accept a schema in order to assert on what a caller routed to it, + which is the opposite of the contract asserted here. + """ + return "tests" in cls.__module__.split(".") + + +def _connections_implementing_chat() -> List[Type[BaseChatModelConnection]]: + """Connections outside a ``tests`` package that supply their own ``chat`` body. + + A class that never overrides ``chat`` inherits only the abstract declaration, so + there is no body to assert on. + """ + return [ + cls + for cls in _subclasses_recursive(BaseChatModelConnection) + if not _is_test_local(cls) and cls.chat is not BaseChatModelConnection.chat + ] + + +def test_every_connection_declares_output_schema_param() -> None: + """Every connection in the tree must declare ``output_schema`` defaulting to None. + + A connection that omits the parameter absorbs a caller's ``output_schema`` into + ``**kwargs``. Connections forward ``**kwargs`` to their provider SDK — and the + cross-language bridge forwards it to Java as ``modelParams`` — so the schema would + reach the request body as an unknown field. Declaring the parameter is what makes + the leak impossible rather than merely unlikely. + + The tree is walked rather than hand-listed: a hand-kept list silently stops + guarding whatever it was never updated to mention, including the abstract base + itself. + """ + offenders: List[str] = [] + connections = [ + BaseChatModelConnection, + *_subclasses_recursive(BaseChatModelConnection), + ] + for cls in connections: + param = inspect.signature(cls.chat).parameters.get("output_schema") + if param is None: + offenders.append(f"{cls.__module__}.{cls.__qualname__}: parameter missing") + elif param.default is not None: + offenders.append( + f"{cls.__module__}.{cls.__qualname__}: default is " + f"{param.default!r}, expected None" + ) + + assert not offenders, ( + "connections violating the output_schema contract:\n" + "\n".join(offenders) + ) + + +def _schema_rejection_failure( + cls: Type[BaseChatModelConnection], schema: OutputSchema +) -> str | None: + """Describe how ``cls.chat`` mishandles ``schema``, or ``None`` if it rejects it. + + ``__new__`` skips ``__init__``, so no credentials, no client and no network are + involved: the rejection has to happen before the first attribute access for the + call to get this far, which is what pins the guard to the top of ``chat``. + """ + name = f"{cls.__module__}.{cls.__qualname__}" + connection = cls.__new__(cls) + try: + cls.chat(connection, messages=[], tools=None, output_schema=schema) + except NotImplementedError: + return None + except Exception as exc: + return f"{name}: raised {type(exc).__name__} before rejecting the schema" + return f"{name}: accepted the schema instead of rejecting it" + + +def test_every_connection_rejects_an_output_schema_it_cannot_translate() -> None: + """A connection with no native translation must reject a schema, not drop it. + + Declaring the parameter only keeps the schema out of the provider request; on its + own it lets a connection silently return an unconstrained response that the caller + would treat as schema-conforming. Rejecting turns that into an error at the call. + + The tree is walked rather than hand-listed, so a connection added to a module that + is already imported is held to the contract for free. Reach still stops at those + imports: ``__subclasses__()`` only sees classes that have been imported, so a + connection in a brand-new module needs that module added at the top of this file. + """ + schema = OutputSchema(output_schema=_Answer) + connections = _connections_implementing_chat() + + assert connections, "the connection walk found nothing to check" + offenders = [ + failure + for cls in connections + if (failure := _schema_rejection_failure(cls, schema)) is not None + ] + + assert not offenders, ( + "connections that do not reject an untranslatable output_schema:\n" + + "\n".join(offenders) + ) diff --git a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py index 6587a8cba..5e7b52a1c 100644 --- a/python/flink_agents/integrations/chat_models/tongyi_chat_model.py +++ b/python/flink_agents/integrations/chat_models/tongyi_chat_model.py @@ -24,6 +24,7 @@ from dashscope import Generation from pydantic import Field +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.api.chat_models.chat_model import ( BaseChatModelConnection, @@ -101,9 +102,17 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Process a sequence of messages, and return a response.""" + """Process a sequence of messages, and return a response. + + A non-``None`` ``output_schema`` is rejected: 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. + """ + self._reject_unsupported_output_schema(output_schema) tongyi_messages = self.__convert_to_tongyi_messages(messages) tongyi_tools: List[Dict[str, Any]] | None = ( diff --git a/python/flink_agents/integrations/embedding_models/openai_embedding_model.py b/python/flink_agents/integrations/embedding_models/openai_embedding_model.py index eac4c8f10..92d09b765 100644 --- a/python/flink_agents/integrations/embedding_models/openai_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/openai_embedding_model.py @@ -24,6 +24,8 @@ from flink_agents.api.embedding_models.embedding_model import ( BaseEmbeddingModelConnection, BaseEmbeddingModelSetup, + EmbeddingResult, + EmbeddingTokenUsage, ) DEFAULT_REQUEST_TIMEOUT = 30.0 @@ -115,6 +117,12 @@ def embed( self, text: str | Sequence[str], **kwargs: Any ) -> list[float] | list[list[float]]: """Generate embedding vector for a single text query.""" + return self.embed_with_usage(text, **kwargs).embeddings + + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings and return OpenAI token usage when available.""" # Extract OpenAI specific parameters model = kwargs.pop("model") encoding_format = kwargs.pop("encoding_format", None) @@ -132,8 +140,24 @@ def embed( user=user if user is not None else NOT_GIVEN, ) + usage = getattr(response, "usage", None) + token_usage = None + if usage is not None: + prompt_tokens = getattr(usage, "prompt_tokens", None) + total_tokens = getattr(usage, "total_tokens", None) + if prompt_tokens is not None or total_tokens is not None: + token_usage = EmbeddingTokenUsage( + prompt_tokens=int(prompt_tokens or 0), + total_tokens=int( + total_tokens if total_tokens is not None else prompt_tokens + ), + ) + embeddings = [list(embedding.embedding) for embedding in response.data] - return embeddings[0] if isinstance(text, str) else embeddings + return EmbeddingResult( + embeddings=embeddings[0] if isinstance(text, str) else embeddings, + token_usage=token_usage, + ) @override def close(self) -> None: diff --git a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py index 49907340f..9d63c2494 100644 --- a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################ import os +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -57,3 +58,34 @@ def get_resource(name: str, type: ResourceType) -> Resource: assert isinstance(response, list) assert len(response) > 0 assert all(isinstance(x, float) for x in response) # + + +def test_openai_embedding_model_returns_token_usage() -> None: + """Test OpenAI embedding usage is returned with the embedding result.""" + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + mock_client = MagicMock() + mock_client.embeddings.create.return_value = SimpleNamespace( + data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3])], + usage=SimpleNamespace(prompt_tokens=5, total_tokens=5), + ) + connection._OpenAIEmbeddingModelConnection__client = mock_client + + def get_resource(name: str, type: ResourceType) -> Resource: + if type == ResourceType.EMBEDDING_MODEL_CONNECTION: + return connection + else: + msg = f"Unknown resource type: {type}" + raise ValueError(msg) + + mock_ctx = MagicMock(spec=ResourceContext) + mock_ctx.get_resource = get_resource + embedding_model = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", resource_context=mock_ctx + ) + embedding_model.open() + + result = embedding_model.embed_with_usage("Hello, Flink Agent!") + assert result.embeddings == [0.1, 0.2, 0.3] + assert result.token_usage is not None + assert result.token_usage.prompt_tokens == 5 + assert result.token_usage.total_tokens == 5 diff --git a/python/flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py b/python/flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py index b60c75596..6262d4bd6 100644 --- a/python/flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/tests/test_tongyi_embedding_model.py @@ -155,6 +155,52 @@ def get_resource(name: str, type: ResourceType) -> Resource: assert len(response) == 5 +def test_tongyi_embedding_returns_token_usage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test DashScope embedding usage is recorded as model token metrics.""" + mock_embedding = [0.1, 0.2, 0.3] + mocked_response = SimpleNamespace( + status_code=HTTPStatus.OK, + output={ + "embeddings": [{"embedding": mock_embedding}], + }, + usage={"total_tokens": 6}, + message="Success", + ) + mock_call = MagicMock(return_value=mocked_response) + monkeypatch.setattr( + "flink_agents.integrations.embedding_models.tongyi_embedding_model.dashscope.TextEmbedding.call", + mock_call, + ) + + connection = TongyiEmbeddingModelConnection( + name="tongyi", + api_key="fake-key", + ) + + def get_resource(name: str, type: ResourceType) -> Resource: + if type == ResourceType.EMBEDDING_MODEL_CONNECTION: + return connection + else: + msg = f"Unknown resource type: {type}" + raise ValueError(msg) + + embedding_model = TongyiEmbeddingModelSetup( + name="tongyi", + model=test_model, + connection="tongyi", + resource_context=_make_ctx(get_resource), + ) + embedding_model.open() + + result = embedding_model.embed_with_usage("Test text") + assert result.embeddings == mock_embedding + assert result.token_usage is not None + assert result.token_usage.prompt_tokens == 6 + assert result.token_usage.total_tokens == 6 + + def test_tongyi_embedding_batch_mock(monkeypatch: pytest.MonkeyPatch) -> None: """Test batch embedding functionality with mocked DashScope API.""" mock_embeddings = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] diff --git a/python/flink_agents/integrations/embedding_models/tongyi_embedding_model.py b/python/flink_agents/integrations/embedding_models/tongyi_embedding_model.py index c63c083d4..bea3075bd 100644 --- a/python/flink_agents/integrations/embedding_models/tongyi_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/tongyi_embedding_model.py @@ -25,12 +25,27 @@ from flink_agents.api.embedding_models.embedding_model import ( BaseEmbeddingModelConnection, BaseEmbeddingModelSetup, + EmbeddingResult, + EmbeddingTokenUsage, ) DEFAULT_REQUEST_TIMEOUT = 30.0 DEFAULT_MODEL = "text-embedding-v4" +def _get_usage_value(obj: Any, *names: str) -> int | None: + """Read a token usage value from dict-like or object-like provider responses.""" + if obj is None: + return None + for name in names: + if isinstance(obj, dict) and name in obj: + return int(obj[name]) + value = getattr(obj, name, None) + if value is not None: + return int(value) + return None + + class TongyiEmbeddingModelConnection(BaseEmbeddingModelConnection): """Tongyi Embedding Model Connection which manages connection to DashScope API. @@ -78,6 +93,12 @@ def embed( self, text: str | Sequence[str], **kwargs: Any ) -> list[float] | list[list[float]]: """Generate embedding vector for text input.""" + return self.embed_with_usage(text, **kwargs).embeddings + + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings and return DashScope token usage when available.""" model = kwargs.pop("model", DEFAULT_MODEL) text_type = kwargs.pop("text_type", None) dimension = kwargs.pop("dimension", None) @@ -103,8 +124,27 @@ def embed( msg = f"DashScope TextEmbedding call failed: {response.message}" raise RuntimeError(msg) + usage = getattr(response, "usage", None) + if usage is None and isinstance(response.output, dict): + usage = response.output.get("usage") + prompt_tokens = _get_usage_value(usage, "input_tokens", "prompt_tokens") + total_tokens = _get_usage_value(usage, "total_tokens") + if prompt_tokens is None: + prompt_tokens = total_tokens + token_usage = None + if prompt_tokens is not None or total_tokens is not None: + token_usage = EmbeddingTokenUsage( + prompt_tokens=int(prompt_tokens or 0), + total_tokens=int( + total_tokens if total_tokens is not None else prompt_tokens + ), + ) + embeddings = [e["embedding"] for e in response.output["embeddings"]] - return embeddings[0] if isinstance(text, str) else embeddings + return EmbeddingResult( + embeddings=embeddings[0] if isinstance(text, str) else embeddings, + token_usage=token_usage, + ) class TongyiEmbeddingModelSetup(BaseEmbeddingModelSetup): diff --git a/python/flink_agents/runtime/java/java_chat_model.py b/python/flink_agents/runtime/java/java_chat_model.py index 3f385a9ad..15df31578 100644 --- a/python/flink_agents/runtime/java/java_chat_model.py +++ b/python/flink_agents/runtime/java/java_chat_model.py @@ -19,6 +19,7 @@ from typing_extensions import override +from flink_agents.api.agents.types import OutputSchema from flink_agents.api.chat_message import ChatMessage from flink_agents.api.chat_models.java_chat_model import ( JavaChatModelConnection, @@ -64,14 +65,22 @@ def chat( self, messages: Sequence[ChatMessage], tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, **kwargs: Any, ) -> ChatMessage: - """Chat method that throws UnsupportedOperationException. + """Chat by forwarding the request to the wrapped Java connection. This connection serves as a Java resource wrapper only. Chat operations should be performed on the Java side using the underlying Java chat model object. + + A non-``None`` ``output_schema`` is rejected: only messages, tools and kwargs + cross to the Java three-argument ``chat``, so a schema forwarded here would be + dropped on the way. Declaring the parameter keeps a caller-supplied schema out + of ``**kwargs``, which crosses to Java as ``modelParams`` — a provider-facing + map, not a channel for framework execution metadata. """ + self._reject_unsupported_output_schema(output_schema) java_messages = [ self._j_resource_adapter.fromPythonChatMessage(message) for message in messages diff --git a/python/flink_agents/runtime/java/java_embedding_model.py b/python/flink_agents/runtime/java/java_embedding_model.py index b2ea48724..230a49441 100644 --- a/python/flink_agents/runtime/java/java_embedding_model.py +++ b/python/flink_agents/runtime/java/java_embedding_model.py @@ -19,6 +19,10 @@ from typing_extensions import override +from flink_agents.api.embedding_models.embedding_model import ( + EmbeddingResult, + EmbeddingTokenUsage, +) from flink_agents.api.embedding_models.java_embedding_model import ( JavaEmbeddingModelConnection, JavaEmbeddingModelSetup, @@ -28,6 +32,28 @@ ) +def _from_java_embedding_result( + j_result: Any, text: str | Sequence[str] +) -> EmbeddingResult[list[float] | list[list[float]]]: + """Convert a Java EmbeddingResult into the Python result contract.""" + j_embeddings = j_result.getEmbeddings() + embeddings = ( + list(j_embeddings) + if isinstance(text, str) + else [list(embedding) for embedding in j_embeddings] + ) + j_usage = j_result.getTokenUsage() + token_usage = ( + None + if j_usage is None + else EmbeddingTokenUsage( + prompt_tokens=j_usage.getPromptTokens(), + total_tokens=j_usage.getTotalTokens(), + ) + ) + return EmbeddingResult(embeddings=embeddings, token_usage=token_usage) + + class JavaEmbeddingModelConnectionImpl(JavaEmbeddingModelConnection): """Java-based implementation of EmbeddingModelConnection that wraps a Java embedding model object. @@ -72,6 +98,16 @@ def embed( ) return list(result) if isinstance(text, str) else [list(emb) for emb in result] + @override + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings through Java and preserve provider token usage.""" + result = self._j_resource.embedWithUsage( + text if isinstance(text, str) else list(text), kwargs + ) + return _from_java_embedding_result(result, text) + class JavaEmbeddingModelSetupImpl(JavaEmbeddingModelSetup): """Java-based implementation of EmbeddingModelSetup that wraps a Java embedding @@ -134,3 +170,13 @@ def embed( text if isinstance(text, str) else list(text), kwargs ) return list(result) if isinstance(text, str) else [list(emb) for emb in result] + + @override + def embed_with_usage( + self, text: str | Sequence[str], **kwargs: Any + ) -> EmbeddingResult[list[float] | list[list[float]]]: + """Generate embeddings through Java and preserve provider token usage.""" + result = self._j_resource.embedWithUsage( + text if isinstance(text, str) else list(text), kwargs + ) + return _from_java_embedding_result(result, text) diff --git a/python/flink_agents/runtime/python_java_utils.py b/python/flink_agents/runtime/python_java_utils.py index 66b9092f7..3c4bb60dc 100644 --- a/python/flink_agents/runtime/python_java_utils.py +++ b/python/flink_agents/runtime/python_java_utils.py @@ -154,7 +154,9 @@ def get_python_tool_metadata( descriptor = PythonFunction(module=module, qualname=qual_name) callable_ = descriptor.as_callable() name = callable_.__name__ - description = (parse(callable_.__doc__).description or "") if callable_.__doc__ else "" + description = ( + (parse(callable_.__doc__).description or "") if callable_.__doc__ else "" + ) callable_injected_args = normalize_injected_args( getattr(callable_, "_injected_args", None) ) @@ -177,9 +179,7 @@ def _dump_injected_args(injected_args: Dict[str, Any]) -> str: ) -def invoke_python_tool( - module: str, qual_name: str, kwargs: Dict[str, Any] -) -> Any: +def invoke_python_tool(module: str, qual_name: str, kwargs: Dict[str, Any]) -> Any: """Invoke a Python callable as a tool, passing the provided keyword arguments. Used by the Java-side ``PythonResourceAdapter.invokePythonTool`` so a Java host can @@ -226,6 +226,32 @@ def from_java_resource(type_name: str, kwargs: Dict[str, Any]) -> Resource: return cls(**kwargs) +def embedding_result_to_java(result: Any) -> Dict[str, Any]: + """Convert an embedding result into Pemja-safe Python primitives.""" + usage = result.token_usage + return { + "embeddings": result.embeddings, + "token_usage": ( + None + if usage is None + else { + "prompt_tokens": usage.prompt_tokens, + "total_tokens": usage.total_tokens, + } + ), + } + + +def call_embedding_with_usage( + embedding_model: Any, kwargs: Dict[str, Any] +) -> Dict[str, Any]: + """Call ``embed_with_usage`` and return Java-safe primitives. + + This avoids PyObject attribute access in the Java caller. + """ + return embedding_result_to_java(embedding_model.embed_with_usage(**kwargs)) + + def normalize_tool_call_id(tool_call: Dict[str, Any]) -> Dict[str, Any]: """Normalize tool call by converting the ID field to string format while preserving all other fields. diff --git a/python/flink_agents/runtime/tests/test_java_embedding_model.py b/python/flink_agents/runtime/tests/test_java_embedding_model.py new file mode 100644 index 000000000..2a3866687 --- /dev/null +++ b/python/flink_agents/runtime/tests/test_java_embedding_model.py @@ -0,0 +1,70 @@ +################################################################################ +# 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 unittest.mock import MagicMock + +import pytest + +from flink_agents.runtime.java.java_embedding_model import ( + JavaEmbeddingModelConnectionImpl, + JavaEmbeddingModelSetupImpl, +) + + +class _JavaTokenUsage: + def getPromptTokens(self) -> int: + return 7 + + def getTotalTokens(self) -> int: + return 9 + + +class _JavaEmbeddingResult: + def getEmbeddings(self) -> list[list[float]]: + return [[0.1, 0.2], [0.3, 0.4]] + + def getTokenUsage(self) -> _JavaTokenUsage: + return _JavaTokenUsage() + + +@pytest.mark.parametrize( + ("wrapper_class", "kwargs"), + [ + (JavaEmbeddingModelConnectionImpl, {}), + ( + JavaEmbeddingModelSetupImpl, + {"connection": "connection", "model": "test-model"}, + ), + ], +) +def test_java_embedding_wrappers_preserve_usage( + wrapper_class: type[JavaEmbeddingModelConnectionImpl | JavaEmbeddingModelSetupImpl], + kwargs: dict[str, str], +) -> None: + j_resource = MagicMock() + j_resource.embedWithUsage.return_value = _JavaEmbeddingResult() + wrapper = wrapper_class(j_resource, MagicMock(), **kwargs) + + result = wrapper.embed_with_usage(["first", "second"], batch_size=2) + + assert result.embeddings == [[0.1, 0.2], [0.3, 0.4]] + assert result.token_usage is not None + assert result.token_usage.prompt_tokens == 7 + assert result.token_usage.total_tokens == 9 + j_resource.embedWithUsage.assert_called_once_with( + ["first", "second"], {"batch_size": 2} + ) diff --git a/python/flink_agents/runtime/tests/test_python_java_utils.py b/python/flink_agents/runtime/tests/test_python_java_utils.py index 7aeeafe9c..fa708353b 100644 --- a/python/flink_agents/runtime/tests/test_python_java_utils.py +++ b/python/flink_agents/runtime/tests/test_python_java_utils.py @@ -18,8 +18,15 @@ import json from flink_agents.api.decorators import tool +from flink_agents.api.embedding_models.embedding_model import ( + EmbeddingResult, + EmbeddingTokenUsage, +) from flink_agents.api.tools import InjectedArg -from flink_agents.runtime.python_java_utils import get_python_tool_metadata +from flink_agents.runtime.python_java_utils import ( + call_embedding_with_usage, + get_python_tool_metadata, +) @tool(injected_args={"tenant_id": InjectedArg.from_config("tenant.id")}) @@ -36,6 +43,27 @@ def test_get_python_tool_metadata_merges_callable_injected_args() -> None: schema = json.loads(flat["inputSchema"]) assert set(schema["properties"]) == {"order_id"} injected_args = json.loads(flat["injectedArgs"]) - assert injected_args == { - "tenant_id": {"source": "config", "key": "tenant.id"} + assert injected_args == {"tenant_id": {"source": "config", "key": "tenant.id"}} + + +class _UsageAwareEmbeddingModel: + def embed_with_usage( + self, text: str, **kwargs: object + ) -> EmbeddingResult[list[float]]: + assert text == "hello" + assert kwargs == {"model": "test-model"} + return EmbeddingResult( + embeddings=[0.1, 0.2], + token_usage=EmbeddingTokenUsage(prompt_tokens=7, total_tokens=9), + ) + + +def test_call_embedding_with_usage_returns_pemja_safe_primitives() -> None: + result = call_embedding_with_usage( + _UsageAwareEmbeddingModel(), {"text": "hello", "model": "test-model"} + ) + + assert result == { + "embeddings": [0.1, 0.2], + "token_usage": {"prompt_tokens": 7, "total_tokens": 9}, } diff --git a/review-guides/runtime-state-recovery.md b/review-guides/runtime-state-recovery.md new file mode 100644 index 000000000..f03c53541 --- /dev/null +++ b/review-guides/runtime-state-recovery.md @@ -0,0 +1,31 @@ +# Review Guide: runtime/ State and Recovery + +Load this guide when a PR changes runtime state, checkpointing, serialization of +stored values, action-state handling, or async, mailbox, and recovery paths. It +narrows the full passes in `code_review.md` to the ones that matter most for this +area; the general passes still apply. + +## Focused checklist + +- Trace every stored value through its full serialize, checkpoint, and restore + round-trip. Confirm the restored value keeps its original type instead of a + lossy or re-encoded form, such as a `byte[]` that returns as a base64 `String`. +- Confirm a serde envelope or version tag scopes the value the way recovery + needs it, not only the way the happy-path write needs it. +- Under the beta policy, question serde or state compatibility flags that only + exist to read old formats. Prefer removing them unless they carry a concrete + obligation. +- For a bugfix, check the failure actually stops reaching the old path after the + fix, not just in the one code path the PR touched. +- Test the real runtime failure path through the operator, mailbox, or action + task, not a helper in isolation. Make failure-injection fixtures deterministic + so the regression test exercises the bug on every run. +- Cover async, retry, delayed-callback, and recovery paths when the change + affects them. + +## Examples from past reviews + +| Case | Pass it exercises | Review | +|---|---|---| +| A versioned serde envelope for stored `MemoryUpdate` values carried an `isEnvelope` backward-compatibility flag. The review asked whether that compatibility path was essential or should be dropped under the beta policy. | Whether a serde compatibility channel must exist, or only keeps an old format alive. | [#874](https://github.com/apache/flink-agents/pull/874#discussion_r3541046130) | +| A hotfix propagated a swallowed action-task `Error` out of the mailbox instead of losing it in a `Future`. The review kept the linkage-error regression fixture deterministic, since `System.nanoTime()` may return a value that skips the injected failure, so the test exercises the real failure path on every run. | Testing the actual runtime failure path with a deterministic fixture. | [#880](https://github.com/apache/flink-agents/pull/880#discussion_r3540383941) | diff --git a/tools/.rat-excludes b/tools/.rat-excludes index dcd7de725..94dacb5f4 100644 --- a/tools/.rat-excludes +++ b/tools/.rat-excludes @@ -25,4 +25,5 @@ skills/* .*\.yaml$ AGENTS.md code_review.md -CLAUDE.md \ No newline at end of file +CLAUDE.md +review-guides/* \ No newline at end of file