PSMDB-2143 Implement OpenAI-compatible embedding provider - #17
Conversation
|
|
There was a problem hiding this comment.
Pull request overview
Implements an OpenAI-compatible embedding provider end-to-end, enabling auto-embedding against any server that speaks the OpenAI /v1/embeddings API (e.g., Ollama/vLLM/TEI/OpenAI/Azure OpenAI), including support for keyless local engines and an operator-editable on-disk model catalog.
Changes:
- Added
OPENAI_COMPATIBLEprovider support across config parsing, model config/credentials types, client factory wiring, and vector-param resolution. - Introduced
OpenAiCompatClientandOpenAiApiSchemato serialize requests and decode embeddings (base64 float32 LE and fallback JSON float arrays). - Shipped and tested a default OpenAI-compatible catalog (including on-disk override behavior) and updated bootstrap/config logic to keep OpenAI-compatible models usable even without Voyage credentials.
Reviewed changes
Copilot reviewed 25 out of 25 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/test/unit/java/com/xgen/mongot/embedding/providers/config/OpenAiApiSchemaTest.java | Adds unit coverage for OpenAI embeddings request/response BSON/JSON handling. |
| src/test/unit/java/com/xgen/mongot/embedding/providers/config/BUILD | Registers the new schema test in the suite. |
| src/test/unit/java/com/xgen/mongot/embedding/providers/clients/OpenAiCompatClientTest.java | Adds extensive unit coverage for auth headers, dimensions forwarding, prefixes, retry/error semantics, and HttpClient renewal behavior. |
| src/test/unit/java/com/xgen/mongot/embedding/providers/clients/EmbeddingClientFactoryTest.java | Verifies factory wiring for OPENAI_COMPATIBLE across tiers and flex-tier isolation. |
| src/test/unit/java/com/xgen/mongot/embedding/providers/clients/BUILD | Registers the new OpenAI-compatible client test in the suite. |
| src/test/unit/java/com/xgen/mongot/config/provider/community/embedding/EmbeddingServiceManagerConfigTest.java | Expands tests for provider parsing, keyless behavior, and on-disk catalog overrides/fallbacks. |
| src/test/unit/java/com/xgen/mongot/config/provider/community/CommunityMongotBootstrapperTest.java | Ensures global endpoint override applies to Voyage only and does not leak into OpenAI-compatible models. |
| src/test/unit/java/com/xgen/mongot/config/provider/community/CommunityConfigTest.java | Updates config deserialization tests for the new modelConfigFile field. |
| src/main/resources/config/community/embedding-service-configs.yml | Adds OpenAI-compatible model entries (bge-m3, nomic-embed-text) and Azure template documentation. |
| src/main/resources/config/community/BUILD | Exposes the catalog filegroup to //deploy:__pkg__ for packaging. |
| src/main/java/com/xgen/mongot/index/definition/VectorAutoEmbedFieldSpecification.java | Resolves vector params via the polymorphic ModelConfig interface (not Voyage-only) and uses “configured” getters. |
| src/main/java/com/xgen/mongot/embedding/providers/configs/OpenAiApiSchema.java | Implements OpenAI embeddings wire schema with base64 float decoding + JSON float-array tolerance. |
| src/main/java/com/xgen/mongot/embedding/providers/configs/EmbeddingServiceConfig.java | Adds OPENAI_COMPATIBLE, introduces OpenAiModelConfig/OpenAiEmbeddingCredentials, and extends ModelConfig with “configured” accessors. |
| src/main/java/com/xgen/mongot/embedding/providers/configs/EmbeddingModelConfig.java | Adds consolidation logic for OpenAI model configs and supports OpenAI credentials overrides. |
| src/main/java/com/xgen/mongot/embedding/providers/configs/EmbeddingConfigFactory.java | Extends polymorphic parsing to OpenAI model configs and credentials. |
| src/main/java/com/xgen/mongot/embedding/providers/configs/BUILD.bazel | Includes OpenAiApiSchema.java in the library. |
| src/main/java/com/xgen/mongot/embedding/providers/clients/OpenAiCompatClient.java | Adds the OpenAI-compatible HTTP client implementation (auth, retry classification, decoding, HttpClient refresh/renewal). |
| src/main/java/com/xgen/mongot/embedding/providers/clients/EmbeddingClientFactory.java | Wires OPENAI_COMPATIBLE to build OpenAiCompatClient. |
| src/main/java/com/xgen/mongot/embedding/providers/clients/BUILD | Includes OpenAiCompatClient.java in the clients library. |
| src/main/java/com/xgen/mongot/config/provider/community/embedding/EmbeddingServiceManagerConfig.java | Adds on-disk catalog override support and keyless model loading behavior; injects provider discriminators/credentials. |
| src/main/java/com/xgen/mongot/config/provider/community/embedding/EmbeddingConfig.java | Adds modelConfigFile to community embedding config schema. |
| src/main/java/com/xgen/mongot/config/provider/community/CommunityMongotBootstrapper.java | Enables auto-embedding without Voyage keys, adds “empty manager” helper, and scopes endpoint override to Voyage only. |
| deploy/community-resources/mongot | Sets -Dmongot.embeddingModelConfigFile to point at the on-disk shipped catalog. |
| deploy/community-resources/config.default.yml | Documents Voyage-only endpoint override and the new on-disk catalog override option. |
| deploy/BUILD | Packages the default embedding model catalog into the deploy tar. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
63f6dde to
d7da0b9
Compare
…schema) Generic embedding client for any server speaking the OpenAI /v1/embeddings protocol (OpenAI, Azure OpenAI, Ollama, vLLM, llama.cpp, LM Studio, LocalAI, HF TEI). Optional API key (keyless local engines), Authorization: Bearer or Azure api-key auth header, fail-fast on HTTP 401/403, tier-aware query/document input prefixes, opt-in dimensions forwarding resolved from the request context, float-only day 1 (quantization requests fail fast). Replaces the hard VoyageModelConfig cast in VectorAutoEmbedFieldSpecification with polymorphic ModelConfig accessors so non-Voyage providers resolve numDimensions, quantization, and similarity. Co-authored-by: Oleksandr Miroshnychenko <alex.miroshnychenko@percona.com>
Auto-embedding now enables whenever an embedding: config section is present: Voyage API keys are optional, VOYAGE catalog entries are dropped with a warning when no keys are configured, and keyless OPENAI_COMPATIBLE models are kept. The model catalog can be overridden on disk via embedding.modelConfigFile or the mongot.embeddingModelConfigFile system property, falling back to the bundled resource. The global providerEndpoint override applies to VOYAGE models only; OPENAI_COMPATIBLE models keep their per-model catalog endpoints. Bundled catalog gains two local Ollama models (bge-m3, nomic-embed-text) plus a commented Azure OpenAI template. Co-authored-by: Oleksandr Miroshnychenko <alex.miroshnychenko@percona.com>
Ship embedding-service-configs.yml next to the JAR in the community tarball and point the launcher at it via -Dmongot.embeddingModelConfigFile so operators can edit the model catalog and restart without rebuilding. Document modelConfigFile and the VOYAGE-only scope of providerEndpoint in config.default.yml. Co-authored-by: Oleksandr Miroshnychenko <alex.miroshnychenko@percona.com>
Operators can put modelConfig/credentials under query, collectionScan, or changeStream in an on-disk catalog without the internal discriminator; inject those the same way as the base fields so parsing succeeds.
Avoid silently falling back to the bundled catalog (localhost endpoints) when an operator-set modelConfigFile is missing or invalid; keep soft fallback only for the launcher system-property path, and assert ERROR logs.
ktrushin
left a comment
ktrushin
left a comment
There was a problem hiding this comment.
I looked through the code and quickly realized it would be naive to think I could find issues in PR that:
- has already passed a couple of AI-assisted reviews;
- is made to the codebase I've never seen before;
- is in a language I used 20 years ago last time.
I tried Claude Code to see whether it was able to find something and it was. Please see the attached file. I guess majority of the "failure scenarios" it suggests are purely theoretical, though. So take them with a grain of salt.
Sorry for resorting to an AI tool when I was supposed to do review exercising my own expertise. But due to the lack of the latter, I think an AI report is better than blind rubber-stamping the PR.
The embedding catalog file is not loaded only when config file has no |
… array order extractVectorsFromResponse zipped embedResponse.data to inputs purely by iteration order, ignoring each item's index field. Unlike the single first-party Voyage API, OPENAI_COMPATIBLE targets many heterogeneous third-party servers; a backend that batches or parallelizes and returns data[] out of input order would silently attach the wrong vector to the wrong document or query text with no error raised. Build a map keyed by the response's index field and look up each non-empty input by its position in the request instead of trusting array order. Adds a regression test with an out-of-order response.
Fixed in separate commit: extractVectorsFromResponse zipped embedResponse.data to inputs purely by Build a map keyed by the response's index field and look up each |
redactApiKey() was applied only to the buildRequest IllegalArgumentException path; the 400/422, 429, 408, and 401/403 branches embedded the raw response.body() in a logged message or exception, bypassing the redaction. A misconfigured or malicious OpenAI-compatible endpoint/proxy that echoes request headers back in an error body would leak the API key into mongot's own logs. Redact response.body() at each of these call sites before it is logged or thrown. Adds a regression test that echoes the key back across all four fixed status codes.
PSMDB-2143 Redact API key from all OpenAiCompatClient error-body paths redactApiKey() was applied only to the buildRequest IllegalArgumentException Redact response.body() at each of these call sites before it is logged or |
endpoint, apiKey, authHeaderName, forwardDimensions, and inputPrefix were plain fields written by updateConfig() with five separate unsynchronized writes, while embed()/buildRequest() read them concurrently from request-handling threads. EmbeddingProviderManager.updateEmbeddingProviderManager can call updateConfig() on the same client instance that concurrently serves embed() calls, so a config reload (e.g. key rotation) had no happens-before relationship with in-flight requests: a reader could observe indefinitely stale values, or a torn combination such as a new endpoint with an old key. Consolidate the five fields into an immutable RequestConfig record held behind a single volatile reference, swapped atomically in updateConfig() and the constructor, and snapshotted once per embed() call. This mirrors the existing volatile HttpClient swap pattern already used in this class.
PSMDB-2143 Swap OpenAiCompatClient request config atomically endpoint, apiKey, authHeaderName, forwardDimensions, and inputPrefix were Consolidate the five fields into an immutable RequestConfig record held |
Status codes outside the explicitly handled set (400/422/429/408/401/403) fell into the generic non-2xx branch and were thrown as EmbeddingProviderTransientException, including permanent client errors like 404 (wrong endpoint path), 405, or 410. A misconfigured providerEndpoint -- a plausible mistake given this provider targets arbitrary self-hosted servers -- would burn the full retry budget on a request that could never succeed. Split the fallback: any other 4xx is a permanent misconfiguration and fails fast as EmbeddingProviderNonTransientException, matching the fail-fast treatment already given to 401/403; 5xx and anything else stays transient.
PSMDB-2143 Fail fast on permanent 4xx errors from OpenAiCompatClient Status codes outside the explicitly handled set (400/422/429/408/401/403) Split the fallback: any other 4xx is a permanent misconfiguration and fails |
I verified the mechanism: injectCredentials computes one provider from the base entry (providerOf(configDoc)) and passes that same value into tagWorkloadProviderFields for all three workload keys (query/collectionScan/changeStream) unconditionally — so a workload override block can never end up tagged with a different provider than its parent entry. That part of the review is accurate. But I don't think this is a real gap, and for the same reason as comment #1: I traced this up through EmbeddingModelConfig, which is the type these tagged docs ultimately get consolidated into. It's a record with one provider field for the entire model (EmbeddingModelConfig(String name, EmbeddingProvider provider, ...)), and create() takes a single provider argument used uniformly to build all three ConsolidatedWorkloadParams (query/changeStream/collectionScan). There's no per-tier provider parameter anywhere in that API, and EmbeddingClientFactory/EmbeddingConfigFactory both dispatch on one provider per model, not per workload. So "provider" is a whole-model property everywhere in this framework — it isn't something this PR chose to leave unsupported per-workload; it's a pre-existing, system-wide architectural invariant. Given that, the review's failure scenario — "an operator writes a workload override meant to use a different provider's credential shape" — describes an operation nothing in the schema, docs, or downstream consolidation logic supports or expects in the first place. Even if tagWorkloadProviderFields were changed to honor a per-workload embeddingProvider override, EmbeddingModelConfig.create() would still process the whole model under the one top-level provider, so the fix wouldn't actually enable anything — it would just move the inevitable failure to a different, equally-opaque spot. The right fix for "an operator tries something the system doesn't support" would be validation/error messaging, which is a different, broader concern than what's described here, not a defect in this diff. |
renewHttpClientAfterConnectionFailure only deduped concurrent failures on the same client instance; it had no cooldown across sequential failures over time. During a sustained outage (a down local engine, which this provider specifically targets), every failed request renewed the client again, each renewal spawning a dedicated OS thread purely to shut down the just-replaced, also-failed client -- unthrottled churn exactly when the system is already degraded. Track the epoch of the last connection-failure renewal separately from the existing periodic-refresh timestamp, and skip renewing again within a cooldown window. The first failure during an outage still renews immediately (initial value 0), so existing behavior is unchanged there.
decodeVectorFromBsonValue's byte-length check only rejected lengths not a multiple of 4; a zero-length base64 string passed the modulo check (0 % 4 == 0) and silently produced a valid-looking zero-dimension Vector instead of a parse error. A misbehaving OpenAI-compatible server response with an empty embedding string would decode successfully rather than failing fast with a clear diagnostic pointing at the actual cause. Reject a zero-length decode explicitly before the modulo check, using the same handleSemanticError idiom already used for the adjacent checks.
PSMDB-2143 Add cooldown to OpenAiCompatClient connection-failure renewal renewHttpClientAfterConnectionFailure only deduped concurrent failures on Track the epoch of the last connection-failure renewal separately from the |
PSMDB-2143 Reject empty base64 embeddings in OpenAiApiSchema decodeVectorFromBsonValue's byte-length check only rejected lengths not a Reject a zero-length decode explicitly before the modulo check, using the |
…st timeout HTTP_CLIENT_SHUTDOWN_AWAIT was 5s while DEFAULT_TIMEOUT (the per-request budget) is 60s. renewHttpClientIfStale swaps the shared HttpClient every 10 minutes unconditionally, regardless of requests still in flight on the old instance; a legitimately slow-but-successful request (well within its own 60s budget) still in flight on the just-replaced client got force-cancelled by shutdownNow() after only 5 more seconds, turning a request that would have succeeded into a spurious connection failure. Unlike the other fixes in this PR, this could happen during completely normal operation, not just misconfiguration or an outage. Raise HTTP_CLIENT_SHUTDOWN_AWAIT to match DEFAULT_TIMEOUT. Any in-flight request will either succeed or hit its own request-level timeout within that window, so graceful shutdown naturally completes on its own once there's nothing left in flight -- shutdownNow() remains as a last-resort fallback for the pathological case where a request ignores its own timeout. The new client is always swapped in synchronously before the old one's shutdown even starts, so this has no effect on new-request latency.
PSMDB-2143 Give OpenAiCompatClient shutdown grace time to match request timeout HTTP_CLIENT_SHUTDOWN_AWAIT was 5s while DEFAULT_TIMEOUT (the per-request Raise HTTP_CLIENT_SHUTDOWN_AWAIT to match DEFAULT_TIMEOUT. Any in-flight |
| IllegalArgumentException cleanedException = | ||
| new IllegalArgumentException(cleanedMessage, e.getCause()); |
There was a problem hiding this comment.
- e's own message may contain the raw header value (e.g. the API key), so it must not be attached as a cause
- e.getCause() is always null here (HttpRequest.Builder throws IllegalArgumentException directly, never wrapping another exception)
- fixed by copying only the stack trace to preserve where the failure occurred
| redactApiKey(response.body(), apiKey)); | ||
| LOG.warn(errorMessage); | ||
| this.invalidRequestCounter.increment(); | ||
| return inputs.stream().map(ignored -> new VectorOrError(errorMessage)).toList(); |
| if (!configField.containsKey("credentials")) { | ||
| BsonDocument credentialsDoc = credentials.indexingCredentials.toBson(); | ||
| credentialsDoc.put("_provider", new BsonString("VOYAGE")); | ||
| configField.put("credentials", credentialsDoc); | ||
| } |
There was a problem hiding this comment.
Out of scope for PSMDB-2143.
- It only affects VOYAGE catalog entries. This ticket's deliverable is the OPENAI_COMPATIBLE provider, which is credential-optional and never passes through this method.
- The failure mode is not silent: a malformed credentials: {} block produces a caught, logged BsonParseException naming the offending file (loadFromFile's broad catch at line ~152), degrading to either "auto-embedding inactive" (explicit override) or a fallback to the bundled catalog (launcher system property) — not a crash or silent data issue.
- Fixing pre-existing VOYAGE-specific robustness gaps is better scoped as its own follow-up than bundled into a PR adding a new, unrelated provider.
…atClient header errors When HttpRequest.Builder rejects an invalid header (e.g. an API key containing a stray control character), its IllegalArgumentException message embeds the raw header value verbatim. Attaching that exception as a cause would leak the unredacted key into logs via the cause chain, so instead copy only its stack trace onto the redacted exception to preserve where the failure occurred.
…esponses extractVectorsFromResponse wrapped every input in a fresh VectorOrError on a 400/422, including inputs that were empty strings filtered out before the request was sent. That diverged from the success path, which maps empty inputs to the EMPTY_INPUT_ERROR singleton, and broke EmbeddingIndexingWorkScheduler's identity check against that singleton that suppresses log noise for the known empty-input case.
No description provided.