From 0de0821b3fef4e7c1092b6e487c17a99f609be73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=A5=E6=A0=96?= Date: Wed, 19 Aug 2026 11:50:38 +0800 Subject: [PATCH] [FLINK-40416][runtime] Add built-in text AI functions --- docs/content.zh/docs/core-concept/ai-model.md | 28 +++- docs/content/docs/core-concept/ai-model.md | 28 +++- .../flink/translator/TransformTranslator.java | 43 +++++- .../flink/FlinkPipelineAiFunctionITCase.java | 49 +++++- .../pipeline/tests/AiFunctionE2eITCase.java | 24 +-- .../cdc/models/dummy/DummyModelClient.java | 18 +++ .../cdc/runtime/ai/AiTextFunctionDef.java | 65 +++++++- .../runtime/functions/impl/AiFunctions.java | 49 +++++- .../cdc/runtime/parser/TransformParser.java | 133 +++++++++++++++- .../functions/impl/AiFunctionsTest.java | 91 +++++++++-- .../runtime/parser/AiFunctionParserTest.java | 145 ++++++++++++++++++ 11 files changed, 621 insertions(+), 52 deletions(-) diff --git a/docs/content.zh/docs/core-concept/ai-model.md b/docs/content.zh/docs/core-concept/ai-model.md index 9b802932653..0a36d17fa56 100644 --- a/docs/content.zh/docs/core-concept/ai-model.md +++ b/docs/content.zh/docs/core-concept/ai-model.md @@ -24,11 +24,30 @@ under the License. # AI 模型 -AI 模型可用于 transform 表达式中的文本补全和 embedding。 +AI 模型可用于 transform 表达式中的文本生成、文本分析和 embedding。 + +## AI Functions + +模型名称必须是字符串常量，并引用 `pipeline.model` 中声明的模型。文本函数要求模型客户端支持文本生成，`AI_EMBED` 要求模型客户端支持 embedding；Pipeline 会在执行前校验引用模型的 capability 是否匹配。 + +所有文本函数都会将模型返回的 JSON 解析为 `VARIANT`。 + +| 函数 | 说明 | JSON 字段 | +|------|------|-----------| +| `AI_COMPLETE(model, input, system_prompt)` | 使用调用方提供的 system prompt 补全文本。 | `result` | +| `AI_CLASSIFY(model, input, labels)` | 将输入分类到给定标签之一。 | `category`、`confidence` | +| `AI_TRANSLATE(model, input, source_lang, target_lang)` | 翻译输入；`source_lang` 可传 `auto` 自动识别。 | `translated_text`、`detected_language` | +| `AI_SUMMARIZE(model, input, max_length)` | 在指定字符数内生成摘要。 | `summary` | +| `AI_SENTIMENT(model, input)` | 分析文本情感。 | `score`、`label`、`confidence` | +| `AI_EXTRACT(model, input, schema)` | 按 schema 字符串描述提取字段。 | `extracted_json` | +| `AI_MASK(model, input, entities)` | 对指定实体类型进行脱敏。 | `masked_text`、`detected_entities` | +| `AI_EMBED(model, input)` | 生成 embedding 向量。 | 不返回 JSON，而是返回 `ARRAY`。 | + +六个专用文本函数使用内置英文 prompt 模板，但输入文本可以是任意语言。输入为 `NULL` 时直接返回 `NULL`，且不会调用模型；模型返回 `NULL` 时也返回 `NULL`。非空文本响应必须是语法合法的 JSON，否则当前记录处理失败，错误信息会标明具体 AI 函数。运行时只校验 JSON 语法，不校验响应字段是否存在或字段类型是否匹配。 ## OpenAI-compatible 模型客户端 -AI 模型客户端可供 transform 中的 `AI_COMPLETE` 和 `AI_EMBED` 函数引用。使用时，需要通过 `--jar` 将模型实现 JAR（例如 `flink-cdc-pipeline-model-openai-compatible`）添加到 Pipeline 命令中。 +AI 模型客户端可供上述 AI Functions 引用。使用时，需要通过 `--jar` 将模型实现 JAR（例如 `flink-cdc-pipeline-model-openai-compatible`）添加到 Pipeline 命令中。 OpenAI-compatible 客户端支持调用实现 OpenAI Chat Completions 和 Embeddings REST API 的服务。 @@ -40,6 +59,7 @@ transform: projection: >- *, AI_COMPLETE('completion_model', content, '总结输入内容') AS summary, + AI_SENTIMENT('completion_model', content) AS sentiment, AI_EMBED('embedding_model', content) AS embedding pipeline: @@ -71,11 +91,11 @@ pipeline: | `model` | 是 | 发送给服务端的模型名称；`model-name` 作为废弃别名仍可使用。 | | `endpoint` | 是 | OpenAI-compatible 服务的 Base URL。 | | `api-key` | 是 | 请求认证使用的 Bearer Token。 | -| `system-prompt` | 否 | 添加在 `AI_COMPLETE` 生成的 system prompt 之前。 | +| `system-prompt` | 否 | 添加在所有文本 AI Function 的 prompt 之前。 | | `user-prompt` | 否 | 在输入之后追加一条 user message。 | | `temperature`、`top-p`、`stop`、`max-tokens` | 否 | 常用文本生成参数。 | | `presence-penalty`、`frequency-penalty`、`n`、`seed` | 否 | 其他文本生成参数。 | -| `response-format` | 否 | 支持 `json_object`；AI completion 的结果必须是合法 JSON。 | +| `response-format` | 否 | 支持 `json_object`；文本 AI Function 的结果必须是合法 JSON。 | | `content-type` | 否 | `text`（默认）或 `image_url`。 | | `dimension` | 否 | 请求的 embedding 维度。 | | `extra-header`、`extra-body` | 否 | JSON 对象格式的厂商自定义请求头或请求体字段。 | diff --git a/docs/content/docs/core-concept/ai-model.md b/docs/content/docs/core-concept/ai-model.md index da3ef0f0083..8aa7666838a 100644 --- a/docs/content/docs/core-concept/ai-model.md +++ b/docs/content/docs/core-concept/ai-model.md @@ -24,11 +24,30 @@ under the License. # AI Model -AI models can be used in transform expressions for text completion and embedding. +AI models can be used in transform expressions for text generation, text analysis, and embedding. + +## 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 `AI_EMBED` requires embedding support. The pipeline validates the referenced model capability before execution. + +All text functions return `VARIANT` values parsed from the model's JSON response. + +| Function | Description | JSON fields | +|----------|-------------|-------------| +| `AI_COMPLETE(model, input, system_prompt)` | Completes the input using a caller-provided system prompt. | `result` | +| `AI_CLASSIFY(model, input, labels)` | Classifies the input into one of the provided labels. | `category`, `confidence` | +| `AI_TRANSLATE(model, input, source_lang, target_lang)` | Translates the input. Use `auto` to detect the source language. | `translated_text`, `detected_language` | +| `AI_SUMMARIZE(model, input, max_length)` | Summarizes the input within the requested character limit. | `summary` | +| `AI_SENTIMENT(model, input)` | Analyzes sentiment. | `score`, `label`, `confidence` | +| `AI_EXTRACT(model, input, schema)` | Extracts fields described by the schema string. | `extracted_json` | +| `AI_MASK(model, input, entities)` | Masks the requested entity types. | `masked_text`, `detected_entities` | +| `AI_EMBED(model, input)` | Creates an embedding vector. | Returns `ARRAY` instead of JSON. | + +The specialized text functions use built-in English prompt templates, but their input may be in any language. If the input is `NULL`, the function returns `NULL` without invoking the model. A `NULL` model response also produces `NULL`. A non-null text response must be syntactically valid JSON; otherwise, record processing fails with an error that identifies the AI function. The runtime validates JSON syntax but does not validate the presence or types of individual response fields. ## OpenAI-compatible Model Client -AI model clients can be referenced by the `AI_COMPLETE` and `AI_EMBED` transform functions. Add the model implementation JAR, such as `flink-cdc-pipeline-model-openai-compatible`, to the pipeline command with `--jar`. +AI model clients can be referenced by the AI functions above. Add the model implementation JAR, such as `flink-cdc-pipeline-model-openai-compatible`, to the pipeline command with `--jar`. The OpenAI-compatible client supports chat completions and text embeddings against endpoints that implement the corresponding OpenAI REST APIs. @@ -40,6 +59,7 @@ transform: projection: >- *, AI_COMPLETE('completion_model', content, 'Summarize the input') AS summary, + AI_SENTIMENT('completion_model', content) AS sentiment, AI_EMBED('embedding_model', content) AS embedding pipeline: @@ -71,11 +91,11 @@ Do not store API keys in source control. Supply them through the secret-manageme | `model` | Yes | Model name sent to the endpoint. `model-name` is accepted as a deprecated alias. | | `endpoint` | Yes | Base URL of the OpenAI-compatible endpoint. | | `api-key` | Yes | Bearer token used to authenticate requests. | -| `system-prompt` | No | Prompt prepended to the system prompt generated by `AI_COMPLETE`. | +| `system-prompt` | No | Prompt prepended to every text AI function prompt. | | `user-prompt` | No | Additional user message appended after the input. | | `temperature`, `top-p`, `stop`, `max-tokens` | No | Common generation parameters. | | `presence-penalty`, `frequency-penalty`, `n`, `seed` | No | Additional generation parameters. | -| `response-format` | No | `json_object` is supported. AI completion results must be valid JSON. | +| `response-format` | No | `json_object` is supported. Text AI function results must be valid JSON. | | `content-type` | No | `text` (default) or `image_url`. | | `dimension` | No | Requested embedding dimension. | | `extra-header`, `extra-body` | No | Provider-specific headers or body fields encoded as JSON objects. | 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 20c8bc54eef..5dc9b268f73 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 @@ -41,6 +41,7 @@ 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; @@ -66,7 +67,8 @@ public DataStream translatePreTransform( if (transforms.isEmpty()) { return input; } - validateModelReferences(transforms, models); + validateModelReferences( + transforms, models, getUserDefinedFunctionNames(udfFunctions, models)); return input.transform( "Transform:Schema", new EventTypeInfo(), @@ -144,7 +146,10 @@ public DataStream translatePostTransform( .filter(ModelDef::isLegacy) .map(this::modelToUDFTuple) .collect(Collectors.toList())); - postTransformFunctionBuilder.addModelClients(loadModelClients(models, env)); + Map modelClients = loadModelClients(models, env); + validateModelCapabilities( + transforms, modelClients, getUserDefinedFunctionNames(udfFunctions, models)); + postTransformFunctionBuilder.addModelClients(modelClients); return input.transform( "Transform:Data", new EventTypeInfo(), postTransformFunctionBuilder.build()) .uid(operatorUidGenerator.generateUid("post-transform")); @@ -186,7 +191,10 @@ private Map loadModelClients( return clients; } - private void validateModelReferences(List transforms, List models) { + private void validateModelReferences( + List transforms, + List models, + Set userDefinedFunctionNames) { Set clientModelNames = models.stream() .filter(model -> !model.isLegacy()) @@ -194,10 +202,37 @@ private void validateModelReferences(List transforms, 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 f01777540e1..d50f1c67d12 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 @@ -34,6 +34,7 @@ import org.apache.flink.cdc.composer.definition.SinkDef; import org.apache.flink.cdc.composer.definition.SourceDef; import org.apache.flink.cdc.composer.definition.TransformDef; +import org.apache.flink.cdc.composer.definition.UdfDef; import org.apache.flink.cdc.connectors.values.ValuesDatabase; import org.apache.flink.cdc.connectors.values.factory.ValuesDataFactory; import org.apache.flink.cdc.connectors.values.sink.ValuesDataSinkOptions; @@ -127,6 +128,42 @@ void testAiCompleteInProjection() throws Exception { "Dummy model closed."); } + @Test + void testSpecializedTextAiFunctionsInProjection() throws Exception { + String[] output = + runAiFunctionTest( + "id, " + + "AI_CLASSIFY('testModel', content, 'positive,negative') AS classified, " + + "AI_TRANSLATE('testModel', content, 'auto', 'en') AS translated, " + + "AI_SUMMARIZE('testModel', content, 100) AS summarized, " + + "AI_SENTIMENT('testModel', content) AS sentiment, " + + "AI_EXTRACT('testModel', content, 'name:string') AS extracted, " + + "AI_MASK('testModel', content, 'name') AS masked", + List.of(ModelDef.of("testModel", "dummy", Collections.emptyMap()))); + + assertThat(output) + .containsExactly( + "CreateTableEvent{tableId=default_namespace.default_schema.mytable1, schema=columns={`id` INT NOT NULL,`classified` VARIANT,`translated` VARIANT,`summarized` VARIANT,`sentiment` VARIANT,`extracted` VARIANT,`masked` VARIANT}, primaryKeys=id, options=()}", + "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, {\"category\":\"dummy\",\"confidence\":1}, {\"detected_language\":\"en\",\"translated_text\":\"dummy translation\"}, {\"summary\":\"dummy summary\"}, {\"confidence\":1,\"label\":\"neutral\",\"score\":0}, {\"extracted_json\":{\"name\":\"dummy\"}}, {\"detected_entities\":\"name\",\"masked_text\":\"d***y\"}], op=INSERT, meta=()}"); + } + + @Test + void testSameNamedUdfTakesPrecedenceOverAiFunction() throws Exception { + String[] output = + runAiFunctionTest( + "id, AI_SENTIMENT(id) AS sentiment", + List.of(ModelDef.of("unusedModel", "dummy", Collections.emptyMap())), + List.of( + new UdfDef( + "ai_sentiment", + "org.apache.flink.cdc.udf.examples.java.AddOneFunctionClass"))); + + assertThat(output) + .containsExactly( + "CreateTableEvent{tableId=default_namespace.default_schema.mytable1, schema=columns={`id` INT NOT NULL,`sentiment` STRING}, primaryKeys=id, options=()}", + "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, 2], op=INSERT, meta=()}"); + } + @Test void testAiEmbedInProjection() throws Exception { String[] output = @@ -140,18 +177,24 @@ void testAiEmbedInProjection() throws Exception { } private String[] runAiFunctionTest(String projection, List models) throws Exception { + return runAiFunctionTest(projection, models, Collections.emptyList()); + } + + private String[] runAiFunctionTest( + String projection, List models, List udfFunctions) throws Exception { URL modelJar = createDummyModelJar().toUri().toURL(); ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader(); try (URLClassLoader modelClassLoader = new DummyModelClassLoader(modelJar, originalClassLoader)) { Thread.currentThread().setContextClassLoader(modelClassLoader); - return runAiFunctionTest(projection, models, modelJar); + return runAiFunctionTest(projection, models, udfFunctions, modelJar); } finally { Thread.currentThread().setContextClassLoader(originalClassLoader); } } - private String[] runAiFunctionTest(String projection, List models, URL modelJar) + private String[] runAiFunctionTest( + String projection, List models, List udfFunctions, URL modelJar) throws Exception { FlinkPipelineComposer composer = FlinkPipelineComposer.ofMiniCluster(); @@ -212,7 +255,7 @@ private String[] runAiFunctionTest(String projection, List models, URL sinkDef, Collections.emptyList(), Collections.singletonList(transformDef), - Collections.emptyList(), + udfFunctions, models, pipelineConfig); diff --git a/flink-cdc-e2e-tests/flink-cdc-pipeline-e2e-tests/src/test/java/org/apache/flink/cdc/pipeline/tests/AiFunctionE2eITCase.java b/flink-cdc-e2e-tests/flink-cdc-pipeline-e2e-tests/src/test/java/org/apache/flink/cdc/pipeline/tests/AiFunctionE2eITCase.java index 8bdb7cb9a94..47cbf7b9a48 100644 --- a/flink-cdc-e2e-tests/flink-cdc-pipeline-e2e-tests/src/test/java/org/apache/flink/cdc/pipeline/tests/AiFunctionE2eITCase.java +++ b/flink-cdc-e2e-tests/flink-cdc-pipeline-e2e-tests/src/test/java/org/apache/flink/cdc/pipeline/tests/AiFunctionE2eITCase.java @@ -30,7 +30,13 @@ class AiFunctionE2eITCase extends PipelineTestEnvironment { private static final String TABLE_1 = "default_namespace.default_schema.table1"; private static final String TABLE_2 = "default_namespace.default_schema.table2"; - private static final String DUMMY_JSON = "{\"result\":\"dummy response\"}"; + private static final String DUMMY_TEXT_RESULTS = + "{\"category\":\"dummy\",\"confidence\":1}, " + + "{\"detected_language\":\"en\",\"translated_text\":\"dummy translation\"}, " + + "{\"summary\":\"dummy summary\"}, " + + "{\"confidence\":1,\"label\":\"neutral\",\"score\":0}, " + + "{\"extracted_json\":{\"name\":\"dummy\"}}, " + + "{\"detected_entities\":\"name\",\"masked_text\":\"d***y\"}"; private static final String DUMMY_EMBEDDING = "[3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0, 6.0]"; @Test @@ -47,7 +53,7 @@ void testAiFunctionsWithDummyModel() throws Exception { + " - source-table: " + TABLE_1 + "\n" - + " projection: col1, AI_COMPLETE('myModel', col1, 'Complete it') AS completed\n" + + " projection: col1, AI_CLASSIFY('myModel', col1, 'a,b') AS classified, AI_TRANSLATE('myModel', col1, 'auto', 'en') AS translated, AI_SUMMARIZE('myModel', col1, 100) AS summarized, AI_SENTIMENT('myModel', col1) AS sentiment, AI_EXTRACT('myModel', col1, 'name:string') AS extracted, AI_MASK('myModel', col1, 'name') AS masked\n" + " - source-table: " + TABLE_2 + "\n" @@ -70,33 +76,33 @@ void testAiFunctionsWithDummyModel() throws Exception { validateResult( "CreateTableEvent{tableId=" + TABLE_1 - + ", schema=columns={`col1` STRING NOT NULL,`completed` VARIANT}, primaryKeys=col1, options=()}", + + ", schema=columns={`col1` STRING NOT NULL,`classified` VARIANT,`translated` VARIANT,`summarized` VARIANT,`sentiment` VARIANT,`extracted` VARIANT,`masked` VARIANT}, primaryKeys=col1, options=()}", "DataChangeEvent{tableId=" + TABLE_1 + ", before=[], after=[1, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], op=INSERT, meta=()}", "DataChangeEvent{tableId=" + TABLE_1 + ", before=[], after=[2, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], op=INSERT, meta=()}", "DataChangeEvent{tableId=" + TABLE_1 + ", before=[], after=[3, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], op=INSERT, meta=()}", "DataChangeEvent{tableId=" + TABLE_1 + ", before=[1, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], after=[], op=DELETE, meta=()}", "DataChangeEvent{tableId=" + TABLE_1 + ", before=[2, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], after=[2, " - + DUMMY_JSON + + DUMMY_TEXT_RESULTS + "], op=UPDATE, meta=()}"); validateResult( diff --git a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-dummy/src/main/java/org/apache/flink/cdc/models/dummy/DummyModelClient.java b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-dummy/src/main/java/org/apache/flink/cdc/models/dummy/DummyModelClient.java index f4c2b7444fc..d92c0a8afb7 100644 --- a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-dummy/src/main/java/org/apache/flink/cdc/models/dummy/DummyModelClient.java +++ b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-dummy/src/main/java/org/apache/flink/cdc/models/dummy/DummyModelClient.java @@ -37,6 +37,24 @@ public String generate(String systemPrompt, String userInput) { if (debug) { System.out.printf("Received prompt: %s%nUser input: %s%n", systemPrompt, userInput); } + if (systemPrompt.contains("\"category\"")) { + return "{\"category\":\"dummy\",\"confidence\":1.0}"; + } + if (systemPrompt.contains("\"translated_text\"")) { + return "{\"translated_text\":\"dummy translation\",\"detected_language\":\"en\"}"; + } + if (systemPrompt.contains("\"summary\"")) { + return "{\"summary\":\"dummy summary\"}"; + } + if (systemPrompt.contains("\"score\"")) { + return "{\"score\":0.0,\"label\":\"neutral\",\"confidence\":1.0}"; + } + if (systemPrompt.contains("\"extracted_json\"")) { + return "{\"extracted_json\":{\"name\":\"dummy\"}}"; + } + if (systemPrompt.contains("\"masked_text\"")) { + return "{\"masked_text\":\"d***y\",\"detected_entities\":\"name\"}"; + } return "{\"result\":\"dummy response\"}"; } diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiTextFunctionDef.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiTextFunctionDef.java index 74687617aff..cbb8fc86cd6 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiTextFunctionDef.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiTextFunctionDef.java @@ -21,13 +21,74 @@ import org.apache.flink.cdc.common.types.DataTypes; import org.apache.flink.cdc.common.types.RowType; +import java.util.Locale; + /** Built-in AI text generation function definitions. */ public enum AiTextFunctionDef { AI_COMPLETE( "AI_COMPLETE", RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"systemPrompt"}), RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"result"}), - "%s\n"); + "%s\n"), + + AI_CLASSIFY( + "AI_CLASSIFY", + RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"labels"}), + RowType.of( + new DataType[] {DataTypes.STRING(), DataTypes.DOUBLE()}, + new String[] {"category", "confidence"}), + "You are a text classifier. Classify the input into exactly one of these labels: %s.\n" + + "Choose only a provided label. Use the dominant meaning when multiple labels apply, " + + "and lower the confidence when no label is a good match.\n"), + + AI_TRANSLATE( + "AI_TRANSLATE", + RowType.of( + new DataType[] {DataTypes.STRING(), DataTypes.STRING()}, + new String[] {"sourceLang", "targetLang"}), + RowType.of( + new DataType[] {DataTypes.STRING(), DataTypes.STRING()}, + new String[] {"translated_text", "detected_language"}), + "You are a translator. Translate the input from %s to %s while preserving its meaning, " + + "formatting, and terminology. If the source language is auto, detect it and report " + + "the detected language code.\n"), + + AI_SUMMARIZE( + "AI_SUMMARIZE", + RowType.of(new DataType[] {DataTypes.INT()}, new String[] {"maxLength"}), + RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"summary"}), + "You are a text summarizer. Summarize the input in no more than %d characters. Preserve " + + "the key facts and conclusions, remove redundancy, and avoid subjective commentary.\n"), + + AI_SENTIMENT( + "AI_SENTIMENT", + RowType.of(new DataType[0], new String[0]), + RowType.of( + new DataType[] {DataTypes.DOUBLE(), DataTypes.STRING(), DataTypes.DOUBLE()}, + new String[] {"score", "label", "confidence"}), + "You are a sentiment analyzer. Analyze the input in context. Return a score from -1.0 " + + "(most negative) to 1.0 (most positive), a label of positive, negative, or neutral, " + + "and a confidence from 0.0 to 1.0. Consider tone, negation, and sarcasm.\n"), + + AI_EXTRACT( + "AI_EXTRACT", + RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"schema"}), + RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"extracted_json"}), + "You are an information extraction system. Extract information from the input according " + + "to this schema: %s. Preserve the requested field names and types, use null for " + + "missing values, and place the extracted JSON object in extracted_json. Supported " + + "types include string, number, integer, boolean, array, object, date, datetime, " + + "email, and phone.\n"), + + AI_MASK( + "AI_MASK", + RowType.of(new DataType[] {DataTypes.STRING()}, new String[] {"entities"}), + RowType.of( + new DataType[] {DataTypes.STRING(), DataTypes.STRING()}, + new String[] {"masked_text", "detected_entities"}), + "You are a data masking system. Detect and consistently mask these entity types in the " + + "input: %s. Preserve the usefulness and structure of non-sensitive content, and " + + "report the entities that were detected.\n"); private final String functionName; private final RowType inputType; @@ -56,6 +117,6 @@ public RowType getOutputType() { } public String buildPrompt(Object... args) { - return String.format(promptTemplate, args); + return String.format(Locale.ROOT, promptTemplate, args); } } 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 7d17c86c4f8..24783b94fb6 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 @@ -36,14 +36,49 @@ 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 aiClassify(AiModelClient model, String input, String labels) { + return generateText(model, 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); + } + + public static BinaryVariant aiSummarize(AiModelClient model, String input, int maxLength) { + return generateText(model, AiTextFunctionDef.AI_SUMMARIZE, input, maxLength); + } + + public static BinaryVariant aiSentiment(AiModelClient model, String input) { + return generateText(model, 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 aiMask(AiModelClient model, String input, String entities) { + return generateText(model, AiTextFunctionDef.AI_MASK, input, entities); + } + + private static BinaryVariant generateText( + AiModelClient model, + 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"); } - AiTextFunctionDef function = AiTextFunctionDef.AI_COMPLETE; String prompt = - function.buildPrompt(systemPrompt) + function.buildPrompt(promptArguments) + "\n" + buildOutputSchemaHint(function.getOutputType()); String json = ((SupportsTextGeneration) model).generate(prompt, input); @@ -53,11 +88,15 @@ public static BinaryVariant aiComplete(AiModelClient model, String input, String try { return BinaryVariantInternalBuilder.parseJson(json, false); } catch (IOException e) { - throw new RuntimeException("Failed to parse AI response as JSON: " + json, e); + throw new RuntimeException( + "AI function " + function.getFunctionName() + " returned invalid JSON.", e); } } public static List aiEmbed(AiModelClient model, String input) { + if (input == null) { + return null; + } if (!(model instanceof SupportsEmbedding)) { throw new UnsupportedOperationException( "Model " + model.getClass().getName() + " does not support embedding"); @@ -67,7 +106,9 @@ public static List aiEmbed(AiModelClient model, String input) { } private static String buildOutputSchemaHint(RowType outputType) { - StringBuilder builder = new StringBuilder("Return valid JSON with this shape:\n{\n"); + StringBuilder builder = + new StringBuilder( + "Return only valid JSON without Markdown fences or additional text, using this shape:\n{\n"); List fieldNames = outputType.getFieldNames(); for (int i = 0; i < fieldNames.size(); i++) { builder.append(" \"") 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 02691968c27..d7f45c49743 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,11 +18,16 @@ 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.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.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; @@ -901,18 +906,58 @@ public static SqlSelect parseFilterExpression(String filterExpression) { /** 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)) { - validateAiModelReferences(parseProjectionExpression(projection), declaredModelNames); + validateAiModelCapabilities( + parseProjectionExpression(projection), modelClients, userDefinedFunctionNames); } if (!isNullOrWhitespaceOnly(filter)) { - validateAiModelReferences(parseFilterExpression(filter), declaredModelNames); + validateAiModelCapabilities( + parseFilterExpression(filter), modelClients, userDefinedFunctionNames); } } - private static void validateAiModelReferences(SqlNode node, Set declaredModelNames) { + private static void validateAiModelReferences( + SqlNode node, Set declaredModelNames, Set userDefinedFunctionNames) { if (node instanceof SqlCall) { SqlCall call = (SqlCall) node; - if (isAiFunction(call.getOperator().getName())) { + String functionName = call.getOperator().getName(); + if (isAiFunction(functionName) + && !isUserDefinedFunction(functionName, userDefinedFunctionNames)) { if (call.operandCount() == 0) { return; } @@ -931,19 +976,91 @@ private static void validateAiModelReferences(SqlNode node, Set declared } for (SqlNode operand : call.getOperandList()) { if (operand != null) { - validateAiModelReferences(operand, declaredModelNames); + validateAiModelReferences( + operand, declaredModelNames, userDefinedFunctionNames); } } } else if (node instanceof SqlNodeList) { for (SqlNode child : (SqlNodeList) node) { - validateAiModelReferences(child, declaredModelNames); + 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) { + SqlNode modelArgument = call.operand(0); + Preconditions.checkArgument( + modelArgument instanceof SqlCharStringLiteral, + "The model argument of %s must be a string constant, but was %s.", + functionName, + modelArgument); + String modelName = ((SqlCharStringLiteral) modelArgument).getNlsString().getValue(); + AiModelClient modelClient = modelClients.get(modelName); + Preconditions.checkArgument( + modelClient != null, + "Model '%s' referenced by %s has not been declared.", + modelName, + functionName); + if (isTextAiFunction(functionName)) { + Preconditions.checkArgument( + modelClient instanceof SupportsTextGeneration, + "Model '%s' referenced by %s does not support text generation.", + 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 boolean isUserDefinedFunction( + String functionName, Set userDefinedFunctionNames) { + return userDefinedFunctionNames.stream().anyMatch(functionName::equalsIgnoreCase); + } + private static boolean isAiFunction(String functionName) { - return "AI_COMPLETE".equalsIgnoreCase(functionName) - || "AI_EMBED".equalsIgnoreCase(functionName); + return isTextAiFunction(functionName) || isEmbeddingAiFunction(functionName); + } + + 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; } public static boolean hasAsterisk(@Nullable String projection) { 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 3cf9873fc89..89f718d8deb 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 @@ -23,6 +23,9 @@ import org.junit.jupiter.api.Test; +import java.util.ArrayList; +import java.util.List; + import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -34,16 +37,27 @@ private static class TestModelClient private static final long serialVersionUID = 1L; - private String lastPrompt; + private final String response; + private final List prompts = new ArrayList<>(); + private int embedCalls; + + private TestModelClient() { + this("{\"result\":\"ABC\"}"); + } + + private TestModelClient(String response) { + this.response = response; + } @Override public String generate(String systemPrompt, String userInput) { - lastPrompt = systemPrompt; - return "{\"result\":\"ABC\"}"; + prompts.add(systemPrompt); + return response; } @Override public float[] embed(String text) { + embedCalls++; return new float[] {0.1f, 0.2f, 0.3f}; } } @@ -53,13 +67,50 @@ private static class UnsupportedModelClient implements AiModelClient { } @Test - void testAiFunctions() { + void testTextAiFunctionsUseEnglishPromptsAndParseJsonResponses() { TestModelClient model = new TestModelClient(); assertThat(AiFunctions.aiComplete(model, "input", "Return three letters")) .hasToString("{\"result\":\"ABC\"}"); - assertThat(model.lastPrompt).contains("Return three letters").contains("\"result\""); + assertThat(AiFunctions.aiClassify(model, "input", "positive,negative")) + .hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiTranslate(model, "input", "auto", "en")) + .hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiSummarize(model, "input", 100)) + .hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiSentiment(model, "input")).hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiExtract(model, "input", "name:string")) + .hasToString("{\"result\":\"ABC\"}"); + assertThat(AiFunctions.aiMask(model, "input", "email,phone")) + .hasToString("{\"result\":\"ABC\"}"); + + assertThat(model.prompts).hasSize(7); + assertThat(model.prompts.get(0)).contains("Return three letters").contains("\"result\""); + assertThat(model.prompts.get(1)) + .contains("text classifier", "positive,negative", "\"category\""); + assertThat(model.prompts.get(2)) + .contains("translator", "auto", "en", "\"translated_text\""); + assertThat(model.prompts.get(3)) + .contains("text summarizer", "100 characters", "\"summary\""); + assertThat(model.prompts.get(4)) + .contains("sentiment analyzer", "\"score\"", "\"confidence\""); + assertThat(model.prompts.get(5)) + .contains("information extraction", "name:string", "\"extracted_json\""); + assertThat(model.prompts.get(6)).contains("data masking", "email,phone", "\"masked_text\""); + assertThat(model.prompts) + .allSatisfy( + prompt -> + assertThat(prompt) + .contains("Return only valid JSON") + .doesNotContainPattern("\\p{IsHan}")); + } + + @Test + void testEmbeddingFunction() { + TestModelClient model = new TestModelClient(); + assertThat(AiFunctions.aiEmbed(model, "input")).containsExactly(0.1f, 0.2f, 0.3f); + assertThat(model.embedCalls).isOne(); } @Test @@ -76,16 +127,28 @@ void testUnsupportedCapabilities() { @Test void testInvalidJsonResponse() { - TestModelClient model = - new TestModelClient() { - @Override - public String generate(String systemPrompt, String userInput) { - return "not-json"; - } - }; + TestModelClient model = new TestModelClient("not-json"); - assertThatThrownBy(() -> AiFunctions.aiComplete(model, "input", "prompt")) + assertThatThrownBy(() -> AiFunctions.aiClassify(model, "input", "positive,negative")) .isInstanceOf(RuntimeException.class) - .hasMessageContaining("Failed to parse AI response as JSON"); + .hasMessage("AI function AI_CLASSIFY returned invalid JSON."); + } + + @Test + void testNullInputSkipsModelInvocation() { + TestModelClient model = new TestModelClient(); + + assertThat(AiFunctions.aiClassify(model, null, "positive,negative")).isNull(); + assertThat(AiFunctions.aiEmbed(model, null)).isNull(); + assertThat(model.prompts).isEmpty(); + assertThat(model.embedCalls).isZero(); + } + + @Test + void testNullModelResponseReturnsNull() { + TestModelClient model = new TestModelClient(null); + + assertThat(AiFunctions.aiSummarize(model, "input", 100)).isNull(); + assertThat(model.prompts).hasSize(1); } } 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 e7af51542e0..76cb744a5d1 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,15 +17,20 @@ 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.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; import org.apache.flink.cdc.runtime.operators.transform.ProjectionColumn; +import org.apache.flink.cdc.runtime.operators.transform.UserDefinedFunctionDescriptor; import org.junit.jupiter.api.Test; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Set; import static org.assertj.core.api.Assertions.assertThat; @@ -40,6 +45,24 @@ class AiFunctionParserTest { Column.physicalColumn("id", DataTypes.INT()), Column.physicalColumn("content", DataTypes.STRING())); + 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]; + } + } + @Test void testTranslateAiFunctions() { List columns = @@ -56,6 +79,69 @@ void testTranslateAiFunctions() { .containsExactly(DataTypes.VARIANT(), DataTypes.ARRAY(DataTypes.FLOAT())); } + @Test + void testTranslateSpecializedTextAiFunctions() { + List columns = + translate( + "AI_CLASSIFY('model', content, 'positive,negative') AS classified, " + + "AI_TRANSLATE('model', content, 'auto', 'en') AS translated, " + + "AI_SUMMARIZE('model', content, 100) AS summarized, " + + "AI_SENTIMENT('model', content) AS sentiment, " + + "AI_EXTRACT('model', content, 'name:string') AS extracted, " + + "AI_MASK('model', content, 'email,phone') AS masked"); + + 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\")"); + assertThat(columns) + .extracting(ProjectionColumn::getDataType) + .containsOnly(DataTypes.VARIANT()); + } + + @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(); + + List columns = + TransformParser.generateProjectionColumns( + "AI_SENTIMENT(id) AS sentiment", + COLUMNS, + List.of( + new UserDefinedFunctionDescriptor( + "ai_sentiment", + "org.apache.flink.cdc.udf.examples.java.AddOneFunctionClass")), + new SupportedMetadataColumn[0]); + + assertThat(columns) + .extracting(ProjectionColumn::getScriptExpression) + .containsExactly("__udf_ai_sentiment.eval($0)"); + assertThat(columns) + .extracting(ProjectionColumn::getDataType) + .containsExactly(DataTypes.STRING()); + } + @Test void testModelArgumentMustBeStringConstant() { assertThatThrownBy( @@ -87,6 +173,53 @@ void testReferencedModelMustBeDeclared() { 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()); + + assertThatCode( + () -> + TransformParser.validateAiModelCapabilities( + "AI_CLASSIFY('textModel', content, 'a,b') AS classified, " + + "AI_EMBED('embeddingModel', content) AS 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"); } @Test @@ -97,6 +230,18 @@ void testFunctionArityValidation() { .hasMessageContaining("Invalid number of arguments to function 'AI_COMPLETE'"); assertThatThrownBy(() -> translate("AI_COMPLETE() AS completed")) .hasMessageContaining("Invalid number of arguments to function 'AI_COMPLETE'"); + assertThatThrownBy(() -> translate("AI_CLASSIFY('model', content) AS classified")) + .hasMessageContaining("Invalid number of arguments to function 'AI_CLASSIFY'"); + assertThatThrownBy(() -> translate("AI_TRANSLATE('model', content, 'auto') AS translated")) + .hasMessageContaining("Invalid number of arguments to function 'AI_TRANSLATE'"); + assertThatThrownBy(() -> translate("AI_SENTIMENT('model') AS sentiment")) + .hasMessageContaining("Invalid number of arguments to function 'AI_SENTIMENT'"); + assertThatThrownBy(() -> translate("AI_EXTRACT('model', content) AS extracted")) + .hasMessageContaining("Invalid number of arguments to function 'AI_EXTRACT'"); + assertThatThrownBy(() -> translate("AI_MASK('model', content) AS masked")) + .hasMessageContaining("Invalid number of arguments to function 'AI_MASK'"); + assertThatThrownBy(() -> translate("AI_SUMMARIZE('model', content, TRUE) AS summarized")) + .hasMessageContaining("Cannot apply 'AI_SUMMARIZE'"); assertThatCode( () -> TransformParser.validateAiModelReferences(