diff --git a/embabel-agent-common/embabel-agent-byok/pom.xml b/embabel-agent-common/embabel-agent-byok/pom.xml index 6eda5df59..57ad419bf 100644 --- a/embabel-agent-common/embabel-agent-byok/pom.xml +++ b/embabel-agent-common/embabel-agent-byok/pom.xml @@ -20,9 +20,14 @@ + - org.jetbrains.kotlin - kotlin-stdlib + com.embabel.agent + embabel-agent-ai diff --git a/embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/ProviderDetection.kt b/embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/ProviderDetection.kt index 0ddd8f182..72dcb9134 100644 --- a/embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/ProviderDetection.kt +++ b/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. diff --git a/embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/embeddingValidation.kt b/embabel-agent-common/embabel-agent-byok/src/main/kotlin/com/embabel/common/byok/embeddingValidation.kt new file mode 100644 index 000000000..65635f058 --- /dev/null +++ b/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) +} diff --git a/embabel-agent-common/embabel-agent-byok/src/test/kotlin/com/embabel/common/byok/EmbeddingValidationTest.kt b/embabel-agent-common/embabel-agent-byok/src/test/kotlin/com/embabel/common/byok/EmbeddingValidationTest.kt new file mode 100644 index 000000000..4480063b4 --- /dev/null +++ b/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): List = 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() + + 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 { + 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 { + validatedEmbeddingService("text-embedding-3-small", "acme", build) + } + } +} diff --git a/embabel-agent-docs/src/main/asciidoc/reference/customizing/page.adoc b/embabel-agent-docs/src/main/asciidoc/reference/customizing/page.adoc index 93b4eb71f..32a727c5d 100644 --- a/embabel-agent-docs/src/main/asciidoc/reference/customizing/page.adoc +++ b/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 diff --git a/embabel-agent-openai/src/main/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactory.kt b/embabel-agent-openai/src/main/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactory.kt index 81b135da6..ce0ae345a 100644 --- a/embabel-agent-openai/src/main/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactory.kt +++ b/embabel-agent-openai/src/main/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactory.kt @@ -25,6 +25,7 @@ import com.embabel.chat.UserMessage import com.embabel.common.ai.model.* import com.embabel.common.byok.ByokFactory import com.embabel.common.byok.InvalidApiKeyException +import com.embabel.common.byok.validatedEmbeddingService import com.embabel.common.util.ObjectProviders import com.openai.client.OpenAIClient import com.openai.client.OpenAIClientAsync @@ -87,6 +88,18 @@ open class OpenAiCompatibleModelFactory( private const val CONNECT_TIMEOUT_MS = 5_000L private const val READ_TIMEOUT_MS = 600_000L + /** + * Message for a key that is present but blank. + * + * Deliberately local. The LLM BYOK path grows the same rule in #1888, and a shared + * helper is the right home for it — but putting it there from here would couple these + * two branches for four lines. Collapse this onto that helper once both have landed. + */ + internal val BLANK_EMBEDDING_KEY_MESSAGE: String = """ + API key is blank. A blank key is treated as absent: + supply a non-empty key, or omit it and let the caller decide there is no key for this request. + """.trimIndent() + /** * Returns a [ByokSpec] for OpenAI. * Validates against [OpenAiModels.GPT_41_MINI] by default. @@ -147,6 +160,40 @@ open class OpenAiCompatibleModelFactory( validationModel: String, validationProvider: String, ): ByokSpec = ByokSpec(baseUrl, apiKey, validationModel, validationProvider) + + /** + * Returns a [ByokEmbeddingSpec] for OpenAI. + * + * [model] is required — see [ByokEmbeddingSpec] for why an embedding model is never + * defaulted or detected. + */ + fun openAiEmbedding( + apiKey: String, + model: String, + pricingModel: PricingModel? = null, + ): ByokEmbeddingSpec = + ByokEmbeddingSpec(null, apiKey, model, OpenAiModels.PROVIDER, pricingModel) + + /** + * Returns a [ByokEmbeddingSpec] for a custom OpenAI-compatible provider. + * + * ```kotlin + * OpenAiCompatibleModelFactory.byokEmbedding( + * baseUrl = "https://api.myprovider.com", + * apiKey = apiKey, + * model = "my-embed-small", + * provider = "MyProvider", + * ) + * ``` + */ + fun byokEmbedding( + baseUrl: String?, + apiKey: String, + model: String, + provider: String, + pricingModel: PricingModel? = null, + ): ByokEmbeddingSpec = + ByokEmbeddingSpec(baseUrl, apiKey, model, provider, pricingModel) } /** @@ -187,6 +234,52 @@ open class OpenAiCompatibleModelFactory( ) } + /** + * A self-contained BYOK spec that validates an API key and returns a ready + * [EmbeddingService], so a key supplied *after* startup can activate an embedding model + * without a restart. + * + * Obtained via [openAiEmbedding] or [byokEmbedding]. + * + * Deliberately unlike [ByokSpec] in two ways, because an embedding model is not the same + * kind of thing as an LLM: + * + * 1. **No provider detection.** An LLM is stateless per call, so racing candidate factories + * and taking whichever accepts the key is a fine answer. An embedding model is a schema + * commitment: the vector index is built at a fixed dimension and every stored vector must + * share that model and dimension. Which provider answers first must not decide the index + * dimension, so this type is not intended for + * [com.embabel.common.byok.detectProvider] and the model is always explicit. + * 2. **Installation or world scope, not per-user.** Per-user LLM keys are coherent; per-user + * *embedding* keys are not, because vectors from different models and dimensions would + * land in one shared index. The dimension is a property of the store, not of the user. + * Do not inherit a per-user credential story here. + * + * The returned service reports the dimension actually observed during validation, so a + * caller writing into an existing index can compare it against that index's width before + * storing anything. That check belongs to whatever owns the index, not here. + * + * Changing the embedding model of an existing corpus still requires re-embedding it. This + * type only builds and validates a service from a runtime-supplied key. + */ + class ByokEmbeddingSpec internal constructor( + private val baseUrl: String?, + private val apiKey: String, + private val model: String, + private val provider: String, + private val pricingModel: PricingModel? = null, + private val observationRegistry: ObservationRegistry = ObservationRegistry.NOOP, + ) : ByokFactory { + + override fun buildValidated(): EmbeddingService = + OpenAiCompatibleModelFactory(baseUrl, apiKey, null, null, observationRegistry = observationRegistry) + .buildValidatedEmbeddingService( + model = model, + provider = provider, + pricingModel = pricingModel, + ) + } + protected val logger: Logger = LoggerFactory.getLogger(javaClass) // Subclasses should add their own more specific logging @@ -307,7 +400,36 @@ open class OpenAiCompatibleModelFactory( ) } - fun openAiCompatibleEmbeddingService( + /** + * Validates the configured API key by embedding a short probe text, then returns a + * production [EmbeddingService] carrying the width the model actually produced. + * + * The probe-and-stamp logic is provider-agnostic and lives in + * [com.embabel.common.byok.validatedEmbeddingService]; this method supplies the + * OpenAI-compatible builder and the blank-key guard. + * + * @throws InvalidApiKeyException if the key is blank or invalid, the provider is + * unreachable, or the model returns no vector. + */ + fun buildValidatedEmbeddingService( + model: String, + provider: String, + pricingModel: PricingModel? = null, + ): EmbeddingService { + if (apiKey.isNullOrBlank()) { + throw InvalidApiKeyException(BLANK_EMBEDDING_KEY_MESSAGE) + } + return validatedEmbeddingService(model = model, provider = provider) { configuredDimensions -> + openAiCompatibleEmbeddingService( + model = model, + provider = provider, + configuredDimensions = configuredDimensions, + pricingModel = pricingModel, + ) + } + } + + open fun openAiCompatibleEmbeddingService( model: String, provider: String, configuredDimensions: Int? = null, diff --git a/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokEmbeddingTest.kt b/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokEmbeddingTest.kt new file mode 100644 index 000000000..12a16d41c --- /dev/null +++ b/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokEmbeddingTest.kt @@ -0,0 +1,236 @@ +/* + * 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.agent.openai + +import com.embabel.agent.api.models.OpenAiModels +import com.embabel.common.ai.model.EmbeddingService +import com.embabel.common.ai.model.PricingModel +import com.embabel.common.byok.ByokFactory +import com.embabel.common.byok.InvalidApiKeyException +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertInstanceOf +import org.junit.jupiter.api.Assertions.assertNotNull +import org.junit.jupiter.api.Assertions.assertSame +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows + +class OpenAiCompatibleModelFactoryByokEmbeddingTest { + + private companion object { + /** + * Named rather than repeated: these are OpenAI catalogue ids, and a model being retired + * should be one edit here rather than a hunt through every assertion. + */ + const val SMALL_MODEL = "text-embedding-3-small" + const val LARGE_MODEL = "text-embedding-3-large" + const val SMALL_MODEL_DIMENSIONS = 1536 + const val LARGE_MODEL_DIMENSIONS = 3072 + } + + /** + * Stands in for a live embedding endpoint so the validation logic can be exercised + * without a network call. + */ + private class FakeEmbeddingService( + override val name: String, + override val provider: String, + 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): List = texts.map(respond) + + override val dimensions: Int + get() = configuredDimensions ?: error("dimensions not configured") + } + + /** + * Replaces the live embedding service with a fake, recording what the factory asked for. + */ + private class FakeFactory( + apiKey: String = "key", + private val respond: (String) -> FloatArray, + ) : OpenAiCompatibleModelFactory(baseUrl = null, apiKey = apiKey) { + + val builds = mutableListOf() + + override fun openAiCompatibleEmbeddingService( + model: String, + provider: String, + configuredDimensions: Int?, + pricingModel: PricingModel?, + ): EmbeddingService = FakeEmbeddingService( + name = model, + provider = provider, + configuredDimensions = configuredDimensions, + pricingModel = pricingModel, + respond = respond, + ).also { builds += it } + } + + @Test + fun `openAiEmbedding returns a ByokFactory of EmbeddingService`() { + val spec: ByokFactory = + OpenAiCompatibleModelFactory.openAiEmbedding("key", SMALL_MODEL) + assertInstanceOf(ByokFactory::class.java, spec) + } + + @Test + fun `byokEmbedding with explicit params returns a ByokFactory`() { + assertInstanceOf( + ByokFactory::class.java, + OpenAiCompatibleModelFactory.byokEmbedding( + baseUrl = "https://api.example.com", + apiKey = "key", + model = "acme-embed-small", + provider = "Acme", + ), + ) + } + + @Test + fun `openAiEmbedding uses the OpenAI provider`() { + val service = FakeFactory { FloatArray(SMALL_MODEL_DIMENSIONS) } + .buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + ) + assertEquals(OpenAiModels.PROVIDER, service.provider) + assertEquals(SMALL_MODEL, service.name) + } + + @Test + fun `valid key returns a service and probes exactly once`() { + var probes = 0 + val factory = FakeFactory { probes++; FloatArray(SMALL_MODEL_DIMENSIONS) { 0.1f } } + + val service = factory.buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + ) + + assertNotNull(service) + assertEquals(1, probes, "buildValidated should probe once") + } + + @Test + fun `probe failure is translated to InvalidApiKeyException`() { + val factory = FakeFactory { throw RuntimeException("401 Unauthorized") } + + val e = assertThrows { + factory.buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + ) + } + assertTrue(e.message!!.contains("401 Unauthorized"), e.message) + } + + @Test + fun `a validation failure names the model, since a typo in it looks the same as a bad key`() { + // The embedding model is caller-supplied and mandatory — never defaulted, never detected — + // so "invalid API key" alone would send someone to re-check a key that was fine. + val factory = FakeFactory { throw RuntimeException("404 model not found") } + + val e = assertThrows { + factory.buildValidatedEmbeddingService( + model = "text-embedding-3-smal", + provider = OpenAiModels.PROVIDER, + ) + } + + assertTrue(e.message!!.contains("text-embedding-3-smal"), e.message) + assertTrue(e.message!!.contains(OpenAiModels.PROVIDER), e.message) + } + + @Test + fun `empty probe vector is rejected`() { + val factory = FakeFactory { FloatArray(0) } + + assertThrows { + factory.buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + ) + } + } + + @Test + fun `returned service carries the probed dimension so no live dimensions call is needed`() { + // The width comes from the model and only from the model: the caller has no way to declare + // one here, so a store can compare it against its index rather than trusting config. + val factory = FakeFactory { FloatArray(LARGE_MODEL_DIMENSIONS) } + + val service = factory.buildValidatedEmbeddingService( + model = LARGE_MODEL, + provider = OpenAiModels.PROVIDER, + ) + + assertEquals(LARGE_MODEL_DIMENSIONS, service.dimensions) + } + + @Test + fun `a blank key is rejected without probing`() { + // Compose passes OPENAI_API_KEY=${OPENAI_API_KEY:-}, so set-but-empty is routine. That + // empty string passes a null check and would otherwise reach the provider, coming back as + // an opaque authentication error a long way from its cause. + listOf("", " ", "\t").forEach { blank -> + val factory = FakeFactory(apiKey = blank) { FloatArray(SMALL_MODEL_DIMENSIONS) } + + val e = assertThrows { + factory.buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + ) + } + + assertEquals(OpenAiCompatibleModelFactory.BLANK_EMBEDDING_KEY_MESSAGE, e.message) + assertEquals(0, factory.builds.size, "a blank key must not build or probe anything") + } + } + + @Test + fun `the BYOK embedding entry points reject a blank key too`() { + listOf( + OpenAiCompatibleModelFactory.openAiEmbedding(" ", SMALL_MODEL), + OpenAiCompatibleModelFactory.byokEmbedding( + baseUrl = "https://api.example.com", + apiKey = "", + model = "acme-embed-small", + provider = "Acme", + ), + ).forEach { spec -> + val e = assertThrows { spec.buildValidated() } + assertEquals(OpenAiCompatibleModelFactory.BLANK_EMBEDDING_KEY_MESSAGE, e.message) + } + } + + @Test + fun `pricing model is carried through to the returned service`() { + val factory = FakeFactory { FloatArray(SMALL_MODEL_DIMENSIONS) } + + val service = factory.buildValidatedEmbeddingService( + model = SMALL_MODEL, + provider = OpenAiModels.PROVIDER, + pricingModel = PricingModel.ALL_YOU_CAN_EAT, + ) + + assertSame(PricingModel.ALL_YOU_CAN_EAT, service.pricingModel) + } +} diff --git a/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokIT.kt b/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokIT.kt index 9fe406e5d..f13ea8e7f 100644 --- a/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokIT.kt +++ b/embabel-agent-openai/src/test/kotlin/com/embabel/agent/openai/OpenAiCompatibleModelFactoryByokIT.kt @@ -15,8 +15,11 @@ */ package com.embabel.agent.openai +import com.embabel.common.byok.InvalidApiKeyException +import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertNotNull import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable /** @@ -56,4 +59,24 @@ class OpenAiCompatibleModelFactoryByokIT { .buildValidated() assertNotNull(service) } + + @Test + @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") + fun `openAiEmbedding buildValidated succeeds with valid key`() { + val service = OpenAiCompatibleModelFactory + .openAiEmbedding(System.getenv("OPENAI_API_KEY"), "text-embedding-3-small") + .buildValidated() + assertNotNull(service) + assertEquals(1536, service.dimensions) + } + + @Test + @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") + fun `openAiEmbedding rejects an invalid key`() { + assertThrows { + OpenAiCompatibleModelFactory + .openAiEmbedding("sk-not-a-real-key", "text-embedding-3-small") + .buildValidated() + } + } }