diff --git a/docs/content.zh/docs/core-concept/ai-model.md b/docs/content.zh/docs/core-concept/ai-model.md index 0a36d17fa56..d809b3224a9 100644 --- a/docs/content.zh/docs/core-concept/ai-model.md +++ b/docs/content.zh/docs/core-concept/ai-model.md @@ -24,11 +24,11 @@ under the License. # AI 模型 -AI 模型可用于 transform 表达式中的文本生成、文本分析和 embedding。 +AI 模型可用于 transform 表达式中的文本生成、文本分析、embedding 和图片理解。 ## AI Functions -模型名称必须是字符串常量,并引用 `pipeline.model` 中声明的模型。文本函数要求模型客户端支持文本生成,`AI_EMBED` 要求模型客户端支持 embedding;Pipeline 会在执行前校验引用模型的 capability 是否匹配。 +模型名称必须是字符串常量,并引用 `pipeline.model` 中声明的模型。文本、embedding 和图片函数分别要求模型客户端实现对应的 capability;Pipeline 会在执行前校验引用模型的 capability 是否匹配。 所有文本函数都会将模型返回的 JSON 解析为 `VARIANT`。 @@ -43,13 +43,24 @@ AI 模型可用于 transform 表达式中的文本生成、文本分析和 embed | `AI_MASK(model, input, entities)` | 对指定实体类型进行脱敏。 | `masked_text`、`detected_entities` | | `AI_EMBED(model, input)` | 生成 embedding 向量。 | 不返回 JSON,而是返回 `ARRAY`。 | +以下多模态函数从 `BYTES` 字段读取图片数据: + +| 函数 | 说明 | 返回类型 | +|------|------|----------| +| `AI_IMAGE_COMPLETE(model, image, prompt)` | 根据图片和自然语言 prompt 生成文本。 | `STRING` | +| `AI_IMAGE_EMBED(model, image)` | 将图片转换为 embedding 向量。 | `ARRAY` | + +OpenAI-compatible 模型客户端通过标准 vision chat 支持 `AI_IMAGE_COMPLETE`。客户端会识别 PNG、JPEG、GIF 和 WebP 图片,并将图片编码为 Base64 data URL。图片为 `NULL` 时直接返回 `NULL`,且不会调用模型;图片为空或格式无法识别时,会在发送请求前报错。 + +`AI_IMAGE_EMBED` 当前只提供框架函数和 provider capability。OpenAI API 目前没有定义标准的图片向量化协议,因此图片 embedding 需要由具体 provider 单独实现。OpenAI-compatible 模型客户端不实现图片 embedding,社区发行包目前也没有可用于生产的图片 embedding provider。需要图片向量化的用户需要等待后续 provider 实现。 + 六个专用文本函数使用内置英文 prompt 模板,但输入文本可以是任意语言。输入为 `NULL` 时直接返回 `NULL`,且不会调用模型;模型返回 `NULL` 时也返回 `NULL`。非空文本响应必须是语法合法的 JSON,否则当前记录处理失败,错误信息会标明具体 AI 函数。运行时只校验 JSON 语法,不校验响应字段是否存在或字段类型是否匹配。 ## OpenAI-compatible 模型客户端 AI 模型客户端可供上述 AI Functions 引用。使用时,需要通过 `--jar` 将模型实现 JAR(例如 `flink-cdc-pipeline-model-openai-compatible`)添加到 Pipeline 命令中。 -OpenAI-compatible 客户端支持调用实现 OpenAI Chat Completions 和 Embeddings REST API 的服务。 +OpenAI-compatible 客户端支持调用实现 OpenAI Chat Completions、vision chat 和 Embeddings REST API 的服务。 system prompt、函数 prompt 和输入文本均支持英文或中文内容。 @@ -60,6 +71,7 @@ transform: *, AI_COMPLETE('completion_model', content, '总结输入内容') AS summary, AI_SENTIMENT('completion_model', content) AS sentiment, + AI_IMAGE_COMPLETE('vision_model', image, '描述这张图片') AS image_description, AI_EMBED('embedding_model', content) AS embedding pipeline: @@ -80,6 +92,12 @@ pipeline: endpoint: https://api.example.com/v1 api-key: dimension: 768 + - name: vision_model + type: openai-compatible + options: + model: gpt-4o-mini + endpoint: https://api.example.com/v1 + api-key: ``` 不要将 API Key 提交到代码仓库中,请通过部署环境的密钥管理机制提供。 diff --git a/docs/content/docs/core-concept/ai-model.md b/docs/content/docs/core-concept/ai-model.md index 8aa7666838a..1ad5a739fe0 100644 --- a/docs/content/docs/core-concept/ai-model.md +++ b/docs/content/docs/core-concept/ai-model.md @@ -24,11 +24,12 @@ under the License. # AI Model -AI models can be used in transform expressions for text generation, text analysis, and embedding. +AI models can be used in transform expressions for text generation, text analysis, embedding, and +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 `AI_EMBED` requires embedding support. The pipeline validates the referenced model capability before execution. +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. All text functions return `VARIANT` values parsed from the model's JSON response. @@ -43,13 +44,31 @@ All text functions return `VARIANT` values parsed from the model's JSON response | `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 following multimodal functions accept image data from a `BYTES` column: + +| Function | Description | Return type | +|----------|-------------|-------------| +| `AI_IMAGE_COMPLETE(model, image, prompt)` | Generates text from an image and a natural-language prompt. | `STRING` | +| `AI_IMAGE_EMBED(model, image)` | Converts an image into an embedding vector. | `ARRAY` | + +The OpenAI-compatible model client supports `AI_IMAGE_COMPLETE` through standard vision chat. It +detects PNG, JPEG, GIF, and WebP images and sends the image as a Base64 data URL. A `NULL` image +returns `NULL` without invoking the model, while empty or unrecognized image data is rejected before +the request is sent. + +`AI_IMAGE_EMBED` currently provides only the framework function and provider capability. The OpenAI +API does not define a standard image embedding protocol, so image embedding requires a +provider-specific implementation. The OpenAI-compatible model client does not implement image +embedding, and the community distribution does not yet include a production provider for it. Users +who need image embedding must wait for a follow-up provider implementation. + 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 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. +The OpenAI-compatible client supports chat completions, vision chat, and text embeddings against endpoints that implement the corresponding OpenAI REST APIs. System prompts, function prompts, and input text may contain either English or Chinese content. @@ -60,6 +79,7 @@ transform: *, AI_COMPLETE('completion_model', content, 'Summarize the input') AS summary, AI_SENTIMENT('completion_model', content) AS sentiment, + AI_IMAGE_COMPLETE('vision_model', image, 'Describe the image') AS image_description, AI_EMBED('embedding_model', content) AS embedding pipeline: @@ -80,6 +100,12 @@ pipeline: endpoint: https://api.example.com/v1 api-key: dimension: 768 + - name: vision_model + type: openai-compatible + options: + model: gpt-4o-mini + endpoint: https://api.example.com/v1 + api-key: ``` Do not store API keys in source control. Supply them through the secret-management mechanism of your deployment environment. diff --git a/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageEmbedding.java b/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageEmbedding.java new file mode 100644 index 00000000000..aa5f1850ff3 --- /dev/null +++ b/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageEmbedding.java @@ -0,0 +1,29 @@ +/* + * 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.common.model.abilities; + +import org.apache.flink.cdc.common.annotation.Experimental; +import org.apache.flink.cdc.common.model.AiModelClient; + +/** Ability interface for {@link AiModelClient} implementations that can embed image bytes. */ +@Experimental +public interface SupportsImageEmbedding { + + /** Converts the given image bytes into a dense float vector. */ + float[] embedImage(byte[] image); +} diff --git a/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageTextGeneration.java b/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageTextGeneration.java new file mode 100644 index 00000000000..1f897411c36 --- /dev/null +++ b/flink-cdc-common/src/main/java/org/apache/flink/cdc/common/model/abilities/SupportsImageTextGeneration.java @@ -0,0 +1,32 @@ +/* + * 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.common.model.abilities; + +import org.apache.flink.cdc.common.annotation.Experimental; +import org.apache.flink.cdc.common.model.AiModelClient; + +/** + * Ability interface for {@link AiModelClient} implementations that can generate text from an image + * combined with a natural-language prompt. + */ +@Experimental +public interface SupportsImageTextGeneration { + + /** Generates text from the given image bytes guided by a natural-language prompt. */ + String generateTextFromImage(byte[] image, String prompt); +} 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 d50f1c67d12..d0610ed9f58 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 @@ -176,6 +176,32 @@ void testAiEmbedInProjection() throws Exception { "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, [3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0, 6.0]], op=INSERT, meta=()}"); } + @Test + void testAiImageCompleteInProjection() throws Exception { + String[] output = + runAiFunctionTest( + "id, AI_IMAGE_COMPLETE('visionModel', image, 'Describe the image') AS description", + List.of(ModelDef.of("visionModel", "dummy", Collections.emptyMap()))); + + assertThat(output) + .containsExactly( + "CreateTableEvent{tableId=default_namespace.default_schema.mytable1, schema=columns={`id` INT NOT NULL,`description` STRING}, primaryKeys=id, options=()}", + "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, A dummy description of the image.], op=INSERT, meta=()}"); + } + + @Test + void testAiImageEmbedInProjection() throws Exception { + String[] output = + runAiFunctionTest( + "id, AI_IMAGE_EMBED('imageEmbedModel', image) AS embedding", + List.of(ModelDef.of("imageEmbedModel", "dummy", Collections.emptyMap()))); + + assertThat(output) + .containsExactly( + "CreateTableEvent{tableId=default_namespace.default_schema.mytable1, schema=columns={`id` INT NOT NULL,`embedding` ARRAY}, primaryKeys=id, options=()}", + "DataChangeEvent{tableId=default_namespace.default_schema.mytable1, before=[], after=[1, [2.0, 7.0, 1.0, 8.0, 2.0, 8.0]], op=INSERT, meta=()}"); + } + private String[] runAiFunctionTest(String projection, List models) throws Exception { return runAiFunctionTest(projection, models, Collections.emptyList()); } @@ -204,6 +230,7 @@ private String[] runAiFunctionTest( Schema.newBuilder() .physicalColumn("id", DataTypes.INT()) .physicalColumn("content", DataTypes.STRING()) + .physicalColumn("image", DataTypes.BYTES()) .primaryKey("id") .build(); BinaryRecordDataGenerator generator = @@ -216,7 +243,9 @@ private String[] runAiFunctionTest( tableId, generator.generate( new Object[] { - 1, BinaryStringData.fromString("I love this product") + 1, + BinaryStringData.fromString("I love this product"), + new byte[] {1, 2, 3, 4} }))); ValuesDataSourceHelper.setSourceEvents(Collections.singletonList(events)); 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 d92c0a8afb7..5a4045e71ad 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 @@ -19,10 +19,17 @@ 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; /** Deterministic AI model client used by tests. */ -public class DummyModelClient implements AiModelClient, SupportsTextGeneration, SupportsEmbedding { +public class DummyModelClient + implements AiModelClient, + SupportsTextGeneration, + SupportsEmbedding, + SupportsImageTextGeneration, + SupportsImageEmbedding { private static final long serialVersionUID = 1L; @@ -63,6 +70,24 @@ public float[] embed(String text) { return new float[] {3f, 1f, 4f, 1f, 5f, 9f, 2f, 6f}; } + @Override + public String generateTextFromImage(byte[] image, String prompt) { + if (debug) { + System.out.printf( + "Received image of %d bytes%nPrompt: %s%n", + image == null ? 0 : image.length, prompt); + } + return "A dummy description of the image."; + } + + @Override + public float[] embedImage(byte[] image) { + if (debug) { + System.out.printf("Received image of %d bytes%n", image == null ? 0 : image.length); + } + return new float[] {2f, 7f, 1f, 8f, 2f, 8f}; + } + @Override public void open() { if (debug) { diff --git a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/main/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClient.java b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/main/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClient.java index 5a1e9041c44..3a9762f40ce 100644 --- a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/main/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClient.java +++ b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/main/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClient.java @@ -17,8 +17,10 @@ package org.apache.flink.cdc.models.openai; +import org.apache.flink.cdc.common.annotation.VisibleForTesting; 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.SupportsImageTextGeneration; import org.apache.flink.cdc.common.model.abilities.SupportsTextGeneration; import com.fasterxml.jackson.core.JsonProcessingException; @@ -33,6 +35,7 @@ import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionContentPart; import com.openai.models.chat.completions.ChatCompletionContentPartImage; +import com.openai.models.chat.completions.ChatCompletionContentPartText; import com.openai.models.chat.completions.ChatCompletionCreateParams; import com.openai.models.embeddings.CreateEmbeddingResponse; import com.openai.models.embeddings.Embedding; @@ -44,6 +47,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.Iterator; @@ -53,10 +57,22 @@ /** AI model client that connects to an OpenAI-compatible endpoint. */ public class OpenAiCompatibleModelClient - implements AiModelClient, SupportsTextGeneration, SupportsEmbedding { + implements AiModelClient, + SupportsTextGeneration, + SupportsEmbedding, + SupportsImageTextGeneration { private static final Logger LOG = LoggerFactory.getLogger(OpenAiCompatibleModelClient.class); + private static final byte[] PNG_SIGNATURE = + new byte[] {(byte) 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}; + private static final byte[] JPEG_SIGNATURE = new byte[] {(byte) 0xFF, (byte) 0xD8, (byte) 0xFF}; + private static final byte[] GIF87A_SIGNATURE = new byte[] {0x47, 0x49, 0x46, 0x38, 0x37, 0x61}; + private static final byte[] GIF89A_SIGNATURE = new byte[] {0x47, 0x49, 0x46, 0x38, 0x39, 0x61}; + private static final byte[] RIFF_SIGNATURE = new byte[] {0x52, 0x49, 0x46, 0x46}; + private static final byte[] WEBP_SIGNATURE = new byte[] {0x57, 0x45, 0x42, 0x50}; + private static final int WEBP_SIGNATURE_OFFSET = 8; + private static final long serialVersionUID = 1L; private final String endpoint; @@ -109,6 +125,15 @@ public float[] embed(String text) { return executeWithRetry("embedding", () -> createEmbedding(text)); } + @Override + public String generateTextFromImage(byte[] image, String prompt) { + if (image == null) { + return null; + } + String imageDataUrl = buildImageDataUrl(image); + return executeWithRetry("image completion", () -> completeImage(imageDataUrl, prompt)); + } + private String complete(String systemPrompt, String userInput) { ChatCompletionCreateParams.Builder builder = ChatCompletionCreateParams.builder().model(model); @@ -150,6 +175,96 @@ private String complete(String systemPrompt, String userInput) { "OpenAI-compatible text completion returned no text content.")); } + private String completeImage(String imageDataUrl, String prompt) { + ChatCompletionContentPart imagePart = + ChatCompletionContentPart.ofImageUrl( + ChatCompletionContentPartImage.builder() + .imageUrl( + ChatCompletionContentPartImage.ImageUrl.builder() + .url(imageDataUrl) + .build()) + .build()); + ChatCompletionContentPart textPart = + ChatCompletionContentPart.ofText( + ChatCompletionContentPartText.builder() + .text(prompt != null ? prompt : "") + .build()); + + ChatCompletionCreateParams.Builder builder = + ChatCompletionCreateParams.builder().model(model); + if (configuredSystemPrompt != null) { + builder.addSystemMessage(configuredSystemPrompt); + } + builder.addUserMessageOfArrayOfContentParts(List.of(textPart, imagePart)); + if (params.userPrompt != null) { + builder.addUserMessage(params.userPrompt); + } + applyCompletionParams(builder); + builder.putAllAdditionalHeaders(headersOrEmpty()); + builder.putAllAdditionalBodyProperties(bodyOrEmpty()); + + ChatCompletion completion = currentClient().chat().completions().create(builder.build()); + if (completion.choices().isEmpty()) { + throw new IllegalStateException( + "OpenAI-compatible image completion returned no choices."); + } + return completion + .choices() + .get(0) + .message() + .content() + .orElseThrow( + () -> + new IllegalStateException( + "OpenAI-compatible image completion returned no text content.")); + } + + private static String buildImageDataUrl(byte[] image) { + return "data:" + + detectImageMimeType(image) + + ";base64," + + Base64.getEncoder().encodeToString(image); + } + + /** + * Detects the MIME type of an image from its leading magic bytes. Throws {@link + * IllegalArgumentException} when the format cannot be recognized or the input is empty. + */ + @VisibleForTesting + static String detectImageMimeType(byte[] bytes) { + if (bytes == null || bytes.length == 0) { + throw new IllegalArgumentException("Image bytes must not be null or empty."); + } + if (matchesSignature(bytes, 0, PNG_SIGNATURE)) { + return "image/png"; + } + if (matchesSignature(bytes, 0, JPEG_SIGNATURE)) { + return "image/jpeg"; + } + if (matchesSignature(bytes, 0, GIF87A_SIGNATURE) + || matchesSignature(bytes, 0, GIF89A_SIGNATURE)) { + return "image/gif"; + } + if (matchesSignature(bytes, 0, RIFF_SIGNATURE) + && matchesSignature(bytes, WEBP_SIGNATURE_OFFSET, WEBP_SIGNATURE)) { + return "image/webp"; + } + throw new IllegalArgumentException( + "Unrecognized image format. Supported formats: PNG, JPEG, GIF, WebP."); + } + + private static boolean matchesSignature(byte[] bytes, int offset, byte[] signature) { + if (bytes.length < offset + signature.length) { + return false; + } + for (int i = 0; i < signature.length; i++) { + if (bytes[offset + i] != signature[i]) { + return false; + } + } + return true; + } + private void applyCompletionParams(ChatCompletionCreateParams.Builder builder) { if (params.temperature != null) { builder.temperature(params.temperature); diff --git a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientFactoryTest.java b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientFactoryTest.java index 9265683d99b..61d4eb72fa5 100644 --- a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientFactoryTest.java +++ b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientFactoryTest.java @@ -21,6 +21,8 @@ import org.apache.flink.cdc.common.model.AiModelClient; import org.apache.flink.cdc.common.model.ModelContext; 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.table.api.ValidationException; @@ -73,7 +75,9 @@ void testCreateClient() { assertThat(client) .isInstanceOf(OpenAiCompatibleModelClient.class) .isInstanceOf(SupportsTextGeneration.class) - .isInstanceOf(SupportsEmbedding.class); + .isInstanceOf(SupportsEmbedding.class) + .isInstanceOf(SupportsImageTextGeneration.class) + .isNotInstanceOf(SupportsImageEmbedding.class); } @Test diff --git a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientTest.java b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientTest.java index e8352a94e82..e291415f13f 100644 --- a/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientTest.java +++ b/flink-cdc-pipeline-model/flink-cdc-pipeline-model-openai-compatible/src/test/java/org/apache/flink/cdc/models/openai/OpenAiCompatibleModelClientTest.java @@ -161,6 +161,95 @@ void testImageUrlCompletion() throws Exception { .isEqualTo("https://example.com/image.png"); } + @Test + void testImageCompletionUsesStandardVisionChatRequest() throws Exception { + server.enqueue(jsonResponse(200, COMPLETION_RESPONSE)); + Map options = baseOptions(); + options.put("system-prompt", "You are a vision assistant."); + options.put("user-prompt", "Answer briefly."); + options.put("temperature", "0.2"); + options.put("extra-header", "{\"X-Vision\":\"enabled\"}"); + options.put("extra-body", "{\"vendor_flag\":true}"); + client = createAndOpenClient(options); + + byte[] png = + new byte[] { + (byte) 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x01, 0x02, 0x03 + }; + assertThat(client.generateTextFromImage(png, "What is in this image?")) + .isEqualTo("{\"result\":\"done\"}"); + + RecordedRequest request = server.takeRequest(); + assertThat(request.getPath()).isEqualTo("/v1/chat/completions"); + assertThat(request.getHeader("X-Vision")).isEqualTo("enabled"); + JsonNode body = objectMapper.readTree(request.getBody().readUtf8()); + assertThat(body.at("/messages/0/role").asText()).isEqualTo("system"); + assertThat(body.at("/messages/0/content").asText()) + .isEqualTo("You are a vision assistant."); + assertThat(body.at("/messages/1/content/0/type").asText()).isEqualTo("text"); + assertThat(body.at("/messages/1/content/0/text").asText()) + .isEqualTo("What is in this image?"); + assertThat(body.at("/messages/1/content/1/type").asText()).isEqualTo("image_url"); + assertThat(body.at("/messages/1/content/1/image_url/url").asText()) + .isEqualTo("data:image/png;base64,iVBORw0KGgoBAgM="); + assertThat(body.at("/messages/2/content").asText()).isEqualTo("Answer briefly."); + assertThat(body.path("temperature").asDouble()).isEqualTo(0.2d); + assertThat(body.path("vendor_flag").asBoolean()).isTrue(); + } + + @Test + void testImageMimeTypeDetection() { + assertThat( + OpenAiCompatibleModelClient.detectImageMimeType( + new byte[] {(byte) 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A})) + .isEqualTo("image/png"); + assertThat( + OpenAiCompatibleModelClient.detectImageMimeType( + new byte[] {(byte) 0xFF, (byte) 0xD8, (byte) 0xFF})) + .isEqualTo("image/jpeg"); + assertThat( + OpenAiCompatibleModelClient.detectImageMimeType( + new byte[] {0x47, 0x49, 0x46, 0x38, 0x39, 0x61})) + .isEqualTo("image/gif"); + assertThat( + OpenAiCompatibleModelClient.detectImageMimeType( + new byte[] { + 0x52, 0x49, 0x46, 0x46, 0, 0, 0, 0, 0x57, 0x45, 0x42, 0x50 + })) + .isEqualTo("image/webp"); + } + + @Test + void testNullAndInvalidImageInputsDoNotSendRequests() { + client = createAndOpenClient(baseOptions()); + + assertThat(client.generateTextFromImage(null, "describe")).isNull(); + assertThatThrownBy(() -> client.generateTextFromImage(new byte[0], "describe")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must not be null or empty"); + assertThatThrownBy( + () -> + client.generateTextFromImage( + new byte[] {0x01, 0x02, 0x03}, "describe")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Unrecognized image format"); + assertThat(server.getRequestCount()).isZero(); + } + + @Test + void testRetryableImageCompletionErrorIsRetried() { + server.enqueue(jsonResponse(429, ERROR_RESPONSE)); + server.enqueue(jsonResponse(200, COMPLETION_RESPONSE)); + Map options = baseOptions(); + options.put("retry-num", "2"); + options.put("retry-backoff-base-interval", "1 ms"); + client = createAndOpenClient(options); + + byte[] jpeg = new byte[] {(byte) 0xFF, (byte) 0xD8, (byte) 0xFF}; + assertThat(client.generateTextFromImage(jpeg, "describe")).contains("done"); + assertThat(server.getRequestCount()).isEqualTo(2); + } + @Test void testTextEmbedding() throws Exception { server.enqueue(jsonResponse(200, EMBEDDING_RESPONSE)); diff --git a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiImageFunctionDef.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiImageFunctionDef.java new file mode 100644 index 00000000000..86982a9a9d1 --- /dev/null +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/ai/AiImageFunctionDef.java @@ -0,0 +1,75 @@ +/* + * 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.types.DataType; +import org.apache.flink.cdc.common.types.DataTypes; +import org.apache.flink.cdc.common.types.RowType; + +/** Built-in AI image function definitions. */ +public enum AiImageFunctionDef { + AI_IMAGE_COMPLETE( + "AI_IMAGE_COMPLETE", + RowType.of( + new DataType[] {DataTypes.BYTES(), DataTypes.STRING()}, + new String[] {"image", "prompt"}), + DataTypes.STRING(), + Capability.IMAGE_TEXT_GENERATION), + + AI_IMAGE_EMBED( + "AI_IMAGE_EMBED", + RowType.of(new DataType[] {DataTypes.BYTES()}, new String[] {"image"}), + DataTypes.ARRAY(DataTypes.FLOAT()), + Capability.IMAGE_EMBEDDING); + + /** Capability required by an image AI function. */ + public enum Capability { + IMAGE_TEXT_GENERATION, + IMAGE_EMBEDDING + } + + private final String functionName; + private final RowType inputType; + private final DataType outputType; + private final Capability capability; + + AiImageFunctionDef( + String functionName, RowType inputType, DataType outputType, Capability capability) { + this.functionName = functionName; + this.inputType = inputType; + this.outputType = outputType; + this.capability = capability; + } + + public String getFunctionName() { + return functionName; + } + + /** Returns the parameter types after the model argument. */ + public RowType getInputType() { + return inputType; + } + + public DataType getOutputType() { + return outputType; + } + + public Capability getCapability() { + return capability; + } +} 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 81aeadc260e..2c21ba1c71f 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 @@ -19,6 +19,8 @@ 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.types.RowType; import org.apache.flink.cdc.common.types.variant.BinaryVariant; @@ -111,6 +113,33 @@ public static List aiEmbed(AiModelClient model, String input) { return embedding == null ? null : Floats.asList(embedding); } + /** Dispatches image-to-text AI functions. */ + public static String aiImageComplete(AiModelClient model, byte[] image, String prompt) { + 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); + } + + /** Dispatches image embedding AI functions. */ + public static List aiImageEmbed(AiModelClient model, byte[] image) { + if (image == null) { + return null; + } + if (!(model instanceof SupportsImageEmbedding)) { + throw new UnsupportedOperationException( + "Model " + model.getClass().getName() + " does not support image embedding"); + } + float[] embedding = ((SupportsImageEmbedding) model).embedImage(image); + return embedding == null ? null : Floats.asList(embedding); + } + private static String truncateInvalidJsonResponse(String response) { if (response.length() <= MAX_INVALID_JSON_RESPONSE_LENGTH) { return response; 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 1497d6cda23..627cde88fa2 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 @@ -30,6 +30,7 @@ import org.apache.flink.cdc.common.utils.Preconditions; import org.apache.flink.cdc.common.utils.StringUtils; 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.UserDefinedFunctionDescriptor; import org.apache.flink.cdc.runtime.parser.metadata.MetadataColumns; @@ -1113,6 +1114,11 @@ private static boolean isAiFunction(String upperCaseName) { return true; } } + for (AiImageFunctionDef def : AiImageFunctionDef.values()) { + if (def.getFunctionName().equals(upperCaseName)) { + return true; + } + } return false; } 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 1ee15901bda..5993691af92 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 @@ -20,6 +20,8 @@ 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; @@ -27,6 +29,7 @@ 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; @@ -998,12 +1001,16 @@ private static void validateAiModelCapabilities( "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, @@ -1024,6 +1031,32 @@ private static void validateAiModelCapabilities( } } + 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( @@ -1040,7 +1073,9 @@ private static boolean isUserDefinedFunction( } private static boolean isAiFunction(String functionName) { - return isTextAiFunction(functionName) || isEmbeddingAiFunction(functionName); + return isTextAiFunction(functionName) + || isEmbeddingAiFunction(functionName) + || findImageAiFunction(functionName) != null; } private static boolean isTextAiFunction(String functionName) { @@ -1061,6 +1096,16 @@ private static boolean isEmbeddingAiFunction(String functionName) { 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/main/java/org/apache/flink/cdc/runtime/parser/metadata/AiFunctionSqlOperatorTable.java b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/metadata/AiFunctionSqlOperatorTable.java index c1d1fa99573..9088d699cf0 100644 --- a/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/metadata/AiFunctionSqlOperatorTable.java +++ b/flink-cdc-runtime/src/main/java/org/apache/flink/cdc/runtime/parser/metadata/AiFunctionSqlOperatorTable.java @@ -20,6 +20,7 @@ import org.apache.flink.cdc.common.types.DataType; import org.apache.flink.cdc.common.types.RowType; 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.typeutils.CalciteDataTypeConverter; @@ -36,7 +37,7 @@ import java.util.ArrayList; import java.util.List; -/** Creates SqlOperatorTable from {@link AiTextFunctionDef} definitions. */ +/** Creates an SqlOperatorTable containing the built-in AI functions. */ public class AiFunctionSqlOperatorTable { private AiFunctionSqlOperatorTable() {} @@ -50,6 +51,9 @@ public static SqlOperatorTable create() { for (AiEmbeddingFunctionDef def : AiEmbeddingFunctionDef.values()) { functions.add(createEmbeddingSqlFunction(def)); } + for (AiImageFunctionDef def : AiImageFunctionDef.values()) { + functions.add(createImageSqlFunction(def)); + } return SqlOperatorTables.of(functions); } @@ -75,6 +79,18 @@ private static SqlFunction createEmbeddingSqlFunction(AiEmbeddingFunctionDef def SqlFunctionCategory.USER_DEFINED_FUNCTION); } + private static SqlFunction createImageSqlFunction(AiImageFunctionDef def) { + return new SqlFunction( + def.getFunctionName(), + SqlKind.OTHER_FUNCTION, + opBinding -> + CalciteDataTypeConverter.convertCalciteType( + opBinding.getTypeFactory(), def.getOutputType()), + null, + OperandTypes.family(toSqlTypeFamiliesWithModelArgument(def.getInputType())), + SqlFunctionCategory.USER_DEFINED_FUNCTION); + } + /** * Converts inputType to SqlTypeFamily array, prepending additional parameters: modelName * (STRING) and input (STRING). @@ -89,6 +105,15 @@ private static SqlTypeFamily[] toSqlTypeFamiliesWithAdditionalParams(RowType inp return families.toArray(new SqlTypeFamily[0]); } + private static SqlTypeFamily[] toSqlTypeFamiliesWithModelArgument(RowType inputType) { + List families = new ArrayList<>(); + families.add(SqlTypeFamily.STRING); + for (DataType fieldType : inputType.getFieldTypes()) { + families.add(toSqlTypeFamily(fieldType)); + } + return families.toArray(new SqlTypeFamily[0]); + } + private static SqlTypeFamily toSqlTypeFamily(DataType dataType) { switch (dataType.getTypeRoot()) { case VARCHAR: @@ -103,6 +128,9 @@ private static SqlTypeFamily toSqlTypeFamily(DataType dataType) { return SqlTypeFamily.APPROXIMATE_NUMERIC; case BOOLEAN: return SqlTypeFamily.BOOLEAN; + case BINARY: + case VARBINARY: + return SqlTypeFamily.BINARY; default: return SqlTypeFamily.ANY; } 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 adf324bf6bf..8eee49ae56f 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 @@ -19,6 +19,8 @@ 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.junit.jupiter.api.Test; @@ -33,13 +35,19 @@ class AiFunctionsTest { private static class TestModelClient - implements AiModelClient, SupportsTextGeneration, SupportsEmbedding { + implements AiModelClient, + SupportsTextGeneration, + SupportsEmbedding, + SupportsImageTextGeneration, + SupportsImageEmbedding { private static final long serialVersionUID = 1L; private final String response; private final List prompts = new ArrayList<>(); private int embedCalls; + private int imageTextCalls; + private int imageEmbedCalls; private TestModelClient() { this("{\"result\":\"ABC\"}"); @@ -60,6 +68,18 @@ public float[] embed(String text) { embedCalls++; return new float[] {0.1f, 0.2f, 0.3f}; } + + @Override + public String generateTextFromImage(byte[] image, String prompt) { + imageTextCalls++; + return "image has " + image.length + " bytes, prompt: " + prompt; + } + + @Override + public float[] embedImage(byte[] image) { + imageEmbedCalls++; + return new float[] {0.9f, 0.8f, 0.7f}; + } } private static class UnsupportedModelClient implements AiModelClient { @@ -113,6 +133,18 @@ void testEmbeddingFunction() { assertThat(model.embedCalls).isOne(); } + @Test + void testImageAiFunctions() { + TestModelClient model = new TestModelClient(); + byte[] image = new byte[] {1, 2, 3, 4}; + + assertThat(AiFunctions.aiImageComplete(model, image, "Describe the image")) + .isEqualTo("image has 4 bytes, prompt: Describe the image"); + assertThat(AiFunctions.aiImageEmbed(model, image)).containsExactly(0.9f, 0.8f, 0.7f); + assertThat(model.imageTextCalls).isOne(); + assertThat(model.imageEmbedCalls).isOne(); + } + @Test void testUnsupportedCapabilities() { UnsupportedModelClient model = new UnsupportedModelClient(); @@ -123,6 +155,12 @@ void testUnsupportedCapabilities() { assertThatThrownBy(() -> AiFunctions.aiEmbed(model, "input")) .isInstanceOf(UnsupportedOperationException.class) .hasMessageContaining("does not support embedding"); + assertThatThrownBy(() -> AiFunctions.aiImageComplete(model, new byte[] {1, 2}, "describe")) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("does not support image text generation"); + assertThatThrownBy(() -> AiFunctions.aiImageEmbed(model, new byte[] {1, 2})) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessageContaining("does not support image embedding"); } @Test @@ -153,8 +191,12 @@ void testNullInputSkipsModelInvocation() { 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(); assertThat(model.prompts).isEmpty(); assertThat(model.embedCalls).isZero(); + assertThat(model.imageTextCalls).isZero(); + assertThat(model.imageEmbedCalls).isZero(); } @Test 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 76cb744a5d1..b6dc28f1ca5 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 @@ -19,6 +19,8 @@ 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; @@ -43,7 +45,8 @@ class AiFunctionParserTest { private static final List COLUMNS = List.of( Column.physicalColumn("id", DataTypes.INT()), - Column.physicalColumn("content", DataTypes.STRING())); + Column.physicalColumn("content", DataTypes.STRING()), + Column.physicalColumn("image", DataTypes.BYTES())); private static class TextModelClient implements AiModelClient, SupportsTextGeneration { private static final long serialVersionUID = 1L; @@ -63,6 +66,26 @@ public float[] embed(String text) { } } + 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 = @@ -104,6 +127,23 @@ void testTranslateSpecializedTextAiFunctions() { .containsOnly(DataTypes.VARIANT()); } + @Test + void testTranslateImageAiFunctions() { + List columns = + translate( + "AI_IMAGE_COMPLETE('vision', image, 'Describe the image') AS description, " + + "AI_IMAGE_EMBED('imageEmbedder', image) AS embedding"); + + assertThat(columns) + .extracting(ProjectionColumn::getScriptExpression) + .containsExactly( + "aiImageComplete(vision, $0, \"Describe the image\")", + "aiImageEmbed(imageEmbedder, $0)"); + assertThat(columns) + .extracting(ProjectionColumn::getDataType) + .containsExactly(DataTypes.STRING(), DataTypes.ARRAY(DataTypes.FLOAT())); + } + @Test void testSameNamedUdfTakesPrecedenceOverAiFunction() { Set udfNames = Set.of("ai_sentiment"); @@ -142,6 +182,47 @@ void testSameNamedUdfTakesPrecedenceOverAiFunction() { .containsExactly(DataTypes.STRING()); } + @Test + void testSameNamedUdfTakesPrecedenceOverImageAiFunction() { + assertSameNamedUdfTakesPrecedenceOverImageAiFunction( + "AI_IMAGE_COMPLETE", "ai_image_complete"); + assertSameNamedUdfTakesPrecedenceOverImageAiFunction("AI_IMAGE_EMBED", "ai_image_embed"); + } + + 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( + projection, + COLUMNS, + List.of( + new UserDefinedFunctionDescriptor( + udfName, + "org.apache.flink.cdc.udf.examples.java.AddOneFunctionClass")), + new SupportedMetadataColumn[0]); + + assertThat(columns) + .extracting(ProjectionColumn::getScriptExpression) + .containsExactly("__udf_" + udfName + ".eval($0)"); + assertThat(columns) + .extracting(ProjectionColumn::getDataType) + .containsExactly(DataTypes.STRING()); + } + @Test void testModelArgumentMustBeStringConstant() { assertThatThrownBy( @@ -190,13 +271,17 @@ void testModelCapabilitiesMustMatchAiFunctions() { Map models = Map.of( "textModel", new TextModelClient(), - "embeddingModel", new EmbeddingModelClient()); + "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_EMBED('embeddingModel', content) AS embedding, " + + "AI_IMAGE_COMPLETE('imageTextModel', image, 'describe') AS description, " + + "AI_IMAGE_EMBED('imageEmbeddingModel', image) AS image_embedding", null, models)) .doesNotThrowAnyException(); @@ -220,6 +305,26 @@ void testModelCapabilitiesMustMatchAiFunctions() { .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 @@ -242,6 +347,11 @@ void testFunctionArityValidation() { .hasMessageContaining("Invalid number of arguments to function 'AI_MASK'"); assertThatThrownBy(() -> translate("AI_SUMMARIZE('model', content, TRUE) AS summarized")) .hasMessageContaining("Cannot apply 'AI_SUMMARIZE'"); + assertThatThrownBy(() -> translate("AI_IMAGE_COMPLETE('model', image) AS description")) + .hasMessageContaining( + "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(