BYOK: build an embedding service from a runtime-supplied key

embabel-agent-common/embabel-agent-byok/pom.xml

@@ -24,6 +24,16 @@
             <groupId>org.jetbrains.kotlin</groupId>
             <artifactId>kotlin-stdlib</artifactId>
         </dependency>
+
+        <!--
+          For EmbeddingService, which embeddingValidation.kt is typed to. Acyclic: embabel-agent-ai
+          does not reference this module, and every consumer of this one already depends on
+          embabel-agent-api, which brings embabel-agent-ai transitively.
+        -->
+        <dependency>
+            <groupId>com.embabel.agent</groupId>
+            <artifactId>embabel-agent-ai</artifactId>
+        </dependency>
     </dependencies>
 
     <build>

embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/ProviderDetection.kt

@@ -40,6 +40,12 @@ import java.util.concurrent.Executors
  * val service = detectProvider(AnthropicModelFactory(apiKey = userKey))
  * ```
  *
+ * Racing is only sound when any candidate that accepts the key is an acceptable answer. That
+ * holds for chat models, which are stateless per call. It does not hold for embedding models:
+ * the vector index is built at a fixed dimension, so which provider replies first must not be
+ * allowed to decide it. Build an embedding service from an explicitly chosen model instead of
+ * racing candidates here.
+ *
  * @param candidates One or more [ByokFactory] instances to race.
  * @return The service returned by the first successful factory.
  * @throws IllegalArgumentException if no candidates are supplied.

embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/embeddingValidation.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2024-2026 Embabel Pty Ltd.
+ *
+ * Licensed 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 com.embabel.common.byok
+
+import com.embabel.common.ai.model.EmbeddingService
+
+/** Short, cheap text used to probe an embedding key. */
+private const val EMBEDDING_VALIDATION_PROBE = "Hi"
+
+/**
+ * Validate a runtime-supplied embedding key by embedding one probe, and return a service that
+ * knows the width the model actually produced.
+ *
+ * Provider-agnostic: [build] supplies the provider's own [EmbeddingService], called once to
+ * probe and once more to produce the result with the observed width stamped on it.
+ *
+ * Stamping matters beyond saving a call. Spring AI's `AbstractEmbeddingModel.dimensions()`
+ * caches, but resolves via `dimensions(this, "Test", "Hello World")` — passing the literal
+ * `"Test"` as the model name, so its known-dimensions table can never hit and it always falls
+ * through to a live `embed()`. Left lazy that happens at an arbitrary later moment and can fail
+ * there. The probe already has the width; take it.
+ *
+ * The width is never taken from the caller: it is whatever the model returned. A caller writing
+ * into an existing index should compare [EmbeddingService.dimensions] against that index's width
+ * before storing anything — vectors of different widths cannot share an index. That check belongs
+ * to whatever owns the index.
+ *
+ * This never sees the API key — [build] closes over whatever credential the provider needs — so
+ * it cannot reject a blank one. Callers holding a key should check that themselves before getting
+ * here, so an empty string fails with something actionable rather than as a provider error.
+ *
+ * @param model the embedding model to validate. Always explicit: unlike an LLM it is never
+ * defaulted and never detected, because the width it commits an index to must not be decided by
+ * whichever provider happens to answer first.
+ * @param provider provider name, used only in the failure message
+ * @param build builds the provider's service, given the width to stamp (null while probing)
+ * @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(
+    model: String,
+    provider: String,
+    build: (configuredDimensions: Int?) -> EmbeddingService,
+): EmbeddingService {
+    // All three failure modes here mean the same thing to the caller - the model could not be
+    // validated - so they share one exit rather than throwing from three places. Construction is
+    // inside the try on purpose: building the probe service is where a malformed base URL or an
+    // unusable credential surfaces for several providers, so it is part of validation rather than
+    // separate from it. The message names the model: it is caller-supplied and mandatory, so a
+    // typo in it arrives as the same provider error, and "invalid API key" alone would send
+    // someone to re-check a key that was fine.
+    val vector = try {
+        build(null).embed(EMBEDDING_VALIDATION_PROBE)
+            .also { require(it.isNotEmpty()) { "the model returned an empty vector" } }
+    } catch (e: Exception) {
+        throw InvalidApiKeyException(
+            "Could not validate embedding model '$model' on $provider: ${e.message ?: "no detail"}",
+        )
+    }
+    // Outside the try, unlike the call above. The credential and the model have both just been
+    // proven to work, so anything thrown here is a bug in the builder and must not be reported as
+    // a rejected key - that would send someone to re-check a key this function just validated.
+    return build(vector.size)
+}

embabel-agent-common/embabel-agent-byok/src/test/kotlin/com/embabel/common/byok/EmbeddingValidationTest.kt

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2024-2026 Embabel Pty Ltd.
+ *
+ * Licensed 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 com.embabel.common.byok
+
+import com.embabel.common.ai.model.EmbeddingService
+import com.embabel.common.ai.model.PricingModel
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.assertThrows
+
+/**
+ * The probe-and-stamp algorithm, independent of any provider. It lives here rather than in a
+ * provider module because nothing about it is provider-specific: the caller supplies the
+ * service, this supplies the validation.
+ */
+class EmbeddingValidationTest {
+
+    private class FakeEmbeddingService(
+        override val name: String,
+        override val provider: String,
+        private val configuredDimensions: Int?,
+        override val pricingModel: PricingModel? = null,
+        private val respond: (String) -> FloatArray,
+    ) : EmbeddingService {
+        override fun embed(text: String): FloatArray = respond(text)
+        override fun embed(texts: List<String>): List<FloatArray> = texts.map(respond)
+        override val dimensions: Int
+            get() = configuredDimensions ?: error("dimensions not configured")
+    }
+
+    /**
+     * Stands in for a provider's builder, recording the width it was asked to stamp on each call.
+     */
+    private class RecordingBuilder(
+        private val respond: (String) -> FloatArray,
+    ) : (Int?) -> EmbeddingService {
+
+        val widthsRequested = mutableListOf<Int?>()
+
+        override fun invoke(configuredDimensions: Int?): EmbeddingService {
+            widthsRequested += configuredDimensions
+            return FakeEmbeddingService(
+                name = "text-embedding-3-small",
+                provider = "acme",
+                configuredDimensions = configuredDimensions,
+                respond = respond,
+            )
+        }
+    }
+
+    @Test
+    fun `probes once, then stamps the width the model actually returned`() {
+        var probes = 0
+        val build = RecordingBuilder { probes++; FloatArray(3072) }
+
+        val service = validatedEmbeddingService("text-embedding-3-large", "acme", build)
+
+        assertEquals(1, probes, "should probe exactly once")
+        assertEquals(3072, service.dimensions)
+        assertEquals(
+            listOf(null, 3072),
+            build.widthsRequested,
+            "probe with no width, then stamp the one observed",
+        )
+    }
+
+    @Test
+    fun `a provider error is translated, and names the model`() {
+        val build = RecordingBuilder { throw RuntimeException("404 model not found") }
+
+        val e = assertThrows<InvalidApiKeyException> {
+            validatedEmbeddingService("text-embedding-3-smal", "acme", build)
+        }
+
+        // The model is caller-supplied and mandatory here, so a typo in it is at least as likely
+        // as a bad key and must be visible in the message.
+        assertTrue(e.message!!.contains("text-embedding-3-smal"), e.message)
+        assertTrue(e.message!!.contains("acme"), e.message)
+    }
+
+    @Test
+    fun `an empty vector is a validation failure, not a zero-width service`() {
+        val build = RecordingBuilder { FloatArray(0) }
+
+        assertThrows<InvalidApiKeyException> {
+            validatedEmbeddingService("text-embedding-3-small", "acme", build)
+        }
+    }
+}

embabel-agent-docs/src/main/asciidoc/reference/customizing/page.adoc

@@ -307,6 +307,64 @@ val service = detectProvider(
 )
 ----
 
+===== Embedding services
+
+Embedding models are resolved through the same BYOK mechanism, so a key that arrives after
+startup can activate an embedding model without a restart.
+
+[source,kotlin]
+----
+// OpenAI
+val embeddings: EmbeddingService =
+    OpenAiCompatibleModelFactory.openAiEmbedding(userKey, "text-embedding-3-small")
+        .buildValidated()
+
+// Any OpenAI-compatible provider
+val acmeEmbeddings: EmbeddingService = OpenAiCompatibleModelFactory.byokEmbedding(
+    baseUrl = "https://api.acme.example.com/v1",
+    apiKey = userKey,
+    model = "acme-embed-small",
+    provider = "Acme",
+).buildValidated()
+----
+
+`buildValidated()` embeds a short probe text. On failure it throws `InvalidApiKeyException`.
+A *blank* key is treated as absent and rejected before any network call — in a container
+`OPENAI_API_KEY` is routinely set-but-empty, and that empty string would otherwise reach the
+provider and come back as an opaque authentication error. The message names the model: an
+embedding model name is caller-supplied and mandatory here, so a typo in it reaches the provider
+the same way a bad key does. On success the returned `EmbeddingService` already knows its
+dimension — it was observed during the probe — so no live `dimensions()` call is made later.
+
+Two things work differently from LLM BYOK, because an embedding model is a schema commitment
+rather than a stateless per-call service:
+
+* **No provider detection.** The vector index is built at a fixed dimension and every stored
+  vector must come from the same model at that dimension. Whichever provider happens to accept
+  a pasted key first must not be allowed to decide that, so there is no embedding equivalent of
+  `detectProvider()` and the model is always named explicitly.
+* **Installation or world scope, not per-user.** Per-user LLM keys are coherent; per-user
+  embedding keys are not, because vectors from different models would land in one shared index.
+  The dimension belongs to the store, not to the user.
+
+The dimension is never configured on the BYOK path — it is whatever the model actually returned
+during validation. When writing into an index that already exists, compare the two yourself
+before storing anything, since vectors of different widths cannot share an index:
+
+[source,kotlin]
+----
+val embeddings = OpenAiCompatibleModelFactory
+    .openAiEmbedding(userKey, "text-embedding-3-small")
+    .buildValidated()
+
+check(embeddings.dimensions == index.dimensions) {
+    "Index is ${index.dimensions}-dimensional; ${embeddings.name} returns ${embeddings.dimensions}"
+}
+----
+
+NOTE: Changing the embedding model of an existing corpus still requires re-embedding it. BYOK
+only builds and validates the service.
+
 ===== Using the validated service
 
 Once you have an `LlmService`, pass it directly to `PromptRunner` or `Ai` via