BYOK: build an embedding service from a runtime-supplied key - #1892
BYOK: build an embedding service from a runtime-supplied key#1892jasperblues wants to merge 8 commits into
Conversation
BYOK modelled LLMs but not embeddings, so an embedding service could only come from ModelProvider, which resolves against models registered at boot. A key that arrived after startup could not produce one at all: a deployment collecting keys at first run had to restart before embeddings worked. Adds ByokFactory<EmbeddingService>, reached via openAiEmbedding() or byokEmbedding(). It probes the key with a single embed() and translates any provider error to InvalidApiKeyException, matching the LLM contract exactly. Two things deliberately differ from the LLM side, because an embedding model is a schema commitment rather than a stateless per-call service: No provider racing. detectProvider() is sound when any candidate accepting the key is an acceptable answer. Here the provider that answers first would decide the index dimension, so the model is always named explicitly. No caller-supplied dimension. The width is whatever the model returned during validation. Callers writing into an existing index compare it against that index themselves; the check belongs to whoever owns the index. Stamping the probed width also removes a hidden network call: Spring AI's AbstractEmbeddingModel.dimensions() caches, but resolves via dimensions(this, "Test", "Hello World") — the literal "Test" as model name, so its known-dimensions table can never hit and it always falls through to a live embed() at an arbitrary later moment. Scope is installation or world level, not per-user: per-user embedding keys would land vectors of different models and widths in one shared index. Closes #1883 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The LLM path treats a blank key as absent; the embedding path did not. Same deployment, same set-but-empty environment variable, two different failures — an actionable error for the chat key and an opaque provider error for the embedding key. The message is duplicated rather than shared. #1888 grows the same rule for LLMs and a shared helper is its right home, but reaching for that from here would couple two branches for four lines. Collapse them once both have landed; the constant says so. Signed-off-by: jasper blues <jasper@liberation-data.com>
cb5cb56 to
43a2be8
Compare
Any provider error during the probe became "invalid API key". On the LLM path the model is defaulted, so the key really is the likely cause. Here it is not: the embedding model is caller-supplied and mandatory, never defaulted and never detected, so a typo in it arrives as the same provider error — and the caller is sent to re-check a key that was fine. Signed-off-by: jasper blues <jasper@liberation-data.com>
igordayen
left a comment
igordayen
left a comment
There was a problem hiding this comment.
@jasperblues - could you please see whether some artifacts can be used for other than OpenAI providers. thank you
Both failure modes — the probe throws, or it comes back empty — mean the same thing to the caller, so they now share one exit instead of throwing from two places. The empty-vector case becomes an ordinary require inside the try, which is also how it picks up the model name in its message. Message constant switched from concatenation to a raw string, per the convention this repo follows everywhere else. Signed-off-by: jasper blues <jasper@liberation-data.com>
Review point, and correct: nothing about "build a service, embed one probe, reject an empty vector, rebuild stamping the width you observed" is OpenAI-specific. It now lives in com.embabel.common.byok as validatedEmbeddingService(model, provider, build), taking the provider's builder as a parameter. This factory supplies its own builder and the blank-key guard; a future Bedrock or Vertex embedding path gets the validation without copying it. That adds a byok -> embabel-agent-ai edge, for EmbeddingService. Acyclic and free in practice: embabel-agent-ai does not reference byok, and byok's three consumers already depend on embabel-agent-api, which brings ai transitively. The spec classes stay where they are. Both construct OpenAiCompatibleModelFactory in buildValidated(), and embabel-agent-openai already depends on byok, so moving them would close a cycle - byok -> openai -> byok. Signed-off-by: jasper blues <jasper@liberation-data.com>
The KDoc promised it throws on a blank key. It never sees a key — build() closes over whatever credential the provider needs — so it cannot check one, and saying otherwise invites a caller to skip their own guard. Says what it actually does now, and where the blank check belongs. Also noted why the second build() sits outside the try: it only constructs a service, so a failure there is a bug in the builder rather than a rejected credential and should not be reported as one. The test's builder returned a Pair of a lambda and a recording list, which reads as noise at every call site. A small named type instead. Signed-off-by: jasper blues <jasper@liberation-data.com>
…te test The comment claimed two failure modes share the exit; there are three, because build(null) is inside the try too. That is deliberate - for several providers a malformed base URL or unusable credential surfaces while constructing the client rather than on the first call - so the comment now says so, and says why the second build stays outside. The dimension test existed twice under different names with an identical body. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
igordayen
left a comment
igordayen
left a comment
There was a problem hiding this comment.
@jasperblues - looks good. major inquiry about validation results caching and connection with (in future) with capability model. thanks
| * @throws InvalidApiKeyException if the probe fails - an invalid key, an unreachable provider, a | ||
| * model the key cannot use - or if the model returns no vector | ||
| */ | ||
| fun validatedEmbeddingService( |
There was a problem hiding this comment.
Is the result cached? Please see how the probe is done for streaming
Also, is it strategically going to be part of the capability model (the one you wanted to have)?
There was a problem hiding this comment.
Yes, effectively cached — and having looked at the streaming probe as you suggested, I would rather not copy it.
Here: the probe runs once inside buildValidated(), and the width it observes is stamped onto the returned service as configuredDimensions. So dimensions on the result never re-probes. One network call per service built, never per read.
Streaming: StreamingCapabilityVerifier.supportsStreaming(chatModel) opens a real stream with the prompt "Say 'test' to confirm streaming works", consumes it under a 100ms timeout, and caches nothing. SpringAiLlmService.supportsStreaming() calls it every time, and that is reachable per-turn — StreamingPromptRunner and DelegatingStreamingPromptRunner both call supportsStreaming() before streaming. So each check can be a live, billable provider request.
I think that is a bug rather than a pattern to follow, and I would rather raise it separately than have this PR match it. Happy to open that if you agree it is one.
On the capability model: my understanding is that is the traits API discussion, which has not moved yet. I have deliberately not designed against it here — nothing in this PR presupposes or forecloses it, and EmbeddingService gains no new capability surface. When traits do land, "has a usable key yet" and "what width does this produce" both look like things that would want expressing there, and I would rather fit them to the real shape than guess at it now.
There was a problem hiding this comment.
Opened it: #1900.
Reading the streaming probe properly for the comparison turned up a second thing beyond the cost, which I have included there — the two catch clauses:
} catch (e: UnsupportedOperationException) {
false
} catch (e: Exception) {
false
}conflate "this model cannot stream" with "rate limited", "network blip", "bad key". A transient provider failure reports the model as non-streaming and the caller quietly takes a non-streaming path, so a wrong answer looks like a correct one. Arguably the worse of the two problems.
Nothing to change on this PR as a result — the embedding path already probes once per service built and stamps the result. Your question is what sent me to look, so thanks for it.
| @@ -24,6 +24,16 @@ | |||
| <groupId>org.jetbrains.kotlin</groupId> | |||
| <artifactId>kotlin-stdlib</artifactId> | |||
There was a problem hiding this comment.
Correct — it was inherited all along, and removed now.
embabel-agent-byok was the only module in the tree declaring kotlin-stdlib explicitly. embabel-agent-ai, embabel-agent-openai and embabel-agent-anthropic all compile Kotlin without it. Deleted, and the module still builds.
| val factory = FakeFactory { FloatArray(1536) } | ||
|
|
||
| val service = factory.buildValidatedEmbeddingService( | ||
| model = "text-embedding-3-small", |
There was a problem hiding this comment.
Done. SMALL_MODEL, LARGE_MODEL and their widths are named in a companion now, rather than the ids being repeated across a dozen assertions.
One literal stays a literal deliberately — the deliberately-mistyped text-embedding-3-smal in the test that a validation failure names the model. Naming that one would hide the point of it.
Review points. Model ids were repeated as literals across a dozen assertions, so a retired catalogue id would be a hunt rather than an edit. Named constants now, widths alongside them. embabel-agent-byok was the only module in the tree declaring kotlin-stdlib explicitly — no sibling does, and it compiles Kotlin fine without it, so it was inherited all along. Also carries a de-duplication that was already in the working tree: two of my tests asserted the same thing about the probed width, and are now one, better named. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: jasper blues <jasper@liberation-data.com>
|
3 participants
Closes #1883.
BYOK modelled LLMs but not embeddings. An
EmbeddingServicecould only come fromModelProvider, which resolves against models registered at boot, so a key supplied at runtime could not produce one at all — a deployment that collects keys at first run had to restart before embeddings became available. That is the restart constraint BYOK already removes on the LLM side, still in place on the embedding side.Guide sidesteps this today by using local ONNX embeddings, which need no key and so are unaffected by the boot-time resolution. That works, but it is a workaround rather than a choice: a deployment that wants a hosted embedding model with a user-supplied key has had no path at all.
What this adds
ByokFactory<EmbeddingService>, reached viaopenAiEmbedding()orbyokEmbedding():It probes with a single
embed()and translates any provider error toInvalidApiKeyException— the same contract, and the same single failure mode, as the LLM side.Two deliberate asymmetries with LLM BYOK
An embedding model is a schema commitment, not a stateless per-call service, so copying the LLM shape verbatim would be wrong in two specific ways.
No provider racing.
detectProvider()is sound when any candidate that accepts the key is an acceptable answer. Here, whichever provider answered first would decide the index dimension. The model is always named explicitly, anddetectProvider's KDoc now says why it is LLM-only.No caller-supplied dimension. The width is whatever the model actually returned during validation. This started out as a
configuredDimensionsparameter with a mismatch check, and was dropped: the returned service now carries a reliable stamped width, so the caller can compare it against their index in one line, at the layer that actually owns the index. Keeping it would have added a second exception type to a contract that has exactly one.Stamping the probed width also removes a hidden network call. Spring AI's
AbstractEmbeddingModel.dimensions()does cache, but it resolves viadimensions(this, "Test", "Hello World")— passing the literal"Test"as the model name, so its known-dimensions lookup table can never hit and it always falls through to a liveembed(), at whatever moment something first touches the property. We are already embedding to validate the key, so the width is free information.Scope is installation or world level, not per-user. Per-user LLM keys are coherent; per-user embedding keys are not, because vectors from different models and widths would land in one shared index. Documented on the type so the per-user shape is not adopted by default.
Note for the consuming side
The mismatch guard did not disappear, it moved. Something now has to do
check(embeddings.dimensions == index.dimensions)before writing into an existing index, and nothing enforces that it does. The docs section shows the check; worth confirming the index layer in embabel/me#735 has one.Testing
OpenAiCompatibleModelFactoryByokEmbeddingTest— 9 cases against a fake embedding endpoint: factory types, single probe, error translation, empty vector rejected, dimension derived from the model, pricing passthrough.OpenAiCompatibleModelFactoryByokIT(env-gated onOPENAI_API_KEY), both verified passing against the real API including the 1536-dimension assertion.embabel-agent-openai,embabel-agent-byok,embabel-agent-ai: 37 run, 0 failures. openai / lmstudio / google-genai autoconfigure consumers rebuild clean.openAiCompatibleEmbeddingServicegainsopen(so the probe path is testable without network) and otherwise keeps its signature, so existing autoconfigure callers are untouched.Docs: new "Embedding services" subsection in
reference/customizing/page.adoc.