purushah opened a new pull request, #1126: URL: https://github.com/apache/flink-agents/pull/1126
Closes #1103. Adds a Java OpenAI embedding model integration, closing the Java/Python parity gap for embeddings: Python already had `OpenAIEmbeddingModelConnection`/`OpenAIEmbeddingModelSetup`; Java had only Ollama and Bedrock, so a Java agent needed the Python wrapper (and a Python runtime on the TaskManagers) for a plain HTTPS call. ## What's in the change - New module `integrations/embedding-models/openai` with `OpenAIEmbeddingModelConnection` and `OpenAIEmbeddingModelSetup`, built on `com.openai:openai-java` (already a dependency of the OpenAI chat module; no new third-party dependency). - Arguments mirror the Python classes one for one. Connection: `api_key` (required), `base_url`, `request_timeout` (seconds, `0` disables), `max_retries`, `organization`, `project`. Setup: `connection`, `model`, `encoding_format`, `dimensions`, `user`, `additional_kwargs` (forwarded as extra body properties, overriding the named ones like the Python `model_kwargs`). - A batch of texts is sent as one request; results are placed by the response `index` so input order is preserved even if the API reorders. `embedWithUsage` reports `prompt_tokens`/`total_tokens`, so the embedding token metrics from #870 cover this provider. - Registration: `ResourceName.EmbeddingModel.OPENAI_CONNECTION/OPENAI_SETUP`, the Python `ResourceName.EmbeddingModel.Java` mirror, YAML aliases (`openai`), `dist`, `ide-support`, and the e2e integration module. - Docs: `embedding_models.md` OpenAI section gains Java tabs (usage, connection and setup parameters); the Python-only hint is removed. - Tests: `OpenAIEmbeddingModelTest` (11 tests, SDK client mocked, following `BedrockEmbeddingModelTest`): defaults, explicit options, missing `api_key`, timeout/retry parsing bounds, setup parameters, single and batch embeddings with token usage, out-of-order response handling, parameter forwarding (encoding format, dimensions, user, extra body properties), mismatched response size, missing model. `EmbeddingIntegrationTest` gains an `OPENAI` provider (gated on `OPENAI_API_KEY`, like the chat model integration test). Naming follows the Python embedding connection (`base_url`, `request_timeout`) rather than the Java chat connection (`api_base_url`, `timeout`), as proposed in the issue, so the two embedding implementations align across languages. ## Test commands and results Unit tests for the new module: ``` mvn -pl integrations/embedding-models/openai test # Tests run: 11, Failures: 0, Errors: 0, Skipped: 0 -- BUILD SUCCESS ``` Complete Java unit suite (all non-e2e modules, as `tools/ut.sh` runs it): ``` mvn test -pl '!e2e-test/flink-agents-end-to-end-tests-integration,!e2e-test/flink-agents-end-to-end-tests-resource-cross-language' # Tests run: 3564, Failures: 0, Errors: 0, Skipped: 0 -- BUILD SUCCESS ``` Complete Python unit suite (e2e tests excluded, as `tools/ut.sh` does) plus the cross-language resource-name consistency check: ``` cd python && uv sync --extra test && uv pip install 'apache-flink~=2.3.0' uv run --no-sync pytest flink_agents -k "not e2e_tests" # 1420 passed, 27 skipped, 82 deselected uv run --no-sync python ../e2e-test/test-scripts/check_resource_consistency.py # ResourceName Cross-language consistency check successful ``` End-to-end, real OpenAI API, embedded Flink 2.3 cluster (`OPENAI_API_KEY` exported from a local file, never printed): ``` mvn -pl e2e-test/flink-agents-end-to-end-tests-integration test -Dtest=EmbeddingIntegrationTest -Pflink-2.3 # Tests run: 2, Failures: 0, Errors: 0, Skipped: 0 -- BUILD SUCCESS # [OLLAMA] 10 embeddings, Dimension=768; [OPENAI] 10 embeddings, Dimension=256 (text-embedding-3-small, dimensions=256) ``` End-to-end, real OpenAI API, standalone local Flink 2.3.0 cluster (`start-cluster.sh`, `flink-agents-dist-flink-2.3` in `lib/`, throwaway job submitted with `flink run`, results written by a `FileSink`): ``` text=Apache Flink processes streams|dim=256|norm=1.000|batch=2|batchDim=256|promptTokens=5|totalTokens=13 text=Embeddings map text to vectors|dim=256|norm=1.000|batch=2|batchDim=256|promptTokens=6|totalTokens=15 text=Model routing picks a model per request|dim=256|norm=1.000|batch=2|batchDim=256|promptTokens=7|totalTokens=17 # job openai-embedding-e2e FINISHED, no exceptions in JobManager/TaskManager logs ``` The requested `dimensions` are honored, OpenAI's unit-normalized vectors come back with norm 1.000, batch calls return one vector per input, and token usage is reported for single and batch calls. Also run: `ruff check` and `ruff format --check` on the touched Python files, Hugo build of the docs, spotless on the touched modules. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
