This is an automated email from the ASF dual-hosted git repository. wenjin272 pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit e8f79d51f6e3eef205e3ff073da1124058c59c04 Author: purshotam shah <[email protected]> AuthorDate: Fri Sep 18 10:50:20 2026 -0700 [integrations][java] Add OpenAI embedding model (#1103) Adds integrations/embedding-models/openai with OpenAIEmbeddingModelConnection and OpenAIEmbeddingModelSetup on com.openai:openai-java, mirroring the Python connection and setup arguments (api_key, base_url, request_timeout, max_retries, organization, project; connection, model, encoding_format, dimensions, user, additional_kwargs). Batches are one request placed by response index with size, range and consistency checks; base64 responses are decoded; token usage is reported through embedWithUsage. Arguments are validated once at setup construction and per call by the same package-private parsers; additional_kwargs may not repeat a typed request field. Registers ResourceName.EmbeddingModel.OPENAI_*, the Python mirror and the `openai` YAML alias in both languages, dist and ide-support dependencies. EmbeddingIntegrationTest gains an OPENAI provider (text-embedding-3-small, dimensions 256) gated on a non-blank OPENAI_API_KEY through OllamaPreparationUtils.hasApiKey, with the Ollama case gated on the local daemon only, and the e2e module depends on the new integration. Docs: Java tabs in the OpenAI section of embedding_models.md, FAQ support matrix, YAML alias table. Generated-by: Claude Code 2.1.273 (Claude Fable 5.1) --- .../flink/agents/api/resource/ResourceName.java | 6 + .../org/apache/flink/agents/api/yaml/Aliases.java | 2 + .../apache/flink/agents/api/yaml/AliasesTest.java | 14 + dist/pom.xml | 5 + docs/content/docs/development/embedding_models.md | 99 ++- docs/content/docs/development/yaml.md | 2 +- docs/content/docs/faq/faq.md | 2 +- .../pom.xml | 5 + .../test/EmbeddingIntegrationAgent.java | 13 + .../integration/test/EmbeddingIntegrationTest.java | 86 +- ide-support/pom.xml | 11 + integrations/embedding-models/{ => openai}/pom.xml | 32 +- .../openai/OpenAIEmbeddingModelConnection.java | 556 +++++++++++++ .../openai/OpenAIEmbeddingModelSetup.java | 112 +++ .../openai/OpenAIEmbeddingModelTest.java | 890 +++++++++++++++++++++ integrations/embedding-models/pom.xml | 1 + python/flink_agents/api/resource.py | 4 + python/flink_agents/api/yaml/aliases.py | 2 + python/flink_agents/api/yaml/tests/test_aliases.py | 17 + 19 files changed, 1797 insertions(+), 62 deletions(-) diff --git a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java index 19dc86de..9ccc798f 100644 --- a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java +++ b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java @@ -173,6 +173,12 @@ public final class ResourceName { public static final String BEDROCK_SETUP = "org.apache.flink.agents.integrations.embeddingmodels.bedrock.BedrockEmbeddingModelSetup"; + // OpenAI + public static final String OPENAI_CONNECTION = + "org.apache.flink.agents.integrations.embeddingmodels.openai.OpenAIEmbeddingModelConnection"; + public static final String OPENAI_SETUP = + "org.apache.flink.agents.integrations.embeddingmodels.openai.OpenAIEmbeddingModelSetup"; + // Python Wrapper public static final String PYTHON_WRAPPER_CONNECTION = "org.apache.flink.agents.api.embedding.model.python.PythonEmbeddingModelConnection"; diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java b/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java index ee70e3fe..a82c71e1 100644 --- a/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java +++ b/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java @@ -126,6 +126,7 @@ public final class Aliases { Map<String, String> embConnJava = new HashMap<>(); embConnJava.put("ollama", ResourceName.EmbeddingModel.OLLAMA_CONNECTION); embConnJava.put("bedrock", ResourceName.EmbeddingModel.BEDROCK_CONNECTION); + embConnJava.put("openai", ResourceName.EmbeddingModel.OPENAI_CONNECTION); Map<String, String> embConnPython = new HashMap<>(); embConnPython.put("ollama", ResourceName.EmbeddingModel.Python.OLLAMA_CONNECTION); embConnPython.put("openai", ResourceName.EmbeddingModel.Python.OPENAI_CONNECTION); @@ -138,6 +139,7 @@ public final class Aliases { Map<String, String> embJava = new HashMap<>(); embJava.put("ollama", ResourceName.EmbeddingModel.OLLAMA_SETUP); embJava.put("bedrock", ResourceName.EmbeddingModel.BEDROCK_SETUP); + embJava.put("openai", ResourceName.EmbeddingModel.OPENAI_SETUP); Map<String, String> embPython = new HashMap<>(); embPython.put("ollama", ResourceName.EmbeddingModel.Python.OLLAMA_SETUP); embPython.put("openai", ResourceName.EmbeddingModel.Python.OPENAI_SETUP); diff --git a/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java b/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java index 9d4508c1..152c2e34 100644 --- a/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java +++ b/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java @@ -75,6 +75,20 @@ class AliasesTest { assertThat(azureOpenAISetup).isEqualTo(ResourceName.ChatModel.AZURE_OPENAI_SETUP); } + @Test + void clazzAliasCoversEmbeddingModelOpenAIInBothLanguages() { + assertThat( + Aliases.resolveClazz( + "openai", ResourceType.EMBEDDING_MODEL_CONNECTION, Language.JAVA)) + .isEqualTo(ResourceName.EmbeddingModel.OPENAI_CONNECTION); + assertThat(Aliases.resolveClazz("openai", ResourceType.EMBEDDING_MODEL, Language.JAVA)) + .isEqualTo(ResourceName.EmbeddingModel.OPENAI_SETUP); + assertThat( + Aliases.resolveClazz( + "openai", ResourceType.EMBEDDING_MODEL_CONNECTION, Language.PYTHON)) + .isEqualTo(ResourceName.EmbeddingModel.Python.OPENAI_CONNECTION); + } + @Test void clazzAliasCoversVectorStoreJavaAndPython() { String opensearch = diff --git a/dist/pom.xml b/dist/pom.xml index df52c7a3..b59f25b4 100644 --- a/dist/pom.xml +++ b/dist/pom.xml @@ -107,6 +107,11 @@ under the License. <artifactId>flink-agents-integrations-embedding-models-bedrock</artifactId> <version>${project.version}</version> </dependency> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-agents-integrations-embedding-models-openai</artifactId> + <version>${project.version}</version> + </dependency> <dependency> <groupId>org.apache.flink</groupId> <artifactId>flink-agents-integrations-vector-stores-elasticsearch</artifactId> diff --git a/docs/content/docs/development/embedding_models.md b/docs/content/docs/development/embedding_models.md index 0966b61d..0de2bb5e 100644 --- a/docs/content/docs/development/embedding_models.md +++ b/docs/content/docs/development/embedding_models.md @@ -419,16 +419,15 @@ Model availability and specifications may change. Always check the official Olla OpenAI provides cloud-based embedding models with state-of-the-art performance. -{{< hint info >}} -OpenAI embedding models are currently supported in the Python API only. To use OpenAI from Java agents, see [Using Cross-Language Providers](#using-cross-language-providers). -{{< /hint >}} - #### Prerequisites 1. Get an API key from [OpenAI Platform](https://platform.openai.com/) #### Usage Example +{{< tabs "OpenAI Embedding Usage Example" >}} + +{{< tab "Python" >}} ```python class MyAgent(Agent): @@ -453,28 +452,106 @@ class MyAgent(Agent): encoding_format="float" ) ``` +{{< /tab >}} + +{{< tab "Java" >}} +```java +public class MyAgent extends Agent { + + @EmbeddingModelConnection + public static ResourceDescriptor openaiConnection() { + return ResourceDescriptor.Builder.newBuilder(ResourceName.EmbeddingModel.OPENAI_CONNECTION) + .addInitialArgument("api_key", System.getenv("OPENAI_API_KEY")) + .addInitialArgument("request_timeout", 30) + .addInitialArgument("max_retries", 3) + .build(); + } + + @EmbeddingModelSetup + public static ResourceDescriptor openaiEmbedding() { + return ResourceDescriptor.Builder.newBuilder(ResourceName.EmbeddingModel.OPENAI_SETUP) + .addInitialArgument("connection", "openaiConnection") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("dimensions", 512) + .build(); + } + + ... +} +``` +{{< /tab >}} + +{{< /tabs >}} #### OpenAIEmbeddingModelConnection Parameters +{{< tabs "OpenAIEmbeddingModelConnection Parameters" >}} + +{{< tab "Python" >}} + | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `api_key` | str | Required | OpenAI API key for authentication | +| `api_key` | str | Required | OpenAI API key for authentication; a blank key is passed through to the SDK for unauthenticated OpenAI-compatible servers (the Java connection rejects a blank key) | | `base_url` | str | `"https://api.openai.com/v1"` | OpenAI API base URL | -| `request_timeout` | float | `30.0` | HTTP request timeout in seconds | +| `request_timeout` | float | `30.0` | HTTP request timeout in seconds; `0` disables the timeout | | `max_retries` | int | `3` | Maximum number of retry attempts | -| `organization` | str | None | Optional organization ID | -| `project` | str | None | Optional project ID | +| `organization` | str | None | Optional organization ID; when unset or blank the Python SDK may read `OPENAI_ORG_ID` from the environment (the Java connection sends no organization) | +| `project` | str | None | Optional project ID; when unset or blank the Python SDK may read `OPENAI_PROJECT_ID` from the environment (the Java connection sends no project) | + +{{< /tab >}} + +{{< tab "Java" >}} + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `api_key` | String | Required | OpenAI API key for authentication | +| `base_url` | String | `"https://api.openai.com/v1"` | OpenAI API base URL | +| `request_timeout` | Number | `30` | HTTP request timeout in seconds; `0` disables the timeout | +| `max_retries` | Number | `3` | Maximum number of retry attempts, a non-negative integer | +| `organization` | String | None | Optional organization ID | +| `project` | String | None | Optional project ID | + +{{< /tab >}} + +{{< /tabs >}} #### OpenAIEmbeddingModelSetup Parameters +{{< tabs "OpenAIEmbeddingModelSetup Parameters" >}} + +{{< tab "Python" >}} + | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `connection` | str | Required | Reference to connection method name | | `model` | str | Required | OpenAI embedding model name | -| `encoding_format` | str | `"float"` | Return format ("float" or "base64") | +| `encoding_format` | str | `"float"` | Wire format, `"float"` or `"base64"`; base64 responses are decoded to floats | | `dimensions` | int | None | Output dimensions (text-embedding-3 models only) | | `user` | str | None | End-user identifier for monitoring | -| `additional_kwargs` | dict | `{}` | Additional parameters for the OpenAI embeddings API | +| `additional_kwargs` | dict | `{}` | Extra request body properties sent with every embeddings request; keys may not be `model`, `input`, `encoding_format`, `dimensions` or `user`. A per-call `additional_kwargs` replaces this map | + +`additional_kwargs` entries are sent to the API as request body properties; an entry that repeats a typed field fails validation (set the typed argument instead), and `dimensions` must be an integer. + +{{< /tab >}} + +{{< tab "Java" >}} + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `connection` | String | Required | Reference to connection method name | +| `model` | String | Required | OpenAI embedding model name | +| `encoding_format` | String | `"float"` | Wire format, `"float"` or `"base64"`; base64 responses are decoded to floats | +| `dimensions` | Number | None | Output dimensions, a positive integer (text-embedding-3 models only) | +| `user` | String | None | End-user identifier for monitoring | +| `additional_kwargs` | Map | `{}` | Extra request body properties sent with every embeddings request; keys may not be `model`, `input`, `encoding_format`, `dimensions` or `user`. A per-call `additional_kwargs` replaces this map | + + +`additional_kwargs` entries that repeat a typed field fail validation instead of overriding it (set the typed argument instead), and `dimensions` must be a positive integer. +{{< /tab >}} + +{{< /tabs >}} + +A batch of texts is embedded in a single request, so it must stay within OpenAI's per-request limits (2048 inputs at the time of writing); split larger batches before calling. Token usage reported by the API is available through `embedWithUsage` in Java and `embed_with_usage` in Python. #### Available Models @@ -559,7 +636,7 @@ Model availability and specifications may change. Always check the official Dash ## Using Cross-Language Providers -Flink Agents supports cross-language embedding model integration, allowing you to use embedding models implemented in one language (Java or Python) from agents written in the other language. This is particularly useful when an embedding model provider is only available in one language (e.g., OpenAI embedding is currently Python-only). +Flink Agents supports cross-language embedding model integration, allowing you to use embedding models implemented in one language (Java or Python) from agents written in the other language. This is particularly useful when an embedding model provider is only available in one language (e.g., Tongyi embedding is currently Python-only). {{< hint warning >}} **Limitations:** diff --git a/docs/content/docs/development/yaml.md b/docs/content/docs/development/yaml.md index 21893514..c0aa035c 100644 --- a/docs/content/docs/development/yaml.md +++ b/docs/content/docs/development/yaml.md @@ -585,7 +585,7 @@ Embedding-model aliases (apply to both `embedding_model_connections` and `embedd | Alias | `type: python` | `type: java` | | --------- | -------------- | -------------- | | `ollama` | Ollama | Ollama | -| `openai` | OpenAI | — | +| `openai` | OpenAI | OpenAI (Java) | | `tongyi` | Tongyi | — | | `bedrock` | — | Bedrock (Java) | diff --git a/docs/content/docs/faq/faq.md b/docs/content/docs/faq/faq.md index 3f4fe812..b26367de 100644 --- a/docs/content/docs/faq/faq.md +++ b/docs/content/docs/faq/faq.md @@ -113,7 +113,7 @@ Flink Agents provides built-in integrations for many ecosystem providers. Some i |---|---|---| | [Amazon Bedrock]({{< ref "docs/development/embedding_models#amazon-bedrock" >}}) | ❌ | ✅ | | [Ollama]({{< ref "docs/development/embedding_models#ollama" >}}) | ✅ | ✅ | -| [OpenAI]({{< ref "docs/development/embedding_models#openai" >}}) | ✅ | ❌ | +| [OpenAI]({{< ref "docs/development/embedding_models#openai" >}}) | ✅ | ✅ | | [Tongyi (DashScope)]({{< ref "docs/development/embedding_models#tongyi-dashscope" >}}) | ✅ | ❌ | **Vector Stores** diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml b/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml index 82d4b941..86370c4e 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml +++ b/e2e-test/flink-agents-end-to-end-tests-integration/pom.xml @@ -95,6 +95,11 @@ under the License. <artifactId>flink-agents-integrations-embedding-models-ollama</artifactId> <version>${project.version}</version> </dependency> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-agents-integrations-embedding-models-openai</artifactId> + <version>${project.version}</version> + </dependency> <dependency> <groupId>org.apache.flink</groupId> <artifactId>flink-agents-integrations-vector-stores-elasticsearch</artifactId> diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationAgent.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationAgent.java index 2167d133..e9d8fda2 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationAgent.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationAgent.java @@ -48,6 +48,7 @@ import java.util.Map; */ public class EmbeddingIntegrationAgent extends Agent { public static final String OLLAMA_MODEL = "nomic-embed-text"; + public static final String OPENAI_MODEL = "text-embedding-3-small"; private static final ObjectMapper MAPPER = new ObjectMapper(); @EmbeddingModelConnection @@ -59,6 +60,12 @@ public class EmbeddingIntegrationAgent extends Agent { .addInitialArgument("host", "http://localhost:11434";) .addInitialArgument("timeout", 60) .build(); + } else if (provider.equals("OPENAI")) { + return ResourceDescriptor.Builder.newBuilder( + ResourceName.EmbeddingModel.OPENAI_CONNECTION) + .addInitialArgument("api_key", System.getenv().get("OPENAI_API_KEY")) + .addInitialArgument("request_timeout", 60) + .build(); } else { throw new RuntimeException(String.format("Unknown model provider %s", provider)); } @@ -72,6 +79,12 @@ public class EmbeddingIntegrationAgent extends Agent { .addInitialArgument("connection", "embeddingConnection") .addInitialArgument("model", OLLAMA_MODEL) .build(); + } else if (provider.equals("OPENAI")) { + return ResourceDescriptor.Builder.newBuilder(ResourceName.EmbeddingModel.OPENAI_SETUP) + .addInitialArgument("connection", "embeddingConnection") + .addInitialArgument("model", OPENAI_MODEL) + .addInitialArgument("dimensions", 256) + .build(); } else { throw new RuntimeException(String.format("Unknown model provider %s", provider)); } diff --git a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationTest.java b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationTest.java index baaf2f68..a0d3ad7a 100644 --- a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationTest.java +++ b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/EmbeddingIntegrationTest.java @@ -32,13 +32,13 @@ import java.io.IOException; import java.util.Map; import static org.apache.flink.agents.integration.test.EmbeddingIntegrationAgent.OLLAMA_MODEL; +import static org.apache.flink.agents.integration.test.OllamaPreparationUtils.hasApiKey; import static org.apache.flink.agents.integration.test.OllamaPreparationUtils.pullModel; /** * Example application that applies {@link EmbeddingIntegrationAgent} to a DataStream of prompts. */ public class EmbeddingIntegrationTest { - private static final String API_KEY = "_API_KEY"; private static final String OLLAMA = "OLLAMA"; private final boolean ollamaReady; @@ -48,52 +48,56 @@ public class EmbeddingIntegrationTest { } @ParameterizedTest() - @ValueSource(strings = {"OLLAMA"}) + @ValueSource(strings = {"OLLAMA", "OPENAI"}) public void testEmbeddingIntegration(String provider) throws Exception { Assumptions.assumeTrue( - (OLLAMA.equals(provider) && ollamaReady) - || System.getenv().get(provider + API_KEY) != null, + OLLAMA.equals(provider) ? ollamaReady : hasApiKey(provider), String.format( "Server or authentication information is not provided for %s", provider)); System.setProperty("MODEL_PROVIDER", provider); - - // Create the execution environment - StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); - env.setParallelism(1); - - // Use prompts that exercise embedding generation and similarity checks - DataStream<String> inputStream = - env.fromData( - "Generate embedding for: 'Machine learning'", - "Generate embedding for: 'Deep learning techniques'", - "Find texts similar to: 'neural networks'", - "Produce embedding and return top-3 similar items for: 'natural language processing'", - "Generate embedding for: 'hello world'", - "Compare similarity between 'cat' and 'dog'", - "Create embedding for: 'space exploration'", - "Find nearest neighbors for: 'artificial intelligence'", - "Generate embedding for: 'data science'", - "Random embedding test"); - - // Create agents execution environment - AgentsExecutionEnvironment agentsEnv = - AgentsExecutionEnvironment.getExecutionEnvironment(env); - - // Apply agent to the DataStream and use the prompt itself as the key - DataStream<Object> outputStream = - agentsEnv - .fromDataStream(inputStream, (KeySelector<String, String>) value -> value) - .apply(new EmbeddingIntegrationAgent()) - .toDataStream(); - - // Collect the results - CloseableIterator<Object> results = outputStream.collectAsync(); - - // Execute the pipeline - agentsEnv.execute(); - - checkResult(results); + try { + + // Create the execution environment + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + // Use prompts that exercise embedding generation and similarity checks + DataStream<String> inputStream = + env.fromData( + "Generate embedding for: 'Machine learning'", + "Generate embedding for: 'Deep learning techniques'", + "Find texts similar to: 'neural networks'", + "Produce embedding and return top-3 similar items for: 'natural language processing'", + "Generate embedding for: 'hello world'", + "Compare similarity between 'cat' and 'dog'", + "Create embedding for: 'space exploration'", + "Find nearest neighbors for: 'artificial intelligence'", + "Generate embedding for: 'data science'", + "Random embedding test"); + + // Create agents execution environment + AgentsExecutionEnvironment agentsEnv = + AgentsExecutionEnvironment.getExecutionEnvironment(env); + + // Apply agent to the DataStream and use the prompt itself as the key + DataStream<Object> outputStream = + agentsEnv + .fromDataStream( + inputStream, (KeySelector<String, String>) value -> value) + .apply(new EmbeddingIntegrationAgent()) + .toDataStream(); + + // Collect the results + CloseableIterator<Object> results = outputStream.collectAsync(); + + // Execute the pipeline + agentsEnv.execute(); + + checkResult(results); + } finally { + System.clearProperty("MODEL_PROVIDER"); + } } @SuppressWarnings("unchecked") diff --git a/ide-support/pom.xml b/ide-support/pom.xml index 5ce239a9..d7b07853 100644 --- a/ide-support/pom.xml +++ b/ide-support/pom.xml @@ -90,6 +90,17 @@ under the License. </exclusion> </exclusions> </dependency> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-agents-integrations-embedding-models-openai</artifactId> + <version>${project.version}</version> + <exclusions> + <exclusion> + <groupId>org.apache.flink</groupId> + <artifactId>flink-agents-api</artifactId> + </exclusion> + </exclusions> + </dependency> <dependency> <groupId>org.apache.flink</groupId> <artifactId>flink-agents-integrations-vector-stores-elasticsearch</artifactId> diff --git a/integrations/embedding-models/pom.xml b/integrations/embedding-models/openai/pom.xml similarity index 55% copy from integrations/embedding-models/pom.xml copy to integrations/embedding-models/openai/pom.xml index 2a585d17..f905294d 100644 --- a/integrations/embedding-models/pom.xml +++ b/integrations/embedding-models/openai/pom.xml @@ -22,17 +22,33 @@ under the License. <parent> <groupId>org.apache.flink</groupId> - <artifactId>flink-agents-integrations</artifactId> + <artifactId>flink-agents-integrations-embedding-models</artifactId> <version>0.4-SNAPSHOT</version> + <relativePath>../pom.xml</relativePath> </parent> - <artifactId>flink-agents-integrations-embedding-models</artifactId> - <name>Flink Agents : Integrations: Embedding Models</name> - <packaging>pom</packaging> + <artifactId>flink-agents-integrations-embedding-models-openai</artifactId> + <name>Flink Agents : Integrations: Embedding Models: OpenAI</name> + <packaging>jar</packaging> - <modules> - <module>bedrock</module> - <module>ollama</module> - </modules> + <dependencies> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-agents-api</artifactId> + <version>${project.version}</version> + </dependency> + + <dependency> + <groupId>com.openai</groupId> + <artifactId>openai-java</artifactId> + <version>${openai.version}</version> + </dependency> + <dependency> + <groupId>org.apache.flink</groupId> + <artifactId>flink-core</artifactId> + <version>${flink.version}</version> + <scope>test</scope> + </dependency> + </dependencies> </project> diff --git a/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelConnection.java b/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelConnection.java new file mode 100644 index 00000000..c58aa952 --- /dev/null +++ b/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelConnection.java @@ -0,0 +1,556 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.embeddingmodels.openai; + +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.core.JsonField; +import com.openai.core.JsonString; +import com.openai.core.JsonValue; +import com.openai.core.Timeout; +import com.openai.models.embeddings.CreateEmbeddingResponse; +import com.openai.models.embeddings.Embedding; +import com.openai.models.embeddings.EmbeddingCreateParams; +import com.openai.models.embeddings.EmbeddingValue; +import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection; +import org.apache.flink.agents.api.embedding.model.EmbeddingModelUtils; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; +import org.apache.flink.agents.api.embedding.model.EmbeddingTokenUsage; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; + +import java.math.BigDecimal; +import java.math.BigInteger; +import java.math.RoundingMode; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.TreeSet; + +/** + * Embedding model connection backed by the OpenAI Embeddings API. + * + * <p>Argument names mirror the Python {@code OpenAIEmbeddingModelConnection}: + * + * <ul> + * <li>{@code api_key} (required): OpenAI API key. + * <li>{@code base_url}: API base URL, default {@code https://api.openai.com/v1}. + * <li>{@code request_timeout}: per-attempt request timeout in seconds, default {@code 30}; {@code + * 0} disables the timeout, as in the other OpenAI connections. + * <li>{@code max_retries}: SDK retry count for failed requests, default {@code 3}. + * <li>{@code organization}, {@code project}: optional OpenAI organization and project IDs. + * </ul> + * + * <p>Per-request parameters come from {@link OpenAIEmbeddingModelSetup#getParameters()} merged with + * any per-call parameters: {@code model} (required), {@code encoding_format} ({@code "float"} or + * {@code "base64"}), {@code dimensions}, {@code user}, and {@code additional_kwargs}, a map whose + * entries are sent as extra request body properties. Its keys must not collide with the typed + * request fields ({@code model}, {@code input}, {@code encoding_format}, {@code dimensions}, {@code + * user}); entries with a {@code null} value are omitted. Blank strings are treated as absent, and + * other per-call keys are ignored. A batch of texts is sent as one request and the embeddings are + * returned in input order. Token usage reported by the API is surfaced through {@link + * #embedWithUsage}; servers that omit it yield a {@code null} usage. + */ +public class OpenAIEmbeddingModelConnection extends BaseEmbeddingModelConnection { + + static final String DEFAULT_BASE_URL = "https://api.openai.com/v1";; + static final int DEFAULT_REQUEST_TIMEOUT_SECONDS = 30; + static final int DEFAULT_MAX_RETRIES = 3; + /** The Java SDK stores timeouts as int milliseconds. */ + static final BigDecimal MAX_REQUEST_TIMEOUT_SECONDS = new BigDecimal("2147483.647"); + + static final String MODEL = "model"; + static final String ENCODING_FORMAT = "encoding_format"; + static final String DIMENSIONS = "dimensions"; + static final String USER = "user"; + /** Setup parameter carrying extra request body properties. */ + static final String ADDITIONAL_KWARGS = "additional_kwargs"; + + static final Set<String> ENCODING_FORMATS = Set.of("float", "base64"); + /** + * Request fields owned by the typed builder calls; {@code additional_kwargs} may not set them. + */ + static final Set<String> RESERVED_ADDITIONAL_KWARGS = + Set.of(MODEL, "input", ENCODING_FORMAT, DIMENSIONS, USER); + + private final OpenAIClient client; + + public OpenAIEmbeddingModelConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + this.client = ClientConfig.parse(descriptor).buildClient(); + } + + // visible for testing + OpenAIEmbeddingModelConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext, OpenAIClient client) { + super(descriptor, resourceContext); + this.client = client; + } + + /** The parsed connection arguments; tests assert on them through {@link #parse}. */ + static final class ClientConfig { + final String apiKey; + final String baseUrl; + final Duration requestTimeout; + final int maxRetries; + final String organization; + final String project; + + private ClientConfig( + String apiKey, + String baseUrl, + Duration requestTimeout, + int maxRetries, + String organization, + String project) { + this.apiKey = apiKey; + this.baseUrl = baseUrl; + this.requestTimeout = requestTimeout; + this.maxRetries = maxRetries; + this.organization = organization; + this.project = project; + } + + /** Blank string arguments are treated as absent. */ + static ClientConfig parse(ResourceDescriptor descriptor) { + String apiKey = requireString(descriptor.getArgument("api_key"), "api_key"); + if (apiKey == null || apiKey.isBlank()) { + throw new IllegalArgumentException("api_key should not be null or empty."); + } + String baseUrl = requireString(descriptor.getArgument("base_url"), "base_url"); + return new ClientConfig( + apiKey, + baseUrl == null || baseUrl.isBlank() ? DEFAULT_BASE_URL : baseUrl, + parseRequestTimeout(descriptor.getArgument("request_timeout")), + parseMaxRetries(descriptor.getArgument("max_retries")), + blankToNull( + requireString(descriptor.getArgument("organization"), "organization")), + blankToNull(requireString(descriptor.getArgument("project"), "project"))); + } + + OpenAIClient buildClient() { + OpenAIOkHttpClient.Builder builder = + new OpenAIOkHttpClient.Builder() + .apiKey(apiKey) + .baseUrl(baseUrl) + .timeout( + Timeout.builder() + .connect(requestTimeout) + .read(requestTimeout) + .write(requestTimeout) + .request(requestTimeout) + .build()) + .maxRetries(maxRetries); + if (organization != null) { + builder.organization(organization); + } + if (project != null) { + builder.project(project); + } + return builder.build(); + } + } + + /** Blank string arguments are treated as absent. */ + static String blankToNull(String value) { + return value == null || value.isBlank() ? null : value; + } + + /** Parses {@code request_timeout} seconds with the same rules as the OpenAI chat connection. */ + static Duration parseRequestTimeout(Object raw) { + if (raw == null) { + return Duration.ofSeconds(DEFAULT_REQUEST_TIMEOUT_SECONDS); + } + BigDecimal seconds = toBigDecimal(raw, "request_timeout"); + if (seconds.signum() < 0) { + throw new IllegalArgumentException("request_timeout must be >= 0, got: " + raw); + } + if (seconds.compareTo(MAX_REQUEST_TIMEOUT_SECONDS) > 0) { + throw new IllegalArgumentException( + "request_timeout exceeds the SDK maximum of " + + MAX_REQUEST_TIMEOUT_SECONDS.toPlainString() + + " seconds, got: " + + raw); + } + // Millisecond precision; positive values round up so a valid timeout never becomes zero, + // which would disable it. Exactly 0 disables the timeout, like the chat connection. + BigInteger millis = + seconds.multiply(BigDecimal.valueOf(1_000L)) + .setScale(0, RoundingMode.CEILING) + .toBigIntegerExact(); + return Duration.ofMillis(millis.longValueExact()); + } + + static int parseMaxRetries(Object raw) { + if (raw == null) { + return DEFAULT_MAX_RETRIES; + } + BigDecimal value = toBigDecimal(raw, "max_retries"); + try { + BigInteger retries = value.toBigIntegerExact(); + if (retries.signum() < 0 + || retries.compareTo(BigInteger.valueOf(Integer.MAX_VALUE)) > 0) { + throw new IllegalArgumentException( + "max_retries must be a non-negative integer, got: " + raw); + } + return retries.intValueExact(); + } catch (ArithmeticException e) { + throw new IllegalArgumentException( + "max_retries must be a non-negative integer, got: " + raw, e); + } + } + + private static BigDecimal toBigDecimal(Object raw, String argumentName) { + if (!(raw instanceof Number)) { + throw new IllegalArgumentException(argumentName + " must be a number, got: " + raw); + } + Number number = (Number) raw; + if ((number instanceof Double || number instanceof Float) + && !Double.isFinite(number.doubleValue())) { + throw new IllegalArgumentException(argumentName + " must be finite, got: " + raw); + } + try { + return new BigDecimal(number.toString()); + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + argumentName + " must be a finite number, got: " + raw, e); + } + } + + // ---- request parameter parsing, shared by the setup (fail fast) and buildParams (per call) + + static String requireString(Object value, String argumentName) { + if (value != null && !(value instanceof String)) { + throw new IllegalArgumentException( + argumentName + " must be a string, got: " + value.getClass().getSimpleName()); + } + return (String) value; + } + + static String parseModel(Object value) { + String model = requireString(value, MODEL); + if (model == null || model.isBlank()) { + throw new IllegalArgumentException( + "OpenAI embedding requires a non-empty 'model' (setup or per-call parameter)."); + } + return model; + } + + static String parseEncodingFormat(Object value) { + String format = requireString(value, ENCODING_FORMAT); + if (format == null || format.isBlank()) { + return OpenAIEmbeddingModelSetup.DEFAULT_ENCODING_FORMAT; + } + if (!ENCODING_FORMATS.contains(format)) { + throw new IllegalArgumentException( + ENCODING_FORMAT + " must be one of " + ENCODING_FORMATS + ", got: " + format); + } + return format; + } + + static Integer parseDimensions(Object value) { + if (value == null) { + return null; + } + BigDecimal number = toBigDecimal(value, DIMENSIONS); + try { + BigInteger dims = number.toBigIntegerExact(); + if (dims.signum() <= 0 || dims.compareTo(BigInteger.valueOf(Integer.MAX_VALUE)) > 0) { + throw new IllegalArgumentException( + DIMENSIONS + " must be a positive integer, got: " + value); + } + return dims.intValueExact(); + } catch (ArithmeticException e) { + throw new IllegalArgumentException( + DIMENSIONS + " must be a positive integer, got: " + value, e); + } + } + + static Map<String, Object> parseAdditionalKwargs(Object value) { + Map<String, Object> result = new LinkedHashMap<>(); + if (value == null) { + return result; + } + if (!(value instanceof Map)) { + throw new IllegalArgumentException( + ADDITIONAL_KWARGS + " must be a map, got: " + value.getClass().getSimpleName()); + } + Set<String> collisions = null; + for (Map.Entry<?, ?> entry : ((Map<?, ?>) value).entrySet()) { + if (entry.getKey() == null || entry.getKey().toString().isBlank()) { + throw new IllegalArgumentException(ADDITIONAL_KWARGS + " contains an empty key."); + } + String key = entry.getKey().toString(); + if (RESERVED_ADDITIONAL_KWARGS.contains(key)) { + if (collisions == null) { + collisions = new TreeSet<>(); + } + collisions.add(key); + } + result.put(key, entry.getValue()); + } + if (collisions != null) { + throw new IllegalArgumentException( + ADDITIONAL_KWARGS + + " must not contain the typed request fields " + + collisions + + "; set them through the corresponding setup argument instead."); + } + return result; + } + + @Override + public float[] embed(String text, Map<String, Object> parameters) { + return embedWithUsage(text, parameters).getEmbeddings(); + } + + @Override + public List<float[]> embed(List<String> texts, Map<String, Object> parameters) { + return embedWithUsage(texts, parameters).getEmbeddings(); + } + + @Override + public EmbeddingResult<float[]> embedWithUsage(String text, Map<String, Object> parameters) { + EmbeddingResult<List<float[]>> result = + embedWithUsage(Collections.singletonList(text), parameters); + return new EmbeddingResult<>(result.getEmbeddings().get(0), result.getTokenUsage()); + } + + @Override + public EmbeddingResult<List<float[]>> embedWithUsage( + List<String> texts, Map<String, Object> parameters) { + if (texts == null) { + throw new IllegalArgumentException("texts must not be null."); + } + List<String> inputs = texts; + for (int i = 0; i < inputs.size(); i++) { + if (inputs.get(i) == null) { + throw new IllegalArgumentException( + "Text at index " + i + " is null; every input must be a string."); + } + } + // Validate the parameters even for an empty batch so a bad per-call override is not + // deferred to the first non-empty batch (the Python connection validates first too). + EmbeddingCreateParams params = + buildParams(inputs, parameters == null ? Collections.emptyMap() : parameters); + if (inputs.isEmpty()) { + return new EmbeddingResult<>(new ArrayList<>(), null); + } + CreateEmbeddingResponse response = client.embeddings().create(params); + + // A compatible server may answer 200 without `data`; the count check below reports it + // instead of the SDK's required-field exception. + JsonField<List<Embedding>> dataField = response._data(); + List<Embedding> data; + if (dataField.isMissing() || dataField.isNull()) { + data = Collections.emptyList(); + } else { + // Present but not a list of embeddings: report the payload, not a count of zero. + data = + dataField + .asKnown() + .orElseThrow( + () -> + new RuntimeException( + "OpenAI returned a malformed data field: " + + dataField.asUnknown().orElse(null))); + } + if (data.size() != texts.size()) { + throw new RuntimeException( + String.format( + "OpenAI returned %d embeddings for %d input texts.", + data.size(), texts.size())); + } + boolean base64 = + params.encodingFormat() + .map(EmbeddingCreateParams.EncodingFormat.BASE64::equals) + .orElse(false); + return new EmbeddingResult<>(orderEmbeddings(data, base64), extractTokenUsage(response)); + } + + /** + * Places each vector at its response {@code index} so input order is preserved even if the API + * reorders; an item without an index (some OpenAI-compatible servers omit it) keeps its + * response position. Every slot must be filled exactly once, so a response whose indices are + * out of range, duplicated, or contradict the positions of index-less items is rejected rather + * than silently mapped. + */ + private static List<float[]> orderEmbeddings(List<Embedding> data, boolean base64) { + float[][] ordered = new float[data.size()][]; + for (int position = 0; position < data.size(); position++) { + Embedding embedding = data.get(position); + if (embedding == null) { + throw new RuntimeException( + "OpenAI returned a malformed embedding at position " + position + "."); + } + JsonField<Long> index = embedding._index(); + long target; + if (index.isMissing() || index.isNull()) { + target = position; + } else { + // A present but non-integer index is a malformed response, not an absent one. + Optional<Long> known = index.asKnown(); + if (known.isEmpty()) { + throw new RuntimeException( + String.format( + "OpenAI returned a non-integer embedding index %s at position %d.", + index.asUnknown().orElse(null), position)); + } + target = known.get(); + } + if (target < 0 || target >= data.size() || ordered[(int) target] != null) { + throw new RuntimeException( + String.format( + "OpenAI returned an unexpected embedding index %d at position %d.", + target, position)); + } + try { + ordered[(int) target] = toVector(embedding, base64); + } catch (RuntimeException e) { + // A null, non-numeric or undecodable vector; name the position as Python does. + throw new RuntimeException( + "OpenAI returned a malformed embedding at position " + position + ".", e); + } + } + return new ArrayList<>(List.of(ordered)); + } + + /** + * Decodes one response vector. Base64 payloads are decoded here rather than through the SDK, + * which silently drops trailing bytes when the length is not a multiple of four; the wire + * format is little-endian float32. + */ + private static float[] toVector(Embedding embedding, boolean base64Requested) { + EmbeddingValue value = embedding.embeddingValue(); + if (value.isBase64()) { + if (!base64Requested) { + // A string vector is only meaningful when base64 was asked for. + throw new IllegalArgumentException( + "a string vector was returned although encoding_format was float."); + } + byte[] bytes = Base64.getDecoder().decode(value.asBase64()); + if (bytes.length % Float.BYTES != 0) { + throw new IllegalArgumentException( + "base64 embedding has " + bytes.length + " bytes, not a multiple of 4."); + } + float[] vector = new float[bytes.length / Float.BYTES]; + ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer().get(vector); + return vector; + } + return EmbeddingModelUtils.toFloatArray(value.asFloats()); + } + + /** + * Converts an {@code additional_kwargs} value, naming the key if the SDK cannot serialize it. + */ + static JsonValue toJsonValue(String key, Object value) { + try { + return JsonValue.from(value); + } catch (RuntimeException e) { + throw new IllegalArgumentException( + ADDITIONAL_KWARGS + + " value for key '" + + key + + "' cannot be sent as JSON: " + + e.getMessage(), + e); + } + } + + static EmbeddingCreateParams buildParams(List<String> texts, Map<String, Object> parameters) { + EmbeddingCreateParams.Builder builder = + EmbeddingCreateParams.builder() + .model(parseModel(parameters.get(MODEL))) + .inputOfArrayOfStrings(texts) + .encodingFormat( + EmbeddingCreateParams.EncodingFormat.of( + parseEncodingFormat(parameters.get(ENCODING_FORMAT)))); + Integer dimensions = parseDimensions(parameters.get(DIMENSIONS)); + if (dimensions != null) { + builder.dimensions(dimensions); + } + String user = requireString(parameters.get(USER), USER); + if (user != null && !user.isBlank()) { + builder.user(user); + } + for (Map.Entry<String, Object> entry : + parseAdditionalKwargs(parameters.get(ADDITIONAL_KWARGS)).entrySet()) { + if (entry.getValue() != null) { + builder.putAdditionalBodyProperty( + entry.getKey(), toJsonValue(entry.getKey(), entry.getValue())); + } + } + return builder.build(); + } + + /** + * Tolerates servers that omit {@code usage} or some of its fields. Embedding requests have no + * completion tokens, so a missing side of the usage equals the other side (as the Tongyi + * connection does) rather than being reported as zero. + */ + private static EmbeddingTokenUsage extractTokenUsage(CreateEmbeddingResponse response) { + Optional<CreateEmbeddingResponse.Usage> usage = response._usage().asKnown(); + if (usage.isEmpty()) { + return null; + } + Optional<Long> prompt = usageCount(usage.get()._promptTokens()); + Optional<Long> total = usageCount(usage.get()._totalTokens()); + if (prompt.isEmpty() && total.isEmpty()) { + return null; + } + long promptTokens = prompt.orElseGet(total::get); + return new EmbeddingTokenUsage(promptTokens, total.orElse(promptTokens)); + } + + /** + * A token count, also accepting the decimal-string form some compatible servers send (as the + * Python connection does); anything else is treated as absent. + */ + private static Optional<Long> usageCount(JsonField<Long> field) { + Optional<Long> known = field.asKnown(); + if (known.isPresent()) { + return known; + } + return field.asUnknown() + .filter(value -> value instanceof JsonString) + .map(value -> ((JsonString) value).value().trim()) + .filter(text -> !text.isEmpty() && text.chars().allMatch(Character::isDigit)) + .flatMap( + text -> { + try { + return Optional.of(Long.parseLong(text)); + } catch (NumberFormatException e) { + return Optional.empty(); + } + }); + } + + @Override + public void close() { + client.close(); + } +} diff --git a/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelSetup.java b/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelSetup.java new file mode 100644 index 00000000..842b20e1 --- /dev/null +++ b/integrations/embedding-models/openai/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelSetup.java @@ -0,0 +1,112 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.embeddingmodels.openai; + +import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; + +import java.util.HashMap; +import java.util.Map; + +/** + * Embedding model setup for the OpenAI Embeddings API. + * + * <p>Arguments mirror the Python {@code OpenAIEmbeddingModelSetup}: {@code connection} and {@code + * model} (required), {@code encoding_format} (default {@code "float"}, or {@code "base64"}), {@code + * dimensions} (text-embedding-3 models only), {@code user}, and {@code additional_kwargs}, a map of + * extra request body properties sent with every request whose keys may not repeat the typed request + * fields. Blank string arguments are treated as absent. Per-call parameters override these entry by + * entry, so a per-call {@code additional_kwargs} map replaces the setup's map rather than merging + * with it. + */ +public class OpenAIEmbeddingModelSetup extends BaseEmbeddingModelSetup { + + static final String DEFAULT_ENCODING_FORMAT = "float"; + static final String CONNECTION = "connection"; + + private final String encodingFormat; + private final Integer dimensions; + private final String user; + private final Map<String, Object> additionalKwargs; + + public OpenAIEmbeddingModelSetup( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + // Fail fast at construction; the connection re-applies the same rules per call so that + // per-call overrides are validated identically. + super(validate(descriptor), resourceContext); + this.encodingFormat = + OpenAIEmbeddingModelConnection.parseEncodingFormat( + descriptor.getArgument(OpenAIEmbeddingModelConnection.ENCODING_FORMAT)); + this.dimensions = + OpenAIEmbeddingModelConnection.parseDimensions( + descriptor.getArgument(OpenAIEmbeddingModelConnection.DIMENSIONS)); + this.user = + OpenAIEmbeddingModelConnection.blankToNull( + OpenAIEmbeddingModelConnection.requireString( + descriptor.getArgument(OpenAIEmbeddingModelConnection.USER), + OpenAIEmbeddingModelConnection.USER)); + this.additionalKwargs = + OpenAIEmbeddingModelConnection.parseAdditionalKwargs( + descriptor.getArgument(OpenAIEmbeddingModelConnection.ADDITIONAL_KWARGS)); + // Values are sent as JSON on every request; probe the conversion here so a value the SDK + // cannot + // serialize fails at construction rather than on the first embed call. + for (Map.Entry<String, Object> entry : additionalKwargs.entrySet()) { + if (entry.getValue() != null) { + OpenAIEmbeddingModelConnection.toJsonValue(entry.getKey(), entry.getValue()); + } + } + } + + /** + * Validates {@code connection} and {@code model} before {@link BaseEmbeddingModelSetup} casts + * them to Strings, so a missing or non-string value fails here with a named error instead of a + * {@link ClassCastException} or a lookup failure at {@code open()}. + */ + private static ResourceDescriptor validate(ResourceDescriptor descriptor) { + Object connection = descriptor.getArgument(CONNECTION); + String name = OpenAIEmbeddingModelConnection.requireString(connection, CONNECTION); + if (name == null || name.isBlank()) { + throw new IllegalArgumentException( + "OpenAI embedding requires a non-empty 'connection' on the embedding model setup."); + } + Object model = descriptor.getArgument(OpenAIEmbeddingModelConnection.MODEL); + OpenAIEmbeddingModelConnection.parseModel(model); + return descriptor; + } + + @Override + public Map<String, Object> getParameters() { + Map<String, Object> parameters = new HashMap<>(); + parameters.put(OpenAIEmbeddingModelConnection.MODEL, model); + parameters.put(OpenAIEmbeddingModelConnection.ENCODING_FORMAT, encodingFormat); + if (dimensions != null) { + parameters.put(OpenAIEmbeddingModelConnection.DIMENSIONS, dimensions); + } + if (user != null) { + parameters.put(OpenAIEmbeddingModelConnection.USER, user); + } + if (!additionalKwargs.isEmpty()) { + parameters.put( + OpenAIEmbeddingModelConnection.ADDITIONAL_KWARGS, + new HashMap<>(additionalKwargs)); + } + return parameters; + } +} diff --git a/integrations/embedding-models/openai/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelTest.java b/integrations/embedding-models/openai/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelTest.java new file mode 100644 index 00000000..892e721f --- /dev/null +++ b/integrations/embedding-models/openai/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/openai/OpenAIEmbeddingModelTest.java @@ -0,0 +1,890 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.embeddingmodels.openai; + +import com.openai.client.OpenAIClient; +import com.openai.core.JsonField; +import com.openai.core.JsonMissing; +import com.openai.core.JsonValue; +import com.openai.models.embeddings.CreateEmbeddingResponse; +import com.openai.models.embeddings.Embedding; +import com.openai.models.embeddings.EmbeddingCreateParams; +import com.openai.models.embeddings.EmbeddingValue; +import com.openai.services.blocking.EmbeddingService; +import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection; +import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup; +import org.apache.flink.agents.api.embedding.model.EmbeddingResult; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.time.Duration; +import java.util.Base64; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class OpenAIEmbeddingModelTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + + private static ResourceDescriptor connDescriptor() { + return ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelConnection.class.getName()) + .addInitialArgument("api_key", "test-key") + .build(); + } + + private static Embedding embedding(long index, float... values) { + Embedding.Builder builder = + Embedding.builder().index(index).object_(JsonValue.from("embedding")); + List<Float> boxed = new java.util.ArrayList<>(); + for (float v : values) { + boxed.add(v); + } + return builder.embedding(boxed).build(); + } + + private static CreateEmbeddingResponse response(long promptTokens, Embedding... data) { + CreateEmbeddingResponse.Builder builder = + CreateEmbeddingResponse.builder() + .model("text-embedding-3-small") + .object_(JsonValue.from("list")) + .usage( + CreateEmbeddingResponse.Usage.builder() + .promptTokens(promptTokens) + .totalTokens(promptTokens) + .build()); + for (Embedding e : data) { + builder.addData(e); + } + return builder.build(); + } + + private static OpenAIEmbeddingModelConnection connectionWith( + OpenAIClient client, EmbeddingService service) { + when(client.embeddings()).thenReturn(service); + return new OpenAIEmbeddingModelConnection(connDescriptor(), NOOP, client); + } + + @Test + @DisplayName("Connection constructor builds a client with defaults") + void testConnectionDefaults() { + OpenAIEmbeddingModelConnection conn = + new OpenAIEmbeddingModelConnection(connDescriptor(), NOOP); + assertThat(conn).isInstanceOf(BaseEmbeddingModelConnection.class); + OpenAIEmbeddingModelConnection.ClientConfig config = + OpenAIEmbeddingModelConnection.ClientConfig.parse(connDescriptor()); + assertThat(config.apiKey).isEqualTo("test-key"); + assertThat(config.baseUrl).isEqualTo(OpenAIEmbeddingModelConnection.DEFAULT_BASE_URL); + assertThat(config.requestTimeout).isEqualTo(Duration.ofSeconds(30)); + assertThat(config.maxRetries).isEqualTo(3); + assertThat(config.organization).isNull(); + assertThat(config.project).isNull(); + conn.close(); + } + + @Test + @DisplayName("Connection accepts organization, project, base_url and numeric options") + void testConnectionExplicitParams() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder( + OpenAIEmbeddingModelConnection.class.getName()) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("base_url", "https://example.invalid/v1";) + .addInitialArgument("request_timeout", 12.5) + .addInitialArgument("max_retries", 0) + .addInitialArgument("organization", "org-1") + .addInitialArgument("project", "proj-1") + .build(); + OpenAIEmbeddingModelConnection conn = new OpenAIEmbeddingModelConnection(desc, NOOP); + OpenAIEmbeddingModelConnection.ClientConfig config = + OpenAIEmbeddingModelConnection.ClientConfig.parse(desc); + assertThat(config.baseUrl).isEqualTo("https://example.invalid/v1";); + assertThat(config.requestTimeout).isEqualTo(Duration.ofMillis(12_500)); + assertThat(config.maxRetries).isZero(); + assertThat(config.organization).isEqualTo("org-1"); + assertThat(config.project).isEqualTo("proj-1"); + conn.close(); + + ResourceDescriptor blanks = + ResourceDescriptor.Builder.newBuilder( + OpenAIEmbeddingModelConnection.class.getName()) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("base_url", " ") + .addInitialArgument("organization", "") + .build(); + OpenAIEmbeddingModelConnection.ClientConfig blankConfig = + OpenAIEmbeddingModelConnection.ClientConfig.parse(blanks); + assertThat(blankConfig.baseUrl).isEqualTo(OpenAIEmbeddingModelConnection.DEFAULT_BASE_URL); + assertThat(blankConfig.organization).isNull(); + } + + @Test + @DisplayName("Missing api_key is rejected at construction") + void testMissingApiKey() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder( + OpenAIEmbeddingModelConnection.class.getName()) + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelConnection(desc, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("api_key"); + } + + @Test + @DisplayName( + "request_timeout and max_retries follow the OpenAI chat connection's parsing rules") + void testNumericOptionParsing() { + assertThat(OpenAIEmbeddingModelConnection.parseRequestTimeout(null)) + .isEqualTo(Duration.ofSeconds(30)); + assertThat(OpenAIEmbeddingModelConnection.parseRequestTimeout(1.5)) + .isEqualTo(Duration.ofMillis(1500)); + // sub-millisecond values round up instead of collapsing to "disabled" + assertThat(OpenAIEmbeddingModelConnection.parseRequestTimeout(0.0001)) + .isEqualTo(Duration.ofMillis(1)); + assertThat(OpenAIEmbeddingModelConnection.parseRequestTimeout(0)).isEqualTo(Duration.ZERO); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseRequestTimeout(-1)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseRequestTimeout(2147483.648)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseRequestTimeout("30")) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseRequestTimeout(Double.NaN)) + .isInstanceOf(IllegalArgumentException.class); + assertThat(OpenAIEmbeddingModelConnection.parseMaxRetries(null)).isEqualTo(3); + assertThat(OpenAIEmbeddingModelConnection.parseMaxRetries(7L)).isEqualTo(7); + assertThat(OpenAIEmbeddingModelConnection.parseMaxRetries(2.0)).isEqualTo(2); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseMaxRetries(1.5)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseMaxRetries(-1)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("dimensions must be a positive integer within int range, also per call") + void testDimensionsValidation() { + assertThat(OpenAIEmbeddingModelConnection.parseDimensions(256L)).isEqualTo(256); + assertThat(OpenAIEmbeddingModelConnection.parseDimensions(256.0)).isEqualTo(256); + for (Object bad : new Object[] {256.7, 0, -1, 4294967552L, 1e10, "256"}) { + assertThatThrownBy(() -> OpenAIEmbeddingModelConnection.parseDimensions(bad)) + .as("dimensions=" + bad) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dimensions"); + } + // per-call overrides go through the same rules as the setup + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + assertThatThrownBy( + () -> + conn.embed( + "hello", + Map.of( + "model", + "text-embedding-3-small", + "dimensions", + 256.7))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dimensions"); + } + + @Test + @DisplayName("Setup getParameters mirrors the Python model_kwargs") + void testSetupParameters() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("dimensions", 256) + .addInitialArgument("user", "user-42") + .addInitialArgument("additional_kwargs", Map.of("custom", "x")) + .build(); + OpenAIEmbeddingModelSetup setup = new OpenAIEmbeddingModelSetup(desc, NOOP); + assertThat(setup).isInstanceOf(BaseEmbeddingModelSetup.class); + assertThat(setup.getParameters()) + .containsEntry("model", "text-embedding-3-small") + .containsEntry("encoding_format", "float") + .containsEntry("dimensions", 256) + .containsEntry("user", "user-42") + .containsEntry("additional_kwargs", Map.of("custom", "x")); + } + + @Test + @DisplayName("Setup omits optional parameters that were not set") + void testSetupParametersDefaults() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .build(); + Map<String, Object> params = new OpenAIEmbeddingModelSetup(desc, NOOP).getParameters(); + assertThat(params) + .containsOnlyKeys("model", "encoding_format") + .containsEntry("encoding_format", "float"); + ResourceDescriptor blankUser = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("user", " ") + .build(); + // Blank strings are treated as absent, so a blank user is not emitted. + assertThat(new OpenAIEmbeddingModelSetup(blankUser, NOOP).getParameters()) + .doesNotContainKey("user"); + } + + @Test + @DisplayName("Single embedding returns the vector and token usage") + void testSingleEmbedding() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(5, embedding(0, 0.1f, 0.2f, 0.3f))); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + EmbeddingResult<float[]> result = + conn.embedWithUsage("hello", Map.of("model", "text-embedding-3-small")); + + assertThat(result.getEmbeddings()).containsExactly(0.1f, 0.2f, 0.3f); + assertThat(result.getTokenUsage()).isNotNull(); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(5L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(5L); + assertThat(conn.embed("hello", Map.of("model", "text-embedding-3-small"))) + .containsExactly(0.1f, 0.2f, 0.3f); + } + + @Test + @DisplayName("Batch embedding is one request and preserves input order") + void testBatchEmbeddingOrderAndSingleRequest() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // Return the vectors out of order to prove the index is honored. + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(9, embedding(1, 0.9f), embedding(0, 0.1f))); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + EmbeddingResult<List<float[]>> result = + conn.embedWithUsage( + List.of("first", "second"), Map.of("model", "text-embedding-3-small")); + + assertThat(result.getEmbeddings()).hasSize(2); + assertThat(result.getEmbeddings().get(0)).containsExactly(0.1f); + assertThat(result.getEmbeddings().get(1)).containsExactly(0.9f); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(9L); + verify(service).create(any(EmbeddingCreateParams.class)); + } + + @Test + @DisplayName("Setup parameters are forwarded to the request, extras as body properties") + void testRequestParametersForwarded() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(1, embedding(0, 1f))); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + conn.embed( + "hello", + Map.of( + "model", + "text-embedding-3-large", + "encoding_format", + "base64", + "dimensions", + 512, + "user", + "u1", + "timeout", + 5000, + "additional_kwargs", + Map.of("custom_flag", true))); + + ArgumentCaptor<EmbeddingCreateParams> captor = + ArgumentCaptor.forClass(EmbeddingCreateParams.class); + verify(service).create(captor.capture()); + EmbeddingCreateParams params = captor.getValue(); + assertThat(params.model().toString()).isEqualTo("text-embedding-3-large"); + assertThat(params.encodingFormat()).contains(EmbeddingCreateParams.EncodingFormat.BASE64); + assertThat(params.dimensions()).contains(512L); + assertThat(params.user()).contains("u1"); + assertThat(params._additionalBodyProperties()) + .containsKey("custom_flag") + .doesNotContainKey("timeout"); + } + + @Test + @DisplayName("A mismatched number of embeddings fails loudly") + void testMismatchedResponseFails() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(1, embedding(0, 1f))); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("2 input texts"); + } + + @Test + @DisplayName("Missing model parameter is rejected before any request") + void testMissingModel() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + assertThatThrownBy(() -> conn.embed("hello", Map.of())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model"); + } + + @Test + @DisplayName("Setup validates argument types and requires a model") + void testSetupValidation() { + ResourceDescriptor noModel = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(noModel, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model"); + ResourceDescriptor stringDims = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("dimensions", "256") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(stringDims, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dimensions"); + ResourceDescriptor badExtra = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("additional_kwargs", "not-a-map") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(badExtra, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("additional_kwargs"); + ResourceDescriptor unserializable = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("additional_kwargs", Map.of("opts", new Object())) + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(unserializable, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("cannot be sent as JSON"); + } + + @Test + @DisplayName("A response without usage yields embeddings and null token usage") + void testMissingUsageIsTolerated() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // The SDK builder refuses a response without usage, so mock the parsed response instead. + CreateEmbeddingResponse noUsage = mock(CreateEmbeddingResponse.class); + when(noUsage._data()).thenReturn(JsonField.of(List.of(embedding(0, 0.5f)))); + @SuppressWarnings("unchecked") + JsonField<CreateEmbeddingResponse.Usage> missing = + (JsonField<CreateEmbeddingResponse.Usage>) (JsonField<?>) JsonMissing.of(); + when(noUsage._usage()).thenReturn(missing); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(noUsage); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + EmbeddingResult<float[]> result = + conn.embedWithUsage("hello", Map.of("model", "text-embedding-3-small")); + assertThat(result.getEmbeddings()).containsExactly(0.5f); + assertThat(result.getTokenUsage()).isNull(); + } + + @Test + @DisplayName("A null input text is rejected with its index before any request") + void testNullTextRejected() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + List<String> texts = new java.util.ArrayList<>(); + texts.add("ok"); + texts.add(null); + assertThatThrownBy(() -> conn.embed(texts, Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index 1"); + } + + @Test + @DisplayName("Partial usage and missing indices from compatible servers are tolerated") + void testPartialUsageAndMissingIndex() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + @SuppressWarnings("unchecked") + JsonField<Long> missingLong = (JsonField<Long>) (JsonField<?>) JsonMissing.of(); + CreateEmbeddingResponse.Usage usage = mock(CreateEmbeddingResponse.Usage.class); + when(usage._promptTokens()).thenReturn(missingLong); + when(usage._totalTokens()).thenReturn(JsonField.of(7L)); + Embedding first = mock(Embedding.class); + when(first._index()).thenReturn(missingLong); + when(first.embeddingValue()).thenReturn(EmbeddingValue.ofFloats(List.of(0.1f))); + Embedding second = mock(Embedding.class); + when(second._index()).thenReturn(missingLong); + when(second.embeddingValue()).thenReturn(EmbeddingValue.ofFloats(List.of(0.2f))); + CreateEmbeddingResponse response = mock(CreateEmbeddingResponse.class); + when(response._data()).thenReturn(JsonField.of(List.of(first, second))); + when(response._usage()).thenReturn(JsonField.of(usage)); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + EmbeddingResult<List<float[]>> result = + conn.embedWithUsage(List.of("a", "b"), Map.of("model", "text-embedding-3-small")); + + assertThat(result.getEmbeddings().get(0)).containsExactly(0.1f); + assertThat(result.getEmbeddings().get(1)).containsExactly(0.2f); + assertThat(result.getTokenUsage()).isNotNull(); + // Missing prompt_tokens equals total_tokens for embedding requests, not zero. + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(7L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(7L); + } + + @Test + @DisplayName("additional_kwargs may not repeat the typed request fields") + void testReservedAdditionalKwargsRejected() { + ResourceDescriptor setup = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument( + "additional_kwargs", Map.of("model", "other", "custom", 1)) + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(setup, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("additional_kwargs") + .hasMessageContaining("[model]"); + + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + assertThatThrownBy( + () -> + conn.embed( + "hello", + Map.of( + "model", + "text-embedding-3-small", + "additional_kwargs", + Map.of("input", "x", "dimensions", 1)))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("[dimensions, input]"); + verify(service, never()).create(any(EmbeddingCreateParams.class)); + } + + @Test + @DisplayName("encoding_format accepts only float and base64") + void testEncodingFormatValidation() { + ResourceDescriptor setup = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", "text-embedding-3-small") + .addInitialArgument("encoding_format", "hex") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(setup, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("encoding_format") + .hasMessageContaining("hex"); + assertThat(OpenAIEmbeddingModelConnection.parseEncodingFormat("base64")) + .isEqualTo("base64"); + assertThat(OpenAIEmbeddingModelConnection.parseEncodingFormat(null)).isEqualTo("float"); + + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + assertThatThrownBy( + () -> + conn.embed( + "hello", + Map.of( + "model", + "text-embedding-3-small", + "encoding_format", + "FLOAT"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("encoding_format"); + verify(service, never()).create(any(EmbeddingCreateParams.class)); + } + + @Test + @DisplayName("Non-string or missing connection/model fail with named errors, not a CCE") + void testSetupConnectionAndModelValidation() { + ResourceDescriptor intModel = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", "conn") + .addInitialArgument("model", 42) + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(intModel, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model must be a string"); + ResourceDescriptor intConnection = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("connection", 42) + .addInitialArgument("model", "text-embedding-3-small") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(intConnection, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("connection must be a string"); + ResourceDescriptor noConnection = + ResourceDescriptor.Builder.newBuilder(OpenAIEmbeddingModelSetup.class.getName()) + .addInitialArgument("model", "text-embedding-3-small") + .build(); + assertThatThrownBy(() -> new OpenAIEmbeddingModelSetup(noConnection, NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("non-empty 'connection'"); + } + + private static Embedding indexlessEmbedding(float value) { + @SuppressWarnings("unchecked") + JsonField<Long> missingLong = (JsonField<Long>) (JsonField<?>) JsonMissing.of(); + Embedding embedding = mock(Embedding.class); + when(embedding._index()).thenReturn(missingLong); + when(embedding.embeddingValue()).thenReturn(EmbeddingValue.ofFloats(List.of(value))); + return embedding; + } + + private static Embedding indexedMock(long index, float value) { + Embedding embedding = mock(Embedding.class); + when(embedding._index()).thenReturn(JsonField.of(index)); + when(embedding.embeddingValue()).thenReturn(EmbeddingValue.ofFloats(List.of(value))); + return embedding; + } + + @SuppressWarnings("unchecked") + private static CreateEmbeddingResponse responseWithoutUsage(Embedding... data) { + CreateEmbeddingResponse response = mock(CreateEmbeddingResponse.class); + when(response._data()).thenReturn(JsonField.of(List.of(data))); + when(response._usage()) + .thenReturn( + (JsonField<CreateEmbeddingResponse.Usage>) (JsonField<?>) JsonMissing.of()); + return response; + } + + @Test + @DisplayName("A partially indexed response is accepted only if its indices match positions") + void testPartiallyIndexedResponse() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + CreateEmbeddingResponse response = + responseWithoutUsage(indexlessEmbedding(0.1f), indexedMock(1, 0.2f)); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + List<float[]> embeddings = + conn.embed(List.of("a", "b"), Map.of("model", "text-embedding-3-small")); + + assertThat(embeddings.get(0)).containsExactly(0.1f); + assertThat(embeddings.get(1)).containsExactly(0.2f); + } + + @Test + @DisplayName("A partially indexed response that contradicts response order is rejected") + void testContradictoryPartialIndexRejected() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // Vector 0.2 claims index 0 while sitting at position 1: neither order can be trusted. + CreateEmbeddingResponse response = + responseWithoutUsage(indexlessEmbedding(0.1f), indexedMock(0, 0.2f)); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("unexpected embedding index 0 at position 1"); + } + + @Test + @DisplayName("An empty batch is validated but sends no request") + void testEmptyBatchValidatesWithoutRequest() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThat(conn.embed(List.of(), Map.of("model", "text-embedding-3-small"))).isEmpty(); + assertThatThrownBy( + () -> + conn.embed( + (List<String>) null, + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("texts must not be null"); + // A null parameter map is treated as empty, so the missing model is the reported error. + assertThatThrownBy(() -> conn.embed("hello", null)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("model"); + assertThatThrownBy( + () -> + conn.embed( + List.of(), + Map.of( + "model", + "text-embedding-3-small", + "encoding_format", + "bogus"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("encoding_format"); + verify(service, never()).create(any(EmbeddingCreateParams.class)); + } + + @Test + @DisplayName("A base64 response is decoded to floats") + void testBase64ResponseDecoded() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // 0.5f and 0.25f as little-endian float32, base64 encoded (the OpenAI wire format). + String base64 = + Base64.getEncoder() + .encodeToString( + ByteBuffer.allocate(8) + .order(ByteOrder.LITTLE_ENDIAN) + .putFloat(0.5f) + .putFloat(0.25f) + .array()); + Embedding encoded = + Embedding.builder() + .index(0L) + .object_(JsonValue.from("embedding")) + .embedding(base64) + .build(); + CreateEmbeddingResponse response = response(3, encoded); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + float[] vector = + conn.embed( + "hello", + Map.of("model", "text-embedding-3-small", "encoding_format", "base64")); + + assertThat(vector).containsExactly(0.5f, 0.25f); + } + + @Test + @DisplayName( + "A base64 vector whose length is not a multiple of four is rejected, not truncated") + void testTruncatedBase64Rejected() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // 6 bytes: one float plus two trailing bytes; the SDK would silently return [1.0]. + Embedding truncated = + Embedding.builder() + .index(0L) + .object_(JsonValue.from("embedding")) + .embedding("AACAPwAA") + .build(); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response(3, truncated)); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + "hello", + Map.of( + "model", + "text-embedding-3-small", + "encoding_format", + "base64"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("malformed embedding at position 0") + .hasRootCauseMessage("base64 embedding has 6 bytes, not a multiple of 4."); + } + + @Test + @DisplayName("Decimal-string token counts from compatible servers are accepted") + void testStringTokenCounts() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + @SuppressWarnings("unchecked") + JsonField<Long> five = (JsonField<Long>) (JsonField<?>) JsonValue.from("5"); + @SuppressWarnings("unchecked") + JsonField<Long> junk = (JsonField<Long>) (JsonField<?>) JsonValue.from("n/a"); + CreateEmbeddingResponse.Usage usage = mock(CreateEmbeddingResponse.Usage.class); + when(usage._promptTokens()).thenReturn(five); + when(usage._totalTokens()).thenReturn(junk); + CreateEmbeddingResponse response = mock(CreateEmbeddingResponse.class); + when(response._data()).thenReturn(JsonField.of(List.of(embedding(0, 0.1f)))); + when(response._usage()).thenReturn(JsonField.of(usage)); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + EmbeddingResult<float[]> result = + conn.embedWithUsage("a", Map.of("model", "text-embedding-3-small")); + + assertThat(result.getTokenUsage()).isNotNull(); + assertThat(result.getTokenUsage().getPromptTokens()).isEqualTo(5L); + assertThat(result.getTokenUsage().getTotalTokens()).isEqualTo(5L); + } + + @Test + @DisplayName("A response without data is reported as a short response") + void testMissingDataReported() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + @SuppressWarnings("unchecked") + JsonField<List<Embedding>> missingData = + (JsonField<List<Embedding>>) (JsonField<?>) JsonMissing.of(); + CreateEmbeddingResponse response = mock(CreateEmbeddingResponse.class); + when(response._data()).thenReturn(missingData); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("returned 0 embeddings for 2 input texts"); + + // A data field that is present but not a list of embeddings names the payload. + @SuppressWarnings("unchecked") + JsonField<List<Embedding>> objectData = + (JsonField<List<Embedding>>) (JsonField<?>) JsonValue.from(Map.of("0", 1)); + CreateEmbeddingResponse badData = mock(CreateEmbeddingResponse.class); + when(badData._data()).thenReturn(objectData); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(badData); + assertThatThrownBy(() -> conn.embed("a", Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("malformed data field"); + + // A string vector when float was requested is malformed, not silently decoded. + Embedding stringVector = + Embedding.builder() + .index(0L) + .object_(JsonValue.from("embedding")) + .embedding("AAAAAAAAAAA=") + .build(); + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(3, stringVector)); + assertThatThrownBy(() -> conn.embed("a", Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("malformed embedding at position 0") + .hasRootCauseMessage( + "a string vector was returned although encoding_format was float."); + + // A null element in data is a malformed embedding at its position. + CreateEmbeddingResponse nullItem = mock(CreateEmbeddingResponse.class); + when(nullItem._data()).thenReturn(JsonField.of(java.util.Arrays.asList((Embedding) null))); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(nullItem); + assertThatThrownBy(() -> conn.embed("a", Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("malformed embedding at position 0"); + + OpenAIClient client2 = mock(OpenAIClient.class); + EmbeddingService service2 = mock(EmbeddingService.class); + OpenAIEmbeddingModelConnection conn2 = connectionWith(client2, service2); + // A per-call additional_kwargs value the SDK cannot serialize names its key. + assertThatThrownBy( + () -> + conn2.embed( + "hello", + Map.of( + "model", + "text-embedding-3-small", + "additional_kwargs", + Map.of("meta", new Object())))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("value for key 'meta' cannot be sent as JSON"); + } + + @Test + @DisplayName("A null or undecodable vector is reported with its position") + void testMalformedVectorReported() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + Embedding broken = mock(Embedding.class); + when(broken._index()).thenReturn(JsonField.of(1L)); + when(broken.embeddingValue()) + .thenThrow(new IllegalStateException("Invalid EmbeddingValue")); + CreateEmbeddingResponse response = responseWithoutUsage(indexedMock(0, 0.1f), broken); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("malformed embedding at position 1") + .hasCauseInstanceOf(IllegalStateException.class); + } + + @Test + @DisplayName("A non-integer index is a malformed response, not an absent index") + void testNonIntegerIndexRejected() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + // A JSON string where a long is expected: present, not null, but not a known Long. + @SuppressWarnings("unchecked") + JsonField<Long> stringIndex = (JsonField<Long>) (JsonField<?>) JsonValue.from("1"); + Embedding malformed = mock(Embedding.class); + when(malformed._index()).thenReturn(stringIndex); + when(malformed.embeddingValue()).thenReturn(EmbeddingValue.ofFloats(List.of(0.9f))); + CreateEmbeddingResponse response = responseWithoutUsage(malformed, indexedMock(0, 0.1f)); + when(service.create(any(EmbeddingCreateParams.class))).thenReturn(response); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("non-integer embedding index"); + } + + @Test + @DisplayName("Indices outside the batch fail with a clear error") + void testOutOfRangeIndexFails() { + OpenAIClient client = mock(OpenAIClient.class); + EmbeddingService service = mock(EmbeddingService.class); + when(service.create(any(EmbeddingCreateParams.class))) + .thenReturn(response(3, embedding(0, 0.1f), embedding(Long.MAX_VALUE, 0.2f))); + OpenAIEmbeddingModelConnection conn = connectionWith(client, service); + + assertThatThrownBy( + () -> + conn.embed( + List.of("a", "b"), + Map.of("model", "text-embedding-3-small"))) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("unexpected embedding index " + Long.MAX_VALUE); + } +} diff --git a/integrations/embedding-models/pom.xml b/integrations/embedding-models/pom.xml index 2a585d17..19325c07 100644 --- a/integrations/embedding-models/pom.xml +++ b/integrations/embedding-models/pom.xml @@ -33,6 +33,7 @@ under the License. <modules> <module>bedrock</module> <module>ollama</module> + <module>openai</module> </modules> </project> diff --git a/python/flink_agents/api/resource.py b/python/flink_agents/api/resource.py index 55c6ee12..8cac7112 100644 --- a/python/flink_agents/api/resource.py +++ b/python/flink_agents/api/resource.py @@ -366,6 +366,10 @@ class ResourceName: BEDROCK_CONNECTION = "org.apache.flink.agents.integrations.embeddingmodels.bedrock.BedrockEmbeddingModelConnection" BEDROCK_SETUP = "org.apache.flink.agents.integrations.embeddingmodels.bedrock.BedrockEmbeddingModelSetup" + # OpenAI + OPENAI_CONNECTION = "org.apache.flink.agents.integrations.embeddingmodels.openai.OpenAIEmbeddingModelConnection" + OPENAI_SETUP = "org.apache.flink.agents.integrations.embeddingmodels.openai.OpenAIEmbeddingModelSetup" + class VectorStore: """VectorStore resource names.""" diff --git a/python/flink_agents/api/yaml/aliases.py b/python/flink_agents/api/yaml/aliases.py index 584ba2fd..76b59ddd 100644 --- a/python/flink_agents/api/yaml/aliases.py +++ b/python/flink_agents/api/yaml/aliases.py @@ -112,6 +112,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]] = { "java": { "ollama": ResourceName.EmbeddingModel.Java.OLLAMA_CONNECTION, "bedrock": ResourceName.EmbeddingModel.Java.BEDROCK_CONNECTION, + "openai": ResourceName.EmbeddingModel.Java.OPENAI_CONNECTION, }, }, ResourceType.EMBEDDING_MODEL: { @@ -123,6 +124,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]] = { "java": { "ollama": ResourceName.EmbeddingModel.Java.OLLAMA_SETUP, "bedrock": ResourceName.EmbeddingModel.Java.BEDROCK_SETUP, + "openai": ResourceName.EmbeddingModel.Java.OPENAI_SETUP, }, }, ResourceType.VECTOR_STORE: { diff --git a/python/flink_agents/api/yaml/tests/test_aliases.py b/python/flink_agents/api/yaml/tests/test_aliases.py index d7ef0bc0..ecac917f 100644 --- a/python/flink_agents/api/yaml/tests/test_aliases.py +++ b/python/flink_agents/api/yaml/tests/test_aliases.py @@ -199,6 +199,23 @@ def test_resolve_clazz_covers_vector_store_java_and_python() -> None: ) +def test_resolve_clazz_covers_embedding_model_openai_in_both_languages() -> None: + assert resolve_clazz("openai", ResourceType.EMBEDDING_MODEL_CONNECTION).endswith( + "openai_embedding_model.OpenAIEmbeddingModelConnection" + ) + assert resolve_clazz("openai", ResourceType.EMBEDDING_MODEL).endswith( + "openai_embedding_model.OpenAIEmbeddingModelSetup" + ) + # The Java and Python simple names collide (OpenAIEmbeddingModelConnection); pin the + # package so a Java entry that resolved to the Python class would fail here. + java_conn = resolve_clazz("openai", ResourceType.EMBEDDING_MODEL_CONNECTION, "java") + assert java_conn.startswith("org.apache.flink.agents") + assert java_conn.endswith("OpenAIEmbeddingModelConnection") + java_setup = resolve_clazz("openai", ResourceType.EMBEDDING_MODEL, "java") + assert java_setup.startswith("org.apache.flink.agents") + assert java_setup.endswith("OpenAIEmbeddingModelSetup") + + def test_java_wrapper_clazz_table_covers_supported_types() -> None: # The Python-side wrappers must exist for every cross-language type expected = {
