diff --git a/dice-storage-autoconfigure/src/main/kotlin/com/embabel/dice/storage/autoconfigure/DiceStorageAutoConfiguration.kt b/dice-storage-autoconfigure/src/main/kotlin/com/embabel/dice/storage/autoconfigure/DiceStorageAutoConfiguration.kt index 9b8759dc..53a247b3 100644 --- a/dice-storage-autoconfigure/src/main/kotlin/com/embabel/dice/storage/autoconfigure/DiceStorageAutoConfiguration.kt +++ b/dice-storage-autoconfigure/src/main/kotlin/com/embabel/dice/storage/autoconfigure/DiceStorageAutoConfiguration.kt @@ -97,7 +97,7 @@ class DiceStorageAutoConfiguration { DrivinePropositionRepository.VECTOR_INDEX, ) return DrivinePropositionRepository( - graphObjectManager, persistenceManager, embeddingService(ai, embeddingServices), transactionManager, + graphObjectManager, persistenceManager, storeEmbeddingService("graph", ai, embeddingServices), transactionManager, ) } @@ -112,10 +112,38 @@ class DiceStorageAutoConfiguration { * eagerly and throws when no model is registered, taking the context down with it. * [DrivinePropositionRepository] only touches the service when it actually embeds, so * an absent-tolerant one is safe to hold. + * + * Several beans and no `@Primary` is the ORDINARY case, not an ambiguity to fail on: the platform + * registers one bean per model it can reach (an OpenAI key alone registers three), and choosing + * among those is exactly what the configured default embedding model is for. A host that wants + * its own service used marks it `@Primary`. */ private fun embeddingService(ai: Ai, embeddingServices: ObjectProvider): EmbeddingService = embeddingServices.getIfUnique() ?: ai.withDefaultEmbeddingService() + /** + * [embeddingService] for a proposition store, saying which one it got and why. Both stores go + * through here, so a host's embedding service reaches whichever backend it runs, and the one + * line an operator needs — is memory embedded with the model I configured? — is at startup + * rather than inferred from a vector-width mismatch later. + */ + private fun storeEmbeddingService( + store: String, + ai: Ai, + embeddingServices: ObjectProvider, + ): EmbeddingService { + val hostService = embeddingServices.getIfUnique() + val service = hostService ?: ai.withDefaultEmbeddingService() + logger.info( + "{} proposition store embeds with '{}' ({}): {}", + store, + service.name, + service.provider, + if (hostService != null) "the application's EmbeddingService bean" else "the platform default embedding service", + ) + return service + } + @Bean @ConditionalOnProperty(prefix = "embabel.dice.store", name = ["type"], havingValue = "graph") @ConditionalOnMissingBean(ChunkHistoryStore::class) @@ -352,9 +380,12 @@ class DiceStorageAutoConfiguration { @Bean @ConditionalOnBean(Ai::class) @ConditionalOnMissingBean(PropositionRepository::class) - fun inMemoryPropositionRepository(ai: Ai): PropositionRepository { + fun inMemoryPropositionRepository( + ai: Ai, + embeddingServices: ObjectProvider, + ): PropositionRepository { logger.info("Wiring in-memory proposition store") - return InMemoryPropositionRepository(ai.withDefaultEmbeddingService()) + return InMemoryPropositionRepository(storeEmbeddingService("in-memory", ai, embeddingServices)) } @Bean diff --git a/dice-storage-autoconfigure/src/test/kotlin/com/embabel/dice/storage/autoconfigure/PropositionStoreEmbeddingServiceTest.kt b/dice-storage-autoconfigure/src/test/kotlin/com/embabel/dice/storage/autoconfigure/PropositionStoreEmbeddingServiceTest.kt new file mode 100644 index 00000000..67ba16b9 --- /dev/null +++ b/dice-storage-autoconfigure/src/test/kotlin/com/embabel/dice/storage/autoconfigure/PropositionStoreEmbeddingServiceTest.kt @@ -0,0 +1,89 @@ +/* + * 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.dice.storage.autoconfigure + +import com.embabel.agent.api.common.Ai +import com.embabel.common.ai.model.EmbeddingService +import com.embabel.dice.proposition.PropositionRepository +import org.junit.jupiter.api.Test +import org.mockito.kotlin.any +import org.mockito.kotlin.mock +import org.mockito.kotlin.never +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever +import org.springframework.boot.autoconfigure.AutoConfigurations +import org.springframework.boot.test.context.runner.ApplicationContextRunner +import org.springframework.context.annotation.Bean +import org.springframework.context.annotation.Configuration +import org.springframework.context.annotation.Primary + +/** + * The proposition store embeds with the host's embedding service whichever backend it runs. + * + * The graph backend already preferred a unique or `@Primary` host bean; the in-memory backend went + * straight to the platform default, so a host that swapped its model at runtime had in-memory + * propositions embedded by a different one. These pin the in-memory half; both halves now share one + * resolution. + */ +class PropositionStoreEmbeddingServiceTest { + + private val platformDefault: EmbeddingService = embeddingService("platform-default") + private val hostService: EmbeddingService = embeddingService("host-model") + + private val runner = ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(DiceStorageAutoConfiguration::class.java)) + .withBean(Ai::class.java, { + mock().also { whenever(it.withDefaultEmbeddingService()).thenReturn(platformDefault) } + }) + + @Test + fun `the in-memory store embeds with the host's primary embedding service`() { + runner + .withUserConfiguration(HostEmbedding::class.java) + .withBean("hostService", EmbeddingService::class.java, { hostService }, { it.isPrimary = true }) + .run { ctx -> + ctx.getBean(PropositionRepository::class.java).save(TestFixtures.proposition("Alice likes the Alps")) + + verify(hostService).embed("Alice likes the Alps") + verify(platformDefault, never()).embed(any()) + } + } + + @Test + fun `the in-memory store embeds with the platform default when the host provides none`() { + runner.run { ctx -> + ctx.getBean(PropositionRepository::class.java).save(TestFixtures.proposition("Alice likes the Alps")) + + verify(platformDefault).embed("Alice likes the Alps") + } + } + + /** A second, non-primary model bean — the platform registers one per model it can reach. */ + @Configuration + open class HostEmbedding { + @Bean + open fun otherModel(): EmbeddingService = embeddingService("other-model") + } + + private companion object { + fun embeddingService(name: String): EmbeddingService = mock { + whenever(it.name).thenReturn(name) + whenever(it.provider).thenReturn("test") + whenever(it.dimensions).thenReturn(3) + whenever(it.embed(any())).thenReturn(floatArrayOf(1f, 0f, 0f)) + } + } +}