From 1430e3a553fcb8107109f1d232cf8873220bed14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=A5=E6=A0=96?= Date: Mon, 7 Sep 2026 01:48:16 +0800 Subject: [PATCH 1/2] [FLINK-40572][runtime] Support dynamic AI model selection --- docs/content.zh/docs/core-concept/ai-model.md | 12 +- docs/content/docs/core-concept/ai-model.md | 12 +- .../flink/translator/TransformTranslator.java | 49 ---- .../flink/FlinkPipelineAiFunctionITCase.java | 14 ++ .../cdc/runtime/ai/AiModelClientResolver.java | 42 ++++ .../runtime/functions/impl/AiFunctions.java | 164 ++++++++++--- .../transform/ProjectionColumnProcessor.java | 11 +- .../TransformExpressionCompiler.java | 25 +- .../transform/TransformFilterProcessor.java | 14 +- .../cdc/runtime/parser/JaninoCompiler.java | 20 +- .../cdc/runtime/parser/TransformParser.java | 209 ---------------- .../functions/impl/AiFunctionsTest.java | 96 ++++++-- .../runtime/parser/AiFunctionParserTest.java | 225 +++--------------- 13 files changed, 332 insertions(+), 561 deletions(-) create mode 100644 flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java diff --git a/docs/content.zh/docs/core-concept/ai-model.md b/docs/content.zh/docs/core-concept/ai-model.md index d809b3224a9..70e016be914 100644 --- a/docs/content.zh/docs/core-concept/ai-model.md +++ b/docs/content.zh/docs/core-concept/ai-model.md @@ -28,7 +28,17 @@ AI 模型可用于 transform 表达式中的文本生成、文本分析、embedd ## AI Functions -模型名称必须是字符串常量，并引用 `pipeline.model` 中声明的模型。文本、embedding 和图片函数分别要求模型客户端实现对应的 capability；Pipeline 会在执行前校验引用模型的 capability 是否匹配。 +模型参数可以是任意 `STRING` 表达式，并会针对每条记录求值。因此可以使用字段、`IF` 或 `CASE` 动态选择模型。仅在实际调用 AI 函数时，才会根据求值结果查找 `pipeline.model` 中声明的模型；如果选中的模型未声明，或没有实现函数所需的 capability，当前记录会在运行时报错。 + +例如，下面的表达式会根据每条记录的优先级选择模型： + +```sql +AI_COMPLETE( + IF(priority = 'high', 'powerful_model', 'economical_model'), + content, + '总结输入内容' +) +``` 所有文本函数都会将模型返回的 JSON 解析为 `VARIANT`。 diff --git a/docs/content/docs/core-concept/ai-model.md b/docs/content/docs/core-concept/ai-model.md index 1ad5a739fe0..7267098c50b 100644 --- a/docs/content/docs/core-concept/ai-model.md +++ b/docs/content/docs/core-concept/ai-model.md @@ -29,7 +29,17 @@ image understanding. ## AI Functions -The model name must be a string constant that refers to a model declared in `pipeline.model`. Text functions require a model client that implements text generation, while embedding and image functions require their corresponding capabilities. The pipeline validates the referenced model capability before execution. +The model argument accepts any `STRING` expression and is evaluated for each record. This enables dynamic model selection with a column, `IF`, or `CASE`. The selected name is resolved against the models declared in `pipeline.model` only when the AI function is invoked. If the selected model is undeclared or does not provide the capability required by the function, processing of that record fails at runtime. + +For example, the following expression chooses a model based on each record's priority: + +```sql +AI_COMPLETE( + IF(priority = 'high', 'powerful_model', 'economical_model'), + content, + 'Summarize the input' +) +``` All text functions return `VARIANT` values parsed from the model's JSON response. diff --git a/flink-cdc-composer/src/main/java/org/apache/flink/cdc/composer/flink/translator/TransformTranslator.java b/flink-cdc-composer/src/main/java/org/apache/flink/cdc/composer/flink/translator/TransformTranslator.java index 5dc9b268f73..99be446f200 100644 --- a/flink-cdc-composer/src/main/java/org/apache/flink/cdc/composer/flink/translator/TransformTranslator.java +++ b/flink-cdc-composer/src/main/java/org/apache/flink/cdc/composer/flink/translator/TransformTranslator.java @@ -35,17 +35,14 @@ import org.apache.flink.cdc.runtime.operators.transform.PostTransformOperatorBuilder; import org.apache.flink.cdc.runtime.operators.transform.PreTransformOperator; import org.apache.flink.cdc.runtime.operators.transform.PreTransformOperatorBuilder; -import org.apache.flink.cdc.runtime.parser.TransformParser; import org.apache.flink.cdc.runtime.typeutils.EventTypeInfo; import org.apache.flink.streaming.api.datastream.DataStream; import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; import java.util.Collections; -import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.stream.Collectors; /** @@ -67,8 +64,6 @@ public DataStream translatePreTransform( if (transforms.isEmpty()) { return input; } - validateModelReferences( - transforms, models, getUserDefinedFunctionNames(udfFunctions, models)); return input.transform( "Transform:Schema", new EventTypeInfo(), @@ -147,8 +142,6 @@ public DataStream translatePostTransform( .map(this::modelToUDFTuple) .collect(Collectors.toList())); Map modelClients = loadModelClients(models, env); - validateModelCapabilities( - transforms, modelClients, getUserDefinedFunctionNames(udfFunctions, models)); postTransformFunctionBuilder.addModelClients(modelClients); return input.transform( "Transform:Data", new EventTypeInfo(), postTransformFunctionBuilder.build()) @@ -191,48 +184,6 @@ private Map loadModelClients( return clients; } - private void validateModelReferences( - List transforms, - List models, - Set userDefinedFunctionNames) { - Set clientModelNames = - models.stream() - .filter(model -> !model.isLegacy()) - .map(ModelDef::getName) - .collect(Collectors.toSet()); - for (TransformDef transform : transforms) { - TransformParser.validateAiModelReferences( - transform.getProjection(), - transform.getFilter(), - clientModelNames, - userDefinedFunctionNames); - } - } - - private void validateModelCapabilities( - List transforms, - Map modelClients, - Set userDefinedFunctionNames) { - for (TransformDef transform : transforms) { - TransformParser.validateAiModelCapabilities( - transform.getProjection(), - transform.getFilter(), - modelClients, - userDefinedFunctionNames); - } - } - - private Set getUserDefinedFunctionNames( - List udfFunctions, List models) { - Set functionNames = new HashSet<>(); - udfFunctions.stream().map(UdfDef::getName).forEach(functionNames::add); - models.stream() - .filter(ModelDef::isLegacy) - .map(ModelDef::getName) - .forEach(functionNames::add); - return functionNames; - } - private Tuple3> udfDefToUDFTuple(UdfDef udf) { return Tuple3.of(udf.getName(), udf.getClasspath(), udf.getOptions()); } diff --git a/flink-cdc-composer/src/test/java/org/apache/flink/cdc/composer/flink/FlinkPipelineAiFunctionITCase.java b/flink-cdc-composer/src/test/java/org/apache/flink/cdc/composer/flink/FlinkPipelineAiFunctionITCase.java index d0610ed9f58..22d79749b74 100644 --- a/flink-cdc-composer/src/test/java/org/apache/flink/cdc/composer/flink/FlinkPipelineAiFunctionITCase.java +++ b/flink-cdc-composer/src/test/java/org/apache/flink/cdc/composer/flink/FlinkPipelineAiFunctionITCase.java @@ -128,6 +128,20 @@ void testAiCompleteInProjection() throws Exception { "Dummy model closed."); } + @Test + void testDynamicModelSelectionInProjection() throws Exception { + String[] output = + runAiFunctionTest( + "id, content, " + + "AI_COMPLETE(IF(id = 1, 'testModel', 'missingModel'), content, 'Complete the text') AS completed", + List.of(ModelDef.of("testModel", "dummy", Collections.emptyMap()))); + + assertThat(output) + .containsExactly( + "CreateTableEvent{tableId=default_namespace.default_schema.mytable1, schema=columns={`id` INT NOT NULL,`content` STRING,`completed` VARIANT}, primaryKeys=id, options=()}", + "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, I love this product, {\"result\":\"dummy response\"}], op=INSERT, meta=()}"); + } + @Test void testSpecializedTextAiFunctionsInProjection() throws Exception { String[] output = diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java new file mode 100644 index 00000000000..5e6f50419a6 --- /dev/null +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.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.cdc.runtime.ai; + +import org.apache.flink.cdc.common.annotation.Internal; +import org.apache.flink.cdc.common.model.AiModelClient; +import org.apache.flink.cdc.common.utils.Preconditions; + +import javax.annotation.Nullable; + +import java.util.Map; + +/** Resolves AI model clients by their logical names during expression evaluation. */ +@Internal +public class AiModelClientResolver { + + private final Map modelClients; + + public AiModelClientResolver(Map modelClients) { + this.modelClients = Preconditions.checkNotNull(modelClients); + } + + @Nullable + public AiModelClient resolve(String modelName) { + return modelClients.get(modelName); + } +} diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java index 2c21ba1c71f..71e7eea689d 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java @@ -25,6 +25,7 @@ import org.apache.flink.cdc.common.types.RowType; import org.apache.flink.cdc.common.types.variant.BinaryVariant; import org.apache.flink.cdc.common.types.variant.BinaryVariantInternalBuilder; +import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.ai.AiTextFunctionDef; import org.apache.flink.shaded.guava31.com.google.common.primitives.Floats; @@ -39,53 +40,93 @@ public class AiFunctions { private AiFunctions() {} - public static BinaryVariant aiComplete(AiModelClient model, String input, String systemPrompt) { - return generateText(model, AiTextFunctionDef.AI_COMPLETE, input, systemPrompt); + public static BinaryVariant aiComplete( + String modelName, + String input, + String systemPrompt, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, modelName, AiTextFunctionDef.AI_COMPLETE, input, systemPrompt); } - public static BinaryVariant aiClassify(AiModelClient model, String input, String labels) { - return generateText(model, AiTextFunctionDef.AI_CLASSIFY, input, labels); + public static BinaryVariant aiClassify( + String modelName, + String input, + String labels, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, modelName, AiTextFunctionDef.AI_CLASSIFY, input, labels); } public static BinaryVariant aiTranslate( - AiModelClient model, String input, String sourceLang, String targetLang) { - return generateText(model, AiTextFunctionDef.AI_TRANSLATE, input, sourceLang, targetLang); + String modelName, + String input, + String sourceLang, + String targetLang, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, + modelName, + AiTextFunctionDef.AI_TRANSLATE, + input, + sourceLang, + targetLang); } - public static BinaryVariant aiSummarize(AiModelClient model, String input, int maxLength) { - return generateText(model, AiTextFunctionDef.AI_SUMMARIZE, input, maxLength); + public static BinaryVariant aiSummarize( + String modelName, + String input, + int maxLength, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, modelName, AiTextFunctionDef.AI_SUMMARIZE, input, maxLength); } - public static BinaryVariant aiSentiment(AiModelClient model, String input) { - return generateText(model, AiTextFunctionDef.AI_SENTIMENT, input); + public static BinaryVariant aiSentiment( + String modelName, String input, AiModelClientResolver modelClientResolver) { + return generateText(modelClientResolver, modelName, AiTextFunctionDef.AI_SENTIMENT, input); } - public static BinaryVariant aiExtract(AiModelClient model, String input, String schema) { - return generateText(model, AiTextFunctionDef.AI_EXTRACT, input, schema); + public static BinaryVariant aiExtract( + String modelName, + String input, + String schema, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, modelName, AiTextFunctionDef.AI_EXTRACT, input, schema); } - public static BinaryVariant aiMask(AiModelClient model, String input, String entities) { - return generateText(model, AiTextFunctionDef.AI_MASK, input, entities); + public static BinaryVariant aiMask( + String modelName, + String input, + String entities, + AiModelClientResolver modelClientResolver) { + return generateText( + modelClientResolver, modelName, AiTextFunctionDef.AI_MASK, input, entities); } private static BinaryVariant generateText( - AiModelClient model, + AiModelClientResolver modelClientResolver, + String modelName, AiTextFunctionDef function, String input, Object... promptArguments) { if (input == null) { return null; } - if (!(model instanceof SupportsTextGeneration)) { - throw new UnsupportedOperationException( - "Model " + model.getClass().getName() + " does not support text generation"); - } + SupportsTextGeneration model = + resolveModel( + modelClientResolver, + modelName, + function.getFunctionName(), + SupportsTextGeneration.class, + "text generation"); String prompt = function.buildPrompt(promptArguments) + "\n" + buildOutputSchemaHint(function.getOutputType()); - String json = ((SupportsTextGeneration) model).generate(prompt, input); + String json = model.generate(prompt, input); if (json == null) { return null; } @@ -101,43 +142,88 @@ private static BinaryVariant generateText( } } - public static List aiEmbed(AiModelClient model, String input) { + public static List aiEmbed( + String modelName, String input, AiModelClientResolver modelClientResolver) { if (input == null) { return null; } - if (!(model instanceof SupportsEmbedding)) { - throw new UnsupportedOperationException( - "Model " + model.getClass().getName() + " does not support embedding"); - } - float[] embedding = ((SupportsEmbedding) model).embed(input); + SupportsEmbedding model = + resolveModel( + modelClientResolver, + modelName, + "AI_EMBED", + SupportsEmbedding.class, + "embedding"); + float[] embedding = model.embed(input); return embedding == null ? null : Floats.asList(embedding); } /** Dispatches image-to-text AI functions. */ - public static String aiImageComplete(AiModelClient model, byte[] image, String prompt) { + public static String aiImageComplete( + String modelName, + byte[] image, + String prompt, + AiModelClientResolver modelClientResolver) { if (image == null) { return null; } - if (!(model instanceof SupportsImageTextGeneration)) { - throw new UnsupportedOperationException( - "Model " - + model.getClass().getName() - + " does not support image text generation"); - } - return ((SupportsImageTextGeneration) model).generateTextFromImage(image, prompt); + SupportsImageTextGeneration model = + resolveModel( + modelClientResolver, + modelName, + "AI_IMAGE_COMPLETE", + SupportsImageTextGeneration.class, + "image text generation"); + return model.generateTextFromImage(image, prompt); } /** Dispatches image embedding AI functions. */ - public static List aiImageEmbed(AiModelClient model, byte[] image) { + public static List aiImageEmbed( + String modelName, byte[] image, AiModelClientResolver modelClientResolver) { if (image == null) { return null; } - if (!(model instanceof SupportsImageEmbedding)) { + SupportsImageEmbedding model = + resolveModel( + modelClientResolver, + modelName, + "AI_IMAGE_EMBED", + SupportsImageEmbedding.class, + "image embedding"); + float[] embedding = model.embedImage(image); + return embedding == null ? null : Floats.asList(embedding); + } + + private static T resolveModel( + AiModelClientResolver modelClientResolver, + String modelName, + String functionName, + Class requiredCapability, + String capabilityName) { + if (modelName == null) { + throw new IllegalArgumentException( + "Model name referenced by " + functionName + " must not be null."); + } + AiModelClient model = modelClientResolver.resolve(modelName); + if (model == null) { + throw new IllegalArgumentException( + "Model '" + + modelName + + "' referenced by " + + functionName + + " has not been declared."); + } + if (!requiredCapability.isInstance(model)) { throw new UnsupportedOperationException( - "Model " + model.getClass().getName() + " does not support image embedding"); + "Model '" + + modelName + + "' referenced by " + + functionName + + " does not support " + + capabilityName + + "."); } - float[] embedding = ((SupportsImageEmbedding) model).embedImage(image); - return embedding == null ? null : Floats.asList(embedding); + return requiredCapability.cast(model); } private static String truncateInvalidJsonResponse(String response) { diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java index 2dc86b37bea..fef5896f57d 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java @@ -21,6 +21,7 @@ import org.apache.flink.cdc.common.model.AiModelClient; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; +import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.codehaus.janino.ExpressionEvaluator; @@ -47,7 +48,7 @@ public class ProjectionColumnProcessor { private final TransformExpressionKey transformExpressionKey; private final Map supportedMetadataColumns; private final List udfFunctionInstances; - private final Map modelClients; + private final AiModelClientResolver modelClientResolver; private final ExpressionEvaluator expressionEvaluator; public ProjectionColumnProcessor( @@ -62,11 +63,11 @@ public ProjectionColumnProcessor( this.projectionColumn = projectionColumn; this.timezone = timezone; this.supportedMetadataColumns = supportedMetadataColumns; - this.modelClients = modelClients; + this.modelClientResolver = new AiModelClientResolver(modelClients); this.transformExpressionKey = generateTransformExpressionKey(); this.expressionEvaluator = TransformExpressionCompiler.compileExpression( - transformExpressionKey, udfDescriptors, modelClients); + transformExpressionKey, udfDescriptors); this.udfFunctionInstances = udfFunctionInstances; } @@ -148,8 +149,8 @@ private Object[] generateParams(Object[] rowData, TransformContext context) { // 3 - Add UDF function instances params.addAll(udfFunctionInstances); - // 4 - Add AI model client instances - params.addAll(modelClients.values()); + // 4 - Add AI model client resolver + params.add(modelClientResolver); return params.toArray(); } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java index 47e42eb90bb..6a2e26e2e07 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java @@ -18,8 +18,9 @@ package org.apache.flink.cdc.runtime.operators.transform; import org.apache.flink.api.common.InvalidProgramException; -import org.apache.flink.cdc.common.model.AiModelClient; +import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.operators.transform.exceptions.TransformException; +import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.apache.flink.util.FlinkRuntimeException; import org.apache.flink.shaded.guava31.com.google.common.cache.Cache; @@ -31,9 +32,7 @@ import org.slf4j.LoggerFactory; import java.util.ArrayList; -import java.util.Collections; import java.util.List; -import java.util.Map; /** * The processor of the transform expression. It processes the expression of projections and @@ -57,20 +56,6 @@ public static void cleanUp() { /** Compiles an expression code to a janino {@link ExpressionEvaluator}. */ public static ExpressionEvaluator compileExpression( TransformExpressionKey key, List udfDescriptors) { - return compileExpression(key, udfDescriptors, Collections.emptyMap()); - } - - /** - * Compiles an expression code to a janino {@link ExpressionEvaluator}, with additional {@link - * AiModelClient} instances appended after UDF instances. - * - *

{@code modelClients} maps model names (e.g. {@code myModel}) to the corresponding client - * instances. - */ - public static ExpressionEvaluator compileExpression( - TransformExpressionKey key, - List udfDescriptors, - Map modelClients) { try { return COMPILED_EXPRESSION_CACHE.get( key, @@ -85,10 +70,8 @@ public static ExpressionEvaluator compileExpression( argumentClasses.add(Class.forName(udfFunction.getClasspath())); } - for (String paramName : modelClients.keySet()) { - argumentNames.add(paramName); - argumentClasses.add(AiModelClient.class); - } + argumentNames.add(JaninoCompiler.DEFAULT_AI_MODEL_CLIENT_RESOLVER); + argumentClasses.add(AiModelClientResolver.class); // Input args expressionEvaluator.setParameters( diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java index 02b2c517166..25f1c546107 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java @@ -23,6 +23,7 @@ import org.apache.flink.cdc.common.pipeline.DecimalPrecisionMode; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; +import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.apache.flink.cdc.runtime.parser.TransformParser; @@ -30,6 +31,7 @@ import java.lang.reflect.InvocationTargetException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashSet; @@ -50,7 +52,7 @@ public class TransformFilterProcessor { private final DecimalPrecisionMode decimalPrecisionMode; private final List udfFunctionInstances; private final Map supportedMetadataColumns; - private final Map modelClients; + private final AiModelClientResolver modelClientResolver; private final TransformExpressionKey transformExpressionKey; private final ExpressionEvaluator expressionEvaluator; @@ -72,7 +74,9 @@ protected TransformFilterProcessor( this.decimalPrecisionMode = decimalPrecisionMode; this.udfFunctionInstances = udfFunctionInstances; this.supportedMetadataColumns = supportedMetadataColumns; - this.modelClients = modelClients; + this.modelClientResolver = + new AiModelClientResolver( + modelClients == null ? Collections.emptyMap() : modelClients); if (isNoOp) { this.transformExpressionKey = null; @@ -87,7 +91,7 @@ protected TransformFilterProcessor( .toArray(new SupportedMetadataColumn[0])); this.expressionEvaluator = TransformExpressionCompiler.compileExpression( - transformExpressionKey, udfDescriptors, modelClients); + transformExpressionKey, udfDescriptors); } } @@ -223,8 +227,8 @@ private Object[] generateParams(Object[] preRow, Object[] postRow, TransformCont // 3 - Add UDF function instances params.addAll(udfFunctionInstances); - // 4 - Add AI model client instances - params.addAll(modelClients.values()); + // 4 - Add AI model client resolver + params.add(modelClientResolver); return params.toArray(); } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java index 627cde88fa2..81c6bc55a1f 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java @@ -119,6 +119,7 @@ public class JaninoCompiler { public static final String DEFAULT_EPOCH_TIME = "__epoch_time__"; public static final String DEFAULT_TIME_ZONE = "__time_zone__"; + public static final String DEFAULT_AI_MODEL_CLIENT_RESOLVER = "__ai_model_client_resolver__"; private static final String[] BUILTIN_FUNCTION_MODULES = { "Ai", "Arithmetic", "Casting", "Comparison", "Logical", "String", "Struct", "Temporal" @@ -264,6 +265,10 @@ private static Java.Rvalue translateSqlBasicCall(Context context, SqlBasicCall s } else if (TIMEZONE_REQUIRED_TEMPORAL_CONVERSION_FUNCTIONS.contains( sqlBasicCall.getOperator().getName().toUpperCase())) { atoms.add(new Java.AmbiguousName(Location.NOWHERE, new String[] {DEFAULT_TIME_ZONE})); + } else if (isAiFunction(functionName)) { + atoms.add( + new Java.AmbiguousName( + Location.NOWHERE, new String[] {DEFAULT_AI_MODEL_CLIENT_RESOLVER})); } return sqlBasicCallToJaninoRvalue(context, sqlBasicCall, atoms.toArray(new Java.Rvalue[0])); } @@ -1002,13 +1007,6 @@ private static Java.Rvalue generateOtherFunctionOperation( return castExpressionToInferredType( context, sqlBasicCall, generateFunctionOperation("element", atoms)); } else { - if (isAiFunction(operationName) && atoms.length >= 1) { - if (!(sqlBasicCall.operand(0) instanceof SqlCharStringLiteral)) { - throw new ParseException( - "The model argument of an AI function must be a string constant."); - } - rewriteAiFunctionModelArg(atoms); - } return new Java.MethodInvocation( Location.NOWHERE, null, @@ -1122,14 +1120,6 @@ private static boolean isAiFunction(String upperCaseName) { return false; } - private static void rewriteAiFunctionModelArg(Java.Rvalue[] atoms) { - String modelName = atoms[0].toString(); - if (modelName.startsWith("\"") && modelName.endsWith("\"")) { - modelName = modelName.substring(1, modelName.length() - 1); - } - atoms[0] = new Java.AmbiguousName(Location.NOWHERE, new String[] {modelName}); - } - private static Java.Rvalue generateTimezoneFreeTemporalFunctionOperation( Context context, String operationName) { return new Java.MethodInvocation( diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/TransformParser.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/TransformParser.java index 5993691af92..ebddd1de8a6 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/TransformParser.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/TransformParser.java @@ -18,19 +18,11 @@ package org.apache.flink.cdc.runtime.parser; import org.apache.flink.api.common.io.ParseException; -import org.apache.flink.cdc.common.model.AiModelClient; -import org.apache.flink.cdc.common.model.abilities.SupportsEmbedding; -import org.apache.flink.cdc.common.model.abilities.SupportsImageEmbedding; -import org.apache.flink.cdc.common.model.abilities.SupportsImageTextGeneration; -import org.apache.flink.cdc.common.model.abilities.SupportsTextGeneration; import org.apache.flink.cdc.common.pipeline.DecimalPrecisionMode; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; import org.apache.flink.cdc.common.types.DataType; import org.apache.flink.cdc.common.utils.Preconditions; -import org.apache.flink.cdc.runtime.ai.AiEmbeddingFunctionDef; -import org.apache.flink.cdc.runtime.ai.AiImageFunctionDef; -import org.apache.flink.cdc.runtime.ai.AiTextFunctionDef; import org.apache.flink.cdc.runtime.operators.transform.ProjectionColumn; import org.apache.flink.cdc.runtime.operators.transform.UserDefinedFunctionDescriptor; import org.apache.flink.cdc.runtime.parser.metadata.AiFunctionSqlOperatorTable; @@ -56,7 +48,6 @@ import org.apache.calcite.schema.impl.ScalarFunctionImpl; import org.apache.calcite.sql.SqlBasicCall; import org.apache.calcite.sql.SqlCall; -import org.apache.calcite.sql.SqlCharStringLiteral; import org.apache.calcite.sql.SqlFunction; import org.apache.calcite.sql.SqlFunctionCategory; import org.apache.calcite.sql.SqlIdentifier; @@ -906,206 +897,6 @@ public static SqlSelect parseFilterExpression(String filterExpression) { return parseSelect(statement.toString()); } - /** Validates model arguments and references in the supported AI functions. */ - public static void validateAiModelReferences( - @Nullable String projection, @Nullable String filter, Set declaredModelNames) { - validateAiModelReferences(projection, filter, declaredModelNames, Collections.emptySet()); - } - - /** Validates model arguments and references in AI functions not shadowed by a UDF. */ - public static void validateAiModelReferences( - @Nullable String projection, - @Nullable String filter, - Set declaredModelNames, - Set userDefinedFunctionNames) { - if (!isNullOrWhitespaceOnly(projection)) { - validateAiModelReferences( - parseProjectionExpression(projection), - declaredModelNames, - userDefinedFunctionNames); - } - if (!isNullOrWhitespaceOnly(filter)) { - validateAiModelReferences( - parseFilterExpression(filter), declaredModelNames, userDefinedFunctionNames); - } - } - - /** Validates that referenced models provide the capability required by each AI function. */ - public static void validateAiModelCapabilities( - @Nullable String projection, - @Nullable String filter, - Map modelClients) { - validateAiModelCapabilities(projection, filter, modelClients, Collections.emptySet()); - } - - /** Validates model capabilities for AI functions not shadowed by a UDF. */ - public static void validateAiModelCapabilities( - @Nullable String projection, - @Nullable String filter, - Map modelClients, - Set userDefinedFunctionNames) { - if (!isNullOrWhitespaceOnly(projection)) { - validateAiModelCapabilities( - parseProjectionExpression(projection), modelClients, userDefinedFunctionNames); - } - if (!isNullOrWhitespaceOnly(filter)) { - validateAiModelCapabilities( - parseFilterExpression(filter), modelClients, userDefinedFunctionNames); - } - } - - private static void validateAiModelReferences( - SqlNode node, Set declaredModelNames, Set userDefinedFunctionNames) { - if (node instanceof SqlCall) { - SqlCall call = (SqlCall) node; - String functionName = call.getOperator().getName(); - if (isAiFunction(functionName) - && !isUserDefinedFunction(functionName, userDefinedFunctionNames)) { - if (call.operandCount() == 0) { - return; - } - String modelName = resolveAiModelName(call); - Preconditions.checkArgument( - declaredModelNames.contains(modelName), - "Model '%s' referenced by %s has not been declared.", - modelName, - call.getOperator().getName()); - } - for (SqlNode operand : call.getOperandList()) { - if (operand != null) { - validateAiModelReferences( - operand, declaredModelNames, userDefinedFunctionNames); - } - } - } else if (node instanceof SqlNodeList) { - for (SqlNode child : (SqlNodeList) node) { - validateAiModelReferences(child, declaredModelNames, userDefinedFunctionNames); - } - } - } - - private static void validateAiModelCapabilities( - SqlNode node, - Map modelClients, - Set userDefinedFunctionNames) { - if (node instanceof SqlCall) { - SqlCall call = (SqlCall) node; - String functionName = call.getOperator().getName(); - if (isAiFunction(functionName) - && !isUserDefinedFunction(functionName, userDefinedFunctionNames) - && call.operandCount() > 0) { - String modelName = resolveAiModelName(call); - AiModelClient modelClient = modelClients.get(modelName); - Preconditions.checkArgument( - modelClient != null, - "Model '%s' referenced by %s has not been declared.", - modelName, - functionName); - AiImageFunctionDef imageFunction = findImageAiFunction(functionName); - if (isTextAiFunction(functionName)) { - Preconditions.checkArgument( - modelClient instanceof SupportsTextGeneration, - "Model '%s' referenced by %s does not support text generation.", - modelName, - functionName); - } else if (imageFunction != null) { - validateImageAiModelCapability( - imageFunction, modelClient, modelName, functionName); - } else { - Preconditions.checkArgument( - modelClient instanceof SupportsEmbedding, - "Model '%s' referenced by %s does not support embedding.", - modelName, - functionName); - } - } - for (SqlNode operand : call.getOperandList()) { - if (operand != null) { - validateAiModelCapabilities(operand, modelClients, userDefinedFunctionNames); - } - } - } else if (node instanceof SqlNodeList) { - for (SqlNode child : (SqlNodeList) node) { - validateAiModelCapabilities(child, modelClients, userDefinedFunctionNames); - } - } - } - - private static void validateImageAiModelCapability( - AiImageFunctionDef function, - AiModelClient modelClient, - String modelName, - String functionName) { - switch (function.getCapability()) { - case IMAGE_TEXT_GENERATION: - Preconditions.checkArgument( - modelClient instanceof SupportsImageTextGeneration, - "Model '%s' referenced by %s does not support image text generation.", - modelName, - functionName); - break; - case IMAGE_EMBEDDING: - Preconditions.checkArgument( - modelClient instanceof SupportsImageEmbedding, - "Model '%s' referenced by %s does not support image embedding.", - modelName, - functionName); - break; - default: - throw new IllegalArgumentException( - "Unsupported capability for image AI function " + functionName); - } - } - - private static String resolveAiModelName(SqlCall call) { - SqlNode modelArgument = call.operand(0); - Preconditions.checkArgument( - modelArgument instanceof SqlCharStringLiteral, - "The model argument of %s must be a string constant, but was %s.", - call.getOperator().getName(), - modelArgument); - return ((SqlCharStringLiteral) modelArgument).getNlsString().getValue(); - } - - private static boolean isUserDefinedFunction( - String functionName, Set userDefinedFunctionNames) { - return userDefinedFunctionNames.stream().anyMatch(functionName::equalsIgnoreCase); - } - - private static boolean isAiFunction(String functionName) { - return isTextAiFunction(functionName) - || isEmbeddingAiFunction(functionName) - || findImageAiFunction(functionName) != null; - } - - private static boolean isTextAiFunction(String functionName) { - for (AiTextFunctionDef function : AiTextFunctionDef.values()) { - if (function.getFunctionName().equalsIgnoreCase(functionName)) { - return true; - } - } - return false; - } - - private static boolean isEmbeddingAiFunction(String functionName) { - for (AiEmbeddingFunctionDef function : AiEmbeddingFunctionDef.values()) { - if (function.getFunctionName().equalsIgnoreCase(functionName)) { - return true; - } - } - return false; - } - - @Nullable - private static AiImageFunctionDef findImageAiFunction(String functionName) { - for (AiImageFunctionDef function : AiImageFunctionDef.values()) { - if (function.getFunctionName().equalsIgnoreCase(functionName)) { - return function; - } - } - return null; - } - public static boolean hasAsterisk(@Nullable String projection) { if (isNullOrWhitespaceOnly(projection)) { // Providing an empty projection expression is equivalent to writing `*` explicitly. diff --git a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java index 8eee49ae56f..05510843723 100644 --- a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java +++ b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java @@ -22,11 +22,14 @@ import org.apache.flink.cdc.common.model.abilities.SupportsImageEmbedding; import org.apache.flink.cdc.common.model.abilities.SupportsImageTextGeneration; import org.apache.flink.cdc.common.model.abilities.SupportsTextGeneration; +import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.junit.jupiter.api.Test; import java.util.ArrayList; +import java.util.Collections; import java.util.List; +import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -89,19 +92,21 @@ private static class UnsupportedModelClient implements AiModelClient { @Test void testTextAiFunctionsUseEnglishPromptsAndParseJsonResponses() { TestModelClient model = new TestModelClient(); + AiModelClientResolver resolver = resolver("textModel", model); - assertThat(AiFunctions.aiComplete(model, "input", "Return three letters")) + assertThat(AiFunctions.aiComplete("textModel", "input", "Return three letters", resolver)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiClassify(model, "input", "positive,negative")) + assertThat(AiFunctions.aiClassify("textModel", "input", "positive,negative", resolver)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiTranslate(model, "input", "auto", "en")) + assertThat(AiFunctions.aiTranslate("textModel", "input", "auto", "en", resolver)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiSummarize(model, "input", 100)) + assertThat(AiFunctions.aiSummarize("textModel", "input", 100, resolver)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiSentiment(model, "input")).hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiExtract(model, "input", "name:string")) + assertThat(AiFunctions.aiSentiment("textModel", "input", resolver)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiMask(model, "input", "email,phone")) + assertThat(AiFunctions.aiExtract("textModel", "input", "name:string", resolver)) + .hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiMask("textModel", "input", "email,phone", resolver)) .hasToString("{\"result\":\"ABC\"}"); assertThat(model.prompts).hasSize(7); @@ -128,19 +133,23 @@ void testTextAiFunctionsUseEnglishPromptsAndParseJsonResponses() { @Test void testEmbeddingFunction() { TestModelClient model = new TestModelClient(); + AiModelClientResolver resolver = resolver("embeddingModel", model); - assertThat(AiFunctions.aiEmbed(model, "input")).containsExactly(0.1f, 0.2f, 0.3f); + assertThat(AiFunctions.aiEmbed("embeddingModel", "input", resolver)) + .containsExactly(0.1f, 0.2f, 0.3f); assertThat(model.embedCalls).isOne(); } @Test void testImageAiFunctions() { TestModelClient model = new TestModelClient(); + AiModelClientResolver resolver = resolver("imageModel", model); byte[] image = new byte[] {1, 2, 3, 4}; - assertThat(AiFunctions.aiImageComplete(model, image, "Describe the image")) + assertThat(AiFunctions.aiImageComplete("imageModel", image, "Describe the image", resolver)) .isEqualTo("image has 4 bytes, prompt: Describe the image"); - assertThat(AiFunctions.aiImageEmbed(model, image)).containsExactly(0.9f, 0.8f, 0.7f); + assertThat(AiFunctions.aiImageEmbed("imageModel", image, resolver)) + .containsExactly(0.9f, 0.8f, 0.7f); assertThat(model.imageTextCalls).isOne(); assertThat(model.imageEmbedCalls).isOne(); } @@ -148,26 +157,57 @@ void testImageAiFunctions() { @Test void testUnsupportedCapabilities() { UnsupportedModelClient model = new UnsupportedModelClient(); + AiModelClientResolver resolver = resolver("unsupported", model); - assertThatThrownBy(() -> AiFunctions.aiComplete(model, "input", "prompt")) + assertThatThrownBy(() -> AiFunctions.aiComplete("unsupported", "input", "prompt", resolver)) .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("Model 'unsupported'") + .hasMessageContaining("AI_COMPLETE") .hasMessageContaining("does not support text generation"); - assertThatThrownBy(() -> AiFunctions.aiEmbed(model, "input")) + assertThatThrownBy(() -> AiFunctions.aiEmbed("unsupported", "input", resolver)) .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("Model 'unsupported'") + .hasMessageContaining("AI_EMBED") .hasMessageContaining("does not support embedding"); - assertThatThrownBy(() -> AiFunctions.aiImageComplete(model, new byte[] {1, 2}, "describe")) + assertThatThrownBy( + () -> + AiFunctions.aiImageComplete( + "unsupported", new byte[] {1, 2}, "describe", resolver)) .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("Model 'unsupported'") + .hasMessageContaining("AI_IMAGE_COMPLETE") .hasMessageContaining("does not support image text generation"); - assertThatThrownBy(() -> AiFunctions.aiImageEmbed(model, new byte[] {1, 2})) + assertThatThrownBy( + () -> AiFunctions.aiImageEmbed("unsupported", new byte[] {1, 2}, resolver)) .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("Model 'unsupported'") + .hasMessageContaining("AI_IMAGE_EMBED") .hasMessageContaining("does not support image embedding"); } + @Test + void testInvalidModelName() { + AiModelClientResolver resolver = new AiModelClientResolver(Collections.emptyMap()); + + assertThatThrownBy( + () -> AiFunctions.aiComplete("missingModel", "input", "prompt", resolver)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage( + "Model 'missingModel' referenced by AI_COMPLETE has not been declared."); + assertThatThrownBy(() -> AiFunctions.aiEmbed(null, "input", resolver)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Model name referenced by AI_EMBED must not be null."); + } + @Test void testInvalidJsonResponse() { TestModelClient model = new TestModelClient("not-json"); + AiModelClientResolver resolver = resolver("model", model); - assertThatThrownBy(() -> AiFunctions.aiClassify(model, "input", "positive,negative")) + assertThatThrownBy( + () -> + AiFunctions.aiClassify( + "model", "input", "positive,negative", resolver)) .isInstanceOf(RuntimeException.class) .hasMessage("AI function AI_CLASSIFY returned invalid JSON: not-json"); } @@ -176,8 +216,12 @@ void testInvalidJsonResponse() { void testInvalidJsonResponseIsTruncated() { String longInvalidJson = "x".repeat(600); TestModelClient model = new TestModelClient(longInvalidJson); + AiModelClientResolver resolver = resolver("model", model); - assertThatThrownBy(() -> AiFunctions.aiClassify(model, "input", "positive,negative")) + assertThatThrownBy( + () -> + AiFunctions.aiClassify( + "model", "input", "positive,negative", resolver)) .isInstanceOf(RuntimeException.class) .hasMessage( "AI function AI_CLASSIFY returned invalid JSON: " @@ -188,11 +232,14 @@ void testInvalidJsonResponseIsTruncated() { @Test void testNullInputSkipsModelInvocation() { TestModelClient model = new TestModelClient(); - - assertThat(AiFunctions.aiClassify(model, null, "positive,negative")).isNull(); - assertThat(AiFunctions.aiEmbed(model, null)).isNull(); - assertThat(AiFunctions.aiImageComplete(model, null, "describe")).isNull(); - assertThat(AiFunctions.aiImageEmbed(model, null)).isNull(); + AiModelClientResolver resolver = resolver("model", model); + AiModelClientResolver emptyResolver = new AiModelClientResolver(Collections.emptyMap()); + + assertThat(AiFunctions.aiClassify("model", null, "positive,negative", resolver)).isNull(); + assertThat(AiFunctions.aiEmbed("model", null, resolver)).isNull(); + assertThat(AiFunctions.aiImageComplete("model", null, "describe", resolver)).isNull(); + assertThat(AiFunctions.aiImageEmbed("model", null, resolver)).isNull(); + assertThat(AiFunctions.aiComplete("missing", null, "prompt", emptyResolver)).isNull(); assertThat(model.prompts).isEmpty(); assertThat(model.embedCalls).isZero(); assertThat(model.imageTextCalls).isZero(); @@ -202,8 +249,13 @@ void testNullInputSkipsModelInvocation() { @Test void testNullModelResponseReturnsNull() { TestModelClient model = new TestModelClient(null); + AiModelClientResolver resolver = resolver("model", model); - assertThat(AiFunctions.aiSummarize(model, "input", 100)).isNull(); + assertThat(AiFunctions.aiSummarize("model", "input", 100, resolver)).isNull(); assertThat(model.prompts).hasSize(1); } + + private static AiModelClientResolver resolver(String modelName, AiModelClient model) { + return new AiModelClientResolver(Map.of(modelName, model)); + } } diff --git a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java index b6dc28f1ca5..2d1fea5dcbd 100644 --- a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java +++ b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java @@ -17,11 +17,6 @@ package org.apache.flink.cdc.runtime.parser; -import org.apache.flink.cdc.common.model.AiModelClient; -import org.apache.flink.cdc.common.model.abilities.SupportsEmbedding; -import org.apache.flink.cdc.common.model.abilities.SupportsImageEmbedding; -import org.apache.flink.cdc.common.model.abilities.SupportsImageTextGeneration; -import org.apache.flink.cdc.common.model.abilities.SupportsTextGeneration; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; import org.apache.flink.cdc.common.types.DataTypes; @@ -32,11 +27,8 @@ import java.util.Collections; import java.util.List; -import java.util.Map; -import java.util.Set; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Parser and Janino tests for the generic AI functions. */ @@ -48,44 +40,6 @@ class AiFunctionParserTest { Column.physicalColumn("content", DataTypes.STRING()), Column.physicalColumn("image", DataTypes.BYTES())); - private static class TextModelClient implements AiModelClient, SupportsTextGeneration { - private static final long serialVersionUID = 1L; - - @Override - public String generate(String systemPrompt, String userInput) { - return "{}"; - } - } - - private static class EmbeddingModelClient implements AiModelClient, SupportsEmbedding { - private static final long serialVersionUID = 1L; - - @Override - public float[] embed(String text) { - return new float[0]; - } - } - - private static class ImageTextModelClient - implements AiModelClient, SupportsImageTextGeneration { - private static final long serialVersionUID = 1L; - - @Override - public String generateTextFromImage(byte[] image, String prompt) { - return "description"; - } - } - - private static class ImageEmbeddingModelClient - implements AiModelClient, SupportsImageEmbedding { - private static final long serialVersionUID = 1L; - - @Override - public float[] embedImage(byte[] image) { - return new float[0]; - } - } - @Test void testTranslateAiFunctions() { List columns = @@ -96,7 +50,8 @@ void testTranslateAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiComplete(completer, $0, \"You are helpful\")", "aiEmbed(embedder, $0)"); + "aiComplete(\"completer\", $0, \"You are helpful\", __ai_model_client_resolver__)", + "aiEmbed(\"embedder\", $0, __ai_model_client_resolver__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsExactly(DataTypes.VARIANT(), DataTypes.ARRAY(DataTypes.FLOAT())); @@ -116,12 +71,18 @@ void testTranslateSpecializedTextAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiClassify(model, $0, \"positive,negative\")", - "aiTranslate(model, $0, \"auto\", \"en\")", - "aiSummarize(model, $0, 100)", - "aiSentiment(model, $0)", - "aiExtract(model, $0, \"name:string\")", - "aiMask(model, $0, \"email,phone\")"); + "aiClassify(\"model\", $0, \"positive,negative\", __ai_model_client_resolver__)", + "aiTranslate(\n" + + " \"model\",\n" + + " $0,\n" + + " \"auto\",\n" + + " \"en\",\n" + + " __ai_model_client_resolver__\n" + + ")", + "aiSummarize(\"model\", $0, 100, __ai_model_client_resolver__)", + "aiSentiment(\"model\", $0, __ai_model_client_resolver__)", + "aiExtract(\"model\", $0, \"name:string\", __ai_model_client_resolver__)", + "aiMask(\"model\", $0, \"email,phone\", __ai_model_client_resolver__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsOnly(DataTypes.VARIANT()); @@ -137,33 +98,30 @@ void testTranslateImageAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiImageComplete(vision, $0, \"Describe the image\")", - "aiImageEmbed(imageEmbedder, $0)"); + "aiImageComplete(\"vision\", $0, \"Describe the image\", __ai_model_client_resolver__)", + "aiImageEmbed(\"imageEmbedder\", $0, __ai_model_client_resolver__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsExactly(DataTypes.STRING(), DataTypes.ARRAY(DataTypes.FLOAT())); } @Test - void testSameNamedUdfTakesPrecedenceOverAiFunction() { - Set udfNames = Set.of("ai_sentiment"); - assertThatCode( - () -> - TransformParser.validateAiModelReferences( - "AI_SENTIMENT(id) AS sentiment", - null, - Collections.emptySet(), - udfNames)) - .doesNotThrowAnyException(); - assertThatCode( - () -> - TransformParser.validateAiModelCapabilities( - "AI_SENTIMENT(id) AS sentiment", - null, - Collections.emptyMap(), - udfNames)) - .doesNotThrowAnyException(); + void testDynamicModelSelection() { + List columns = + translate( + "AI_COMPLETE(IF(id = 1, 'powerful', 'cheap'), content, 'prompt') AS completed"); + assertThat(columns) + .extracting(ProjectionColumn::getScriptExpression) + .containsExactly( + "aiComplete(isTrue(valueEquals($0, 1)) ? \"powerful\" : \"cheap\", $1, \"prompt\", __ai_model_client_resolver__)"); + assertThat(columns) + .extracting(ProjectionColumn::getDataType) + .containsExactly(DataTypes.VARIANT()); + } + + @Test + void testSameNamedUdfTakesPrecedenceOverAiFunction() { List columns = TransformParser.generateProjectionColumns( "AI_SENTIMENT(id) AS sentiment", @@ -192,18 +150,6 @@ void testSameNamedUdfTakesPrecedenceOverImageAiFunction() { private static void assertSameNamedUdfTakesPrecedenceOverImageAiFunction( String functionName, String udfName) { String projection = functionName + "(id) AS udf_output"; - Set udfNames = Set.of(udfName); - - assertThatCode( - () -> - TransformParser.validateAiModelReferences( - projection, null, Collections.emptySet(), udfNames)) - .doesNotThrowAnyException(); - assertThatCode( - () -> - TransformParser.validateAiModelCapabilities( - projection, null, Collections.emptyMap(), udfNames)) - .doesNotThrowAnyException(); List columns = TransformParser.generateProjectionColumns( @@ -223,110 +169,6 @@ private static void assertSameNamedUdfTakesPrecedenceOverImageAiFunction( .containsExactly(DataTypes.STRING()); } - @Test - void testModelArgumentMustBeStringConstant() { - assertThatThrownBy( - () -> - TransformParser.validateAiModelReferences( - "AI_COMPLETE(content, content, 'prompt') AS completed", - null, - Set.of("content"))) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("must be a string constant"); - } - - @Test - void testReferencedModelMustBeDeclared() { - assertThatThrownBy( - () -> - TransformParser.validateAiModelReferences( - "AI_EMBED('missing', content) AS embedding", - null, - Set.of("declared"))) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'missing'") - .hasMessageContaining("has not been declared"); - - assertThatCode( - () -> - TransformParser.validateAiModelReferences( - "AI_EMBED('declared', content) AS embedding", - null, - Set.of("declared"))) - .doesNotThrowAnyException(); - - assertThatThrownBy( - () -> - TransformParser.validateAiModelReferences( - "AI_CLASSIFY('missing', content, 'a,b') AS classified", - null, - Set.of("declared"))) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'missing'") - .hasMessageContaining("AI_CLASSIFY"); - } - - @Test - void testModelCapabilitiesMustMatchAiFunctions() { - Map models = - Map.of( - "textModel", new TextModelClient(), - "embeddingModel", new EmbeddingModelClient(), - "imageTextModel", new ImageTextModelClient(), - "imageEmbeddingModel", new ImageEmbeddingModelClient()); - - assertThatCode( - () -> - TransformParser.validateAiModelCapabilities( - "AI_CLASSIFY('textModel', content, 'a,b') AS classified, " - + "AI_EMBED('embeddingModel', content) AS embedding, " - + "AI_IMAGE_COMPLETE('imageTextModel', image, 'describe') AS description, " - + "AI_IMAGE_EMBED('imageEmbeddingModel', image) AS image_embedding", - null, - models)) - .doesNotThrowAnyException(); - assertThatThrownBy( - () -> - TransformParser.validateAiModelCapabilities( - "AI_SENTIMENT('embeddingModel', content) AS sentiment", - null, - models)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'embeddingModel'") - .hasMessageContaining("AI_SENTIMENT") - .hasMessageContaining("does not support text generation"); - assertThatThrownBy( - () -> - TransformParser.validateAiModelCapabilities( - "AI_EMBED('textModel', content) AS embedding", - null, - models)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'textModel'") - .hasMessageContaining("AI_EMBED") - .hasMessageContaining("does not support embedding"); - assertThatThrownBy( - () -> - TransformParser.validateAiModelCapabilities( - "AI_IMAGE_COMPLETE('textModel', image, 'describe') AS description", - null, - models)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'textModel'") - .hasMessageContaining("AI_IMAGE_COMPLETE") - .hasMessageContaining("does not support image text generation"); - assertThatThrownBy( - () -> - TransformParser.validateAiModelCapabilities( - "AI_IMAGE_EMBED('embeddingModel', image) AS embedding", - null, - models)) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Model 'embeddingModel'") - .hasMessageContaining("AI_IMAGE_EMBED") - .hasMessageContaining("does not support image embedding"); - } - @Test void testFunctionArityValidation() { assertThatThrownBy(() -> translate("AI_EMBED('model') AS embedding")) @@ -352,11 +194,6 @@ void testFunctionArityValidation() { "Invalid number of arguments to function 'AI_IMAGE_COMPLETE'"); assertThatThrownBy(() -> translate("AI_IMAGE_EMBED('model') AS embedding")) .hasMessageContaining("Invalid number of arguments to function 'AI_IMAGE_EMBED'"); - assertThatCode( - () -> - TransformParser.validateAiModelReferences( - "AI_COMPLETE() AS completed", null, Collections.emptySet())) - .doesNotThrowAnyException(); } private List translate(String expression) { From 92c5cc97c7316c1e0fdb19d3ff15820e4ecb5002 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=A5=E6=A0=96?= Date: Wed, 9 Sep 2026 11:45:29 +0800 Subject: [PATCH 2/2] [FLINK-40572][runtime] Address review comments --- .../cdc/runtime/ai/AiModelClientResolver.java | 42 -------- .../runtime/functions/impl/AiFunctions.java | 79 +++++++-------- .../transform/ProjectionColumnProcessor.java | 9 +- .../TransformExpressionCompiler.java | 6 +- .../transform/TransformFilterProcessor.java | 11 +-- .../cdc/runtime/parser/JaninoCompiler.java | 4 +- .../functions/impl/AiFunctionsTest.java | 95 +++++++++++-------- .../runtime/parser/AiFunctionParserTest.java | 22 ++--- 8 files changed, 109 insertions(+), 159 deletions(-) delete mode 100644 flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java deleted file mode 100644 index 5e6f50419a6..00000000000 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiModelClientResolver.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * 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.cdc.runtime.ai; - -import org.apache.flink.cdc.common.annotation.Internal; -import org.apache.flink.cdc.common.model.AiModelClient; -import org.apache.flink.cdc.common.utils.Preconditions; - -import javax.annotation.Nullable; - -import java.util.Map; - -/** Resolves AI model clients by their logical names during expression evaluation. */ -@Internal -public class AiModelClientResolver { - - private final Map modelClients; - - public AiModelClientResolver(Map modelClients) { - this.modelClients = Preconditions.checkNotNull(modelClients); - } - - @Nullable - public AiModelClient resolve(String modelName) { - return modelClients.get(modelName); - } -} diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java index 71e7eea689d..2d874ac606e 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctions.java @@ -25,13 +25,13 @@ import org.apache.flink.cdc.common.types.RowType; import org.apache.flink.cdc.common.types.variant.BinaryVariant; import org.apache.flink.cdc.common.types.variant.BinaryVariantInternalBuilder; -import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.ai.AiTextFunctionDef; import org.apache.flink.shaded.guava31.com.google.common.primitives.Floats; import java.io.IOException; import java.util.List; +import java.util.Map; /** Built-in AI functions available to transform expressions. */ public class AiFunctions { @@ -44,18 +44,17 @@ public static BinaryVariant aiComplete( String modelName, String input, String systemPrompt, - AiModelClientResolver modelClientResolver) { + Map modelClients) { return generateText( - modelClientResolver, modelName, AiTextFunctionDef.AI_COMPLETE, input, systemPrompt); + modelClients, modelName, AiTextFunctionDef.AI_COMPLETE, input, systemPrompt); } public static BinaryVariant aiClassify( String modelName, String input, String labels, - AiModelClientResolver modelClientResolver) { - return generateText( - modelClientResolver, modelName, AiTextFunctionDef.AI_CLASSIFY, input, labels); + Map modelClients) { + return generateText(modelClients, modelName, AiTextFunctionDef.AI_CLASSIFY, input, labels); } public static BinaryVariant aiTranslate( @@ -63,9 +62,9 @@ public static BinaryVariant aiTranslate( String input, String sourceLang, String targetLang, - AiModelClientResolver modelClientResolver) { + Map modelClients) { return generateText( - modelClientResolver, + modelClients, modelName, AiTextFunctionDef.AI_TRANSLATE, input, @@ -77,36 +76,34 @@ public static BinaryVariant aiSummarize( String modelName, String input, int maxLength, - AiModelClientResolver modelClientResolver) { + Map modelClients) { return generateText( - modelClientResolver, modelName, AiTextFunctionDef.AI_SUMMARIZE, input, maxLength); + modelClients, modelName, AiTextFunctionDef.AI_SUMMARIZE, input, maxLength); } public static BinaryVariant aiSentiment( - String modelName, String input, AiModelClientResolver modelClientResolver) { - return generateText(modelClientResolver, modelName, AiTextFunctionDef.AI_SENTIMENT, input); + String modelName, String input, Map modelClients) { + return generateText(modelClients, modelName, AiTextFunctionDef.AI_SENTIMENT, input); } public static BinaryVariant aiExtract( String modelName, String input, String schema, - AiModelClientResolver modelClientResolver) { - return generateText( - modelClientResolver, modelName, AiTextFunctionDef.AI_EXTRACT, input, schema); + Map modelClients) { + return generateText(modelClients, modelName, AiTextFunctionDef.AI_EXTRACT, input, schema); } public static BinaryVariant aiMask( String modelName, String input, String entities, - AiModelClientResolver modelClientResolver) { - return generateText( - modelClientResolver, modelName, AiTextFunctionDef.AI_MASK, input, entities); + Map modelClients) { + return generateText(modelClients, modelName, AiTextFunctionDef.AI_MASK, input, entities); } private static BinaryVariant generateText( - AiModelClientResolver modelClientResolver, + Map modelClients, String modelName, AiTextFunctionDef function, String input, @@ -116,11 +113,10 @@ private static BinaryVariant generateText( } SupportsTextGeneration model = resolveModel( - modelClientResolver, + modelClients, modelName, function.getFunctionName(), - SupportsTextGeneration.class, - "text generation"); + SupportsTextGeneration.class); String prompt = function.buildPrompt(promptArguments) @@ -143,17 +139,12 @@ private static BinaryVariant generateText( } public static List aiEmbed( - String modelName, String input, AiModelClientResolver modelClientResolver) { + String modelName, String input, Map modelClients) { if (input == null) { return null; } SupportsEmbedding model = - resolveModel( - modelClientResolver, - modelName, - "AI_EMBED", - SupportsEmbedding.class, - "embedding"); + resolveModel(modelClients, modelName, "AI_EMBED", SupportsEmbedding.class); float[] embedding = model.embed(input); return embedding == null ? null : Floats.asList(embedding); } @@ -163,48 +154,42 @@ public static String aiImageComplete( String modelName, byte[] image, String prompt, - AiModelClientResolver modelClientResolver) { + Map modelClients) { if (image == null) { return null; } SupportsImageTextGeneration model = resolveModel( - modelClientResolver, + modelClients, modelName, "AI_IMAGE_COMPLETE", - SupportsImageTextGeneration.class, - "image text generation"); + SupportsImageTextGeneration.class); return model.generateTextFromImage(image, prompt); } /** Dispatches image embedding AI functions. */ public static List aiImageEmbed( - String modelName, byte[] image, AiModelClientResolver modelClientResolver) { + String modelName, byte[] image, Map modelClients) { if (image == null) { return null; } SupportsImageEmbedding model = resolveModel( - modelClientResolver, - modelName, - "AI_IMAGE_EMBED", - SupportsImageEmbedding.class, - "image embedding"); + modelClients, modelName, "AI_IMAGE_EMBED", SupportsImageEmbedding.class); float[] embedding = model.embedImage(image); return embedding == null ? null : Floats.asList(embedding); } private static T resolveModel( - AiModelClientResolver modelClientResolver, + Map modelClients, String modelName, String functionName, - Class requiredCapability, - String capabilityName) { + Class requiredCapability) { if (modelName == null) { throw new IllegalArgumentException( "Model name referenced by " + functionName + " must not be null."); } - AiModelClient model = modelClientResolver.resolve(modelName); + AiModelClient model = modelClients.get(modelName); if (model == null) { throw new IllegalArgumentException( "Model '" @@ -217,11 +202,11 @@ private static T resolveModel( throw new UnsupportedOperationException( "Model '" + modelName - + "' referenced by " + + "' could not be used in " + functionName - + " does not support " - + capabilityName - + "."); + + " because it does not implement " + + requiredCapability.getSimpleName() + + " interface."); } return requiredCapability.cast(model); } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java index fef5896f57d..d4697007ea3 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/ProjectionColumnProcessor.java @@ -21,7 +21,6 @@ import org.apache.flink.cdc.common.model.AiModelClient; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; -import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.codehaus.janino.ExpressionEvaluator; @@ -48,7 +47,7 @@ public class ProjectionColumnProcessor { private final TransformExpressionKey transformExpressionKey; private final Map supportedMetadataColumns; private final List udfFunctionInstances; - private final AiModelClientResolver modelClientResolver; + private final Map modelClients; private final ExpressionEvaluator expressionEvaluator; public ProjectionColumnProcessor( @@ -63,7 +62,7 @@ public ProjectionColumnProcessor( this.projectionColumn = projectionColumn; this.timezone = timezone; this.supportedMetadataColumns = supportedMetadataColumns; - this.modelClientResolver = new AiModelClientResolver(modelClients); + this.modelClients = modelClients; this.transformExpressionKey = generateTransformExpressionKey(); this.expressionEvaluator = TransformExpressionCompiler.compileExpression( @@ -149,8 +148,8 @@ private Object[] generateParams(Object[] rowData, TransformContext context) { // 3 - Add UDF function instances params.addAll(udfFunctionInstances); - // 4 - Add AI model client resolver - params.add(modelClientResolver); + // 4 - Add AI model clients + params.add(modelClients); return params.toArray(); } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java index 6a2e26e2e07..33a70a9ce96 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformExpressionCompiler.java @@ -18,7 +18,6 @@ package org.apache.flink.cdc.runtime.operators.transform; import org.apache.flink.api.common.InvalidProgramException; -import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.operators.transform.exceptions.TransformException; import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.apache.flink.util.FlinkRuntimeException; @@ -33,6 +32,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; /** * The processor of the transform expression. It processes the expression of projections and @@ -70,8 +70,8 @@ public static ExpressionEvaluator compileExpression( argumentClasses.add(Class.forName(udfFunction.getClasspath())); } - argumentNames.add(JaninoCompiler.DEFAULT_AI_MODEL_CLIENT_RESOLVER); - argumentClasses.add(AiModelClientResolver.class); + argumentNames.add(JaninoCompiler.DEFAULT_AI_MODEL_CLIENTS); + argumentClasses.add(Map.class); // Input args expressionEvaluator.setParameters( diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java index 25f1c546107..6e4eac45f36 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/operators/transform/TransformFilterProcessor.java @@ -23,7 +23,6 @@ import org.apache.flink.cdc.common.pipeline.DecimalPrecisionMode; import org.apache.flink.cdc.common.schema.Column; import org.apache.flink.cdc.common.source.SupportedMetadataColumn; -import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.apache.flink.cdc.runtime.parser.JaninoCompiler; import org.apache.flink.cdc.runtime.parser.TransformParser; @@ -52,7 +51,7 @@ public class TransformFilterProcessor { private final DecimalPrecisionMode decimalPrecisionMode; private final List udfFunctionInstances; private final Map supportedMetadataColumns; - private final AiModelClientResolver modelClientResolver; + private final Map modelClients; private final TransformExpressionKey transformExpressionKey; private final ExpressionEvaluator expressionEvaluator; @@ -74,9 +73,7 @@ protected TransformFilterProcessor( this.decimalPrecisionMode = decimalPrecisionMode; this.udfFunctionInstances = udfFunctionInstances; this.supportedMetadataColumns = supportedMetadataColumns; - this.modelClientResolver = - new AiModelClientResolver( - modelClients == null ? Collections.emptyMap() : modelClients); + this.modelClients = modelClients == null ? Collections.emptyMap() : modelClients; if (isNoOp) { this.transformExpressionKey = null; @@ -227,8 +224,8 @@ private Object[] generateParams(Object[] preRow, Object[] postRow, TransformCont // 3 - Add UDF function instances params.addAll(udfFunctionInstances); - // 4 - Add AI model client resolver - params.add(modelClientResolver); + // 4 - Add AI model clients + params.add(modelClients); return params.toArray(); } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java index 81c6bc55a1f..340e7e61eb8 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/JaninoCompiler.java @@ -119,7 +119,7 @@ public class JaninoCompiler { public static final String DEFAULT_EPOCH_TIME = "__epoch_time__"; public static final String DEFAULT_TIME_ZONE = "__time_zone__"; - public static final String DEFAULT_AI_MODEL_CLIENT_RESOLVER = "__ai_model_client_resolver__"; + public static final String DEFAULT_AI_MODEL_CLIENTS = "__ai_model_clients__"; private static final String[] BUILTIN_FUNCTION_MODULES = { "Ai", "Arithmetic", "Casting", "Comparison", "Logical", "String", "Struct", "Temporal" @@ -268,7 +268,7 @@ private static Java.Rvalue translateSqlBasicCall(Context context, SqlBasicCall s } else if (isAiFunction(functionName)) { atoms.add( new Java.AmbiguousName( - Location.NOWHERE, new String[] {DEFAULT_AI_MODEL_CLIENT_RESOLVER})); + Location.NOWHERE, new String[] {DEFAULT_AI_MODEL_CLIENTS})); } return sqlBasicCallToJaninoRvalue(context, sqlBasicCall, atoms.toArray(new Java.Rvalue[0])); } diff --git a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java index 05510843723..eac36122833 100644 --- a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java +++ b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/functions/impl/AiFunctionsTest.java @@ -22,7 +22,6 @@ import org.apache.flink.cdc.common.model.abilities.SupportsImageEmbedding; import org.apache.flink.cdc.common.model.abilities.SupportsImageTextGeneration; import org.apache.flink.cdc.common.model.abilities.SupportsTextGeneration; -import org.apache.flink.cdc.runtime.ai.AiModelClientResolver; import org.junit.jupiter.api.Test; @@ -92,21 +91,23 @@ private static class UnsupportedModelClient implements AiModelClient { @Test void testTextAiFunctionsUseEnglishPromptsAndParseJsonResponses() { TestModelClient model = new TestModelClient(); - AiModelClientResolver resolver = resolver("textModel", model); + Map modelClients = modelClients("textModel", model); - assertThat(AiFunctions.aiComplete("textModel", "input", "Return three letters", resolver)) + assertThat( + AiFunctions.aiComplete( + "textModel", "input", "Return three letters", modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiClassify("textModel", "input", "positive,negative", resolver)) + assertThat(AiFunctions.aiClassify("textModel", "input", "positive,negative", modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiTranslate("textModel", "input", "auto", "en", resolver)) + assertThat(AiFunctions.aiTranslate("textModel", "input", "auto", "en", modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiSummarize("textModel", "input", 100, resolver)) + assertThat(AiFunctions.aiSummarize("textModel", "input", 100, modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiSentiment("textModel", "input", resolver)) + assertThat(AiFunctions.aiSentiment("textModel", "input", modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiExtract("textModel", "input", "name:string", resolver)) + assertThat(AiFunctions.aiExtract("textModel", "input", "name:string", modelClients)) .hasToString("{\"result\":\"ABC\"}"); - assertThat(AiFunctions.aiMask("textModel", "input", "email,phone", resolver)) + assertThat(AiFunctions.aiMask("textModel", "input", "email,phone", modelClients)) .hasToString("{\"result\":\"ABC\"}"); assertThat(model.prompts).hasSize(7); @@ -133,9 +134,9 @@ void testTextAiFunctionsUseEnglishPromptsAndParseJsonResponses() { @Test void testEmbeddingFunction() { TestModelClient model = new TestModelClient(); - AiModelClientResolver resolver = resolver("embeddingModel", model); + Map modelClients = modelClients("embeddingModel", model); - assertThat(AiFunctions.aiEmbed("embeddingModel", "input", resolver)) + assertThat(AiFunctions.aiEmbed("embeddingModel", "input", modelClients)) .containsExactly(0.1f, 0.2f, 0.3f); assertThat(model.embedCalls).isOne(); } @@ -143,12 +144,14 @@ void testEmbeddingFunction() { @Test void testImageAiFunctions() { TestModelClient model = new TestModelClient(); - AiModelClientResolver resolver = resolver("imageModel", model); + Map modelClients = modelClients("imageModel", model); byte[] image = new byte[] {1, 2, 3, 4}; - assertThat(AiFunctions.aiImageComplete("imageModel", image, "Describe the image", resolver)) + assertThat( + AiFunctions.aiImageComplete( + "imageModel", image, "Describe the image", modelClients)) .isEqualTo("image has 4 bytes, prompt: Describe the image"); - assertThat(AiFunctions.aiImageEmbed("imageModel", image, resolver)) + assertThat(AiFunctions.aiImageEmbed("imageModel", image, modelClients)) .containsExactly(0.9f, 0.8f, 0.7f); assertThat(model.imageTextCalls).isOne(); assertThat(model.imageEmbedCalls).isOne(); @@ -157,44 +160,51 @@ void testImageAiFunctions() { @Test void testUnsupportedCapabilities() { UnsupportedModelClient model = new UnsupportedModelClient(); - AiModelClientResolver resolver = resolver("unsupported", model); + Map modelClients = modelClients("unsupported", model); - assertThatThrownBy(() -> AiFunctions.aiComplete("unsupported", "input", "prompt", resolver)) + assertThatThrownBy( + () -> + AiFunctions.aiComplete( + "unsupported", "input", "prompt", modelClients)) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining("Model 'unsupported'") .hasMessageContaining("AI_COMPLETE") - .hasMessageContaining("does not support text generation"); - assertThatThrownBy(() -> AiFunctions.aiEmbed("unsupported", "input", resolver)) + .hasMessageContaining("does not implement SupportsTextGeneration interface"); + assertThatThrownBy(() -> AiFunctions.aiEmbed("unsupported", "input", modelClients)) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining("Model 'unsupported'") .hasMessageContaining("AI_EMBED") - .hasMessageContaining("does not support embedding"); + .hasMessageContaining("does not implement SupportsEmbedding interface"); assertThatThrownBy( () -> AiFunctions.aiImageComplete( - "unsupported", new byte[] {1, 2}, "describe", resolver)) + "unsupported", new byte[] {1, 2}, "describe", modelClients)) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining("Model 'unsupported'") .hasMessageContaining("AI_IMAGE_COMPLETE") - .hasMessageContaining("does not support image text generation"); + .hasMessageContaining("does not implement SupportsImageTextGeneration interface"); assertThatThrownBy( - () -> AiFunctions.aiImageEmbed("unsupported", new byte[] {1, 2}, resolver)) + () -> + AiFunctions.aiImageEmbed( + "unsupported", new byte[] {1, 2}, modelClients)) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining("Model 'unsupported'") .hasMessageContaining("AI_IMAGE_EMBED") - .hasMessageContaining("does not support image embedding"); + .hasMessageContaining("does not implement SupportsImageEmbedding interface"); } @Test void testInvalidModelName() { - AiModelClientResolver resolver = new AiModelClientResolver(Collections.emptyMap()); + Map modelClients = Collections.emptyMap(); assertThatThrownBy( - () -> AiFunctions.aiComplete("missingModel", "input", "prompt", resolver)) + () -> + AiFunctions.aiComplete( + "missingModel", "input", "prompt", modelClients)) .isInstanceOf(IllegalArgumentException.class) .hasMessage( "Model 'missingModel' referenced by AI_COMPLETE has not been declared."); - assertThatThrownBy(() -> AiFunctions.aiEmbed(null, "input", resolver)) + assertThatThrownBy(() -> AiFunctions.aiEmbed(null, "input", modelClients)) .isInstanceOf(IllegalArgumentException.class) .hasMessage("Model name referenced by AI_EMBED must not be null."); } @@ -202,12 +212,12 @@ void testInvalidModelName() { @Test void testInvalidJsonResponse() { TestModelClient model = new TestModelClient("not-json"); - AiModelClientResolver resolver = resolver("model", model); + Map modelClients = modelClients("model", model); assertThatThrownBy( () -> AiFunctions.aiClassify( - "model", "input", "positive,negative", resolver)) + "model", "input", "positive,negative", modelClients)) .isInstanceOf(RuntimeException.class) .hasMessage("AI function AI_CLASSIFY returned invalid JSON: not-json"); } @@ -216,12 +226,12 @@ void testInvalidJsonResponse() { void testInvalidJsonResponseIsTruncated() { String longInvalidJson = "x".repeat(600); TestModelClient model = new TestModelClient(longInvalidJson); - AiModelClientResolver resolver = resolver("model", model); + Map modelClients = modelClients("model", model); assertThatThrownBy( () -> AiFunctions.aiClassify( - "model", "input", "positive,negative", resolver)) + "model", "input", "positive,negative", modelClients)) .isInstanceOf(RuntimeException.class) .hasMessage( "AI function AI_CLASSIFY returned invalid JSON: " @@ -232,14 +242,15 @@ void testInvalidJsonResponseIsTruncated() { @Test void testNullInputSkipsModelInvocation() { TestModelClient model = new TestModelClient(); - AiModelClientResolver resolver = resolver("model", model); - AiModelClientResolver emptyResolver = new AiModelClientResolver(Collections.emptyMap()); - - assertThat(AiFunctions.aiClassify("model", null, "positive,negative", resolver)).isNull(); - assertThat(AiFunctions.aiEmbed("model", null, resolver)).isNull(); - assertThat(AiFunctions.aiImageComplete("model", null, "describe", resolver)).isNull(); - assertThat(AiFunctions.aiImageEmbed("model", null, resolver)).isNull(); - assertThat(AiFunctions.aiComplete("missing", null, "prompt", emptyResolver)).isNull(); + Map modelClients = modelClients("model", model); + Map emptyModelClients = Collections.emptyMap(); + + assertThat(AiFunctions.aiClassify("model", null, "positive,negative", modelClients)) + .isNull(); + assertThat(AiFunctions.aiEmbed("model", null, modelClients)).isNull(); + assertThat(AiFunctions.aiImageComplete("model", null, "describe", modelClients)).isNull(); + assertThat(AiFunctions.aiImageEmbed("model", null, modelClients)).isNull(); + assertThat(AiFunctions.aiComplete("missing", null, "prompt", emptyModelClients)).isNull(); assertThat(model.prompts).isEmpty(); assertThat(model.embedCalls).isZero(); assertThat(model.imageTextCalls).isZero(); @@ -249,13 +260,13 @@ void testNullInputSkipsModelInvocation() { @Test void testNullModelResponseReturnsNull() { TestModelClient model = new TestModelClient(null); - AiModelClientResolver resolver = resolver("model", model); + Map modelClients = modelClients("model", model); - assertThat(AiFunctions.aiSummarize("model", "input", 100, resolver)).isNull(); + assertThat(AiFunctions.aiSummarize("model", "input", 100, modelClients)).isNull(); assertThat(model.prompts).hasSize(1); } - private static AiModelClientResolver resolver(String modelName, AiModelClient model) { - return new AiModelClientResolver(Map.of(modelName, model)); + private static Map modelClients(String modelName, AiModelClient model) { + return Map.of(modelName, model); } } diff --git a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java index 2d1fea5dcbd..6f60d0883e7 100644 --- a/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java +++ b/flink-cdc-runtime/src/test/java/org/apache/flink/cdc/runtime/parser/AiFunctionParserTest.java @@ -50,8 +50,8 @@ void testTranslateAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiComplete(\"completer\", $0, \"You are helpful\", __ai_model_client_resolver__)", - "aiEmbed(\"embedder\", $0, __ai_model_client_resolver__)"); + "aiComplete(\"completer\", $0, \"You are helpful\", __ai_model_clients__)", + "aiEmbed(\"embedder\", $0, __ai_model_clients__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsExactly(DataTypes.VARIANT(), DataTypes.ARRAY(DataTypes.FLOAT())); @@ -71,18 +71,18 @@ void testTranslateSpecializedTextAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiClassify(\"model\", $0, \"positive,negative\", __ai_model_client_resolver__)", + "aiClassify(\"model\", $0, \"positive,negative\", __ai_model_clients__)", "aiTranslate(\n" + " \"model\",\n" + " $0,\n" + " \"auto\",\n" + " \"en\",\n" - + " __ai_model_client_resolver__\n" + + " __ai_model_clients__\n" + ")", - "aiSummarize(\"model\", $0, 100, __ai_model_client_resolver__)", - "aiSentiment(\"model\", $0, __ai_model_client_resolver__)", - "aiExtract(\"model\", $0, \"name:string\", __ai_model_client_resolver__)", - "aiMask(\"model\", $0, \"email,phone\", __ai_model_client_resolver__)"); + "aiSummarize(\"model\", $0, 100, __ai_model_clients__)", + "aiSentiment(\"model\", $0, __ai_model_clients__)", + "aiExtract(\"model\", $0, \"name:string\", __ai_model_clients__)", + "aiMask(\"model\", $0, \"email,phone\", __ai_model_clients__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsOnly(DataTypes.VARIANT()); @@ -98,8 +98,8 @@ void testTranslateImageAiFunctions() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiImageComplete(\"vision\", $0, \"Describe the image\", __ai_model_client_resolver__)", - "aiImageEmbed(\"imageEmbedder\", $0, __ai_model_client_resolver__)"); + "aiImageComplete(\"vision\", $0, \"Describe the image\", __ai_model_clients__)", + "aiImageEmbed(\"imageEmbedder\", $0, __ai_model_clients__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsExactly(DataTypes.STRING(), DataTypes.ARRAY(DataTypes.FLOAT())); @@ -114,7 +114,7 @@ void testDynamicModelSelection() { assertThat(columns) .extracting(ProjectionColumn::getScriptExpression) .containsExactly( - "aiComplete(isTrue(valueEquals($0, 1)) ? \"powerful\" : \"cheap\", $1, \"prompt\", __ai_model_client_resolver__)"); + "aiComplete(isTrue(valueEquals($0, 1)) ? \"powerful\" : \"cheap\", $1, \"prompt\", __ai_model_clients__)"); assertThat(columns) .extracting(ProjectionColumn::getDataType) .containsExactly(DataTypes.VARIANT());