diff --git a/spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironment.java b/spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironment.java index 6e02dfa20..ddfba20fb 100644 --- a/spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironment.java +++ b/spring-ai-spring-boot-docker-compose/src/main/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironment.java @@ -19,10 +19,20 @@ import java.util.Map; class ChromaEnvironment { + // Chroma version <= 0.4.x + private static final String CHROMA_SERVER_AUTH_CREDENTIALS = "CHROMA_SERVER_AUTH_CREDENTIALS"; + + // Chroma version >= 0.5.x + private static final String CHROMA_SERVER_AUTHN_CREDENTIALS = "CHROMA_SERVER_AUTHN_CREDENTIALS"; + private final String keyToken; ChromaEnvironment(Map env) { - this.keyToken = env.get("CHROMA_SERVER_AUTHN_CREDENTIALS"); + if (env.containsKey(CHROMA_SERVER_AUTH_CREDENTIALS)) { + this.keyToken = env.get(CHROMA_SERVER_AUTH_CREDENTIALS); + return; + } + this.keyToken = env.get(CHROMA_SERVER_AUTHN_CREDENTIALS); } public String getKeyToken() { diff --git a/spring-ai-spring-boot-docker-compose/src/test/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironmentTests.java b/spring-ai-spring-boot-docker-compose/src/test/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironmentTests.java index 8f4eac436..df3e514d0 100644 --- a/spring-ai-spring-boot-docker-compose/src/test/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironmentTests.java +++ b/spring-ai-spring-boot-docker-compose/src/test/java/org/springframework/ai/docker/compose/service/connection/chroma/ChromaEnvironmentTests.java @@ -31,7 +31,13 @@ class ChromaEnvironmentTests { } @Test - void getKeyTokenWhenHasCredential() { + void getKeyTokenFromAuthCredentialsWhenHasCredential() { + ChromaEnvironment environment = new ChromaEnvironment(Map.of("CHROMA_SERVER_AUTH_CREDENTIALS", "secret")); + assertThat(environment.getKeyToken()).isEqualTo("secret"); + } + + @Test + void getKeyTokenFromAuthnCredentialsWhenHasCredential() { ChromaEnvironment environment = new ChromaEnvironment(Map.of("CHROMA_SERVER_AUTHN_CREDENTIALS", "secret")); assertThat(environment.getKeyToken()).isEqualTo("secret"); } diff --git a/spring-ai-spring-boot-testcontainers/src/main/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaContainerConnectionDetailsFactory.java b/spring-ai-spring-boot-testcontainers/src/main/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaContainerConnectionDetailsFactory.java index e60b7bebe..5efb6b1af 100644 --- a/spring-ai-spring-boot-testcontainers/src/main/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaContainerConnectionDetailsFactory.java +++ b/spring-ai-spring-boot-testcontainers/src/main/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaContainerConnectionDetailsFactory.java @@ -20,6 +20,8 @@ import org.springframework.boot.testcontainers.service.connection.ContainerConne import org.springframework.boot.testcontainers.service.connection.ContainerConnectionSource; import org.testcontainers.chromadb.ChromaDBContainer; +import java.util.Map; + /** * @author Eddú Meléndez */ @@ -37,6 +39,12 @@ class ChromaContainerConnectionDetailsFactory private static final class ChromaDBContainerConnectionDetails extends ContainerConnectionDetails implements ChromaConnectionDetails { + // Chroma version <= 0.4.x + private static final String CHROMA_SERVER_AUTH_CREDENTIALS = "CHROMA_SERVER_AUTH_CREDENTIALS"; + + // Chroma version >= 0.5.x + private static final String CHROMA_SERVER_AUTHN_CREDENTIALS = "CHROMA_SERVER_AUTHN_CREDENTIALS"; + private ChromaDBContainerConnectionDetails(ContainerConnectionSource source) { super(source); } @@ -53,7 +61,11 @@ class ChromaContainerConnectionDetailsFactory @Override public String getKeyToken() { - return getContainer().getEnvMap().get("CHROMA_SERVER_AUTHN_CREDENTIALS"); + Map envVars = getContainer().getEnvMap(); + if (envVars.containsKey(CHROMA_SERVER_AUTH_CREDENTIALS)) { + return envVars.get(CHROMA_SERVER_AUTH_CREDENTIALS); + } + return envVars.get(CHROMA_SERVER_AUTHN_CREDENTIALS); } } diff --git a/spring-ai-spring-boot-testcontainers/src/test/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaWithToken2ContainerConnectionDetailsFactoryTest.java b/spring-ai-spring-boot-testcontainers/src/test/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaWithToken2ContainerConnectionDetailsFactoryTest.java new file mode 100644 index 000000000..9c134acb2 --- /dev/null +++ b/spring-ai-spring-boot-testcontainers/src/test/java/org/springframework/ai/testcontainers/service/connection/chroma/ChromaWithToken2ContainerConnectionDetailsFactoryTest.java @@ -0,0 +1,96 @@ +/* + * Copyright 2023 - 2024 the original author or authors. + * + * 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 + * + * https://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 org.springframework.ai.testcontainers.service.connection.chroma; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.autoconfigure.vectorstore.chroma.ChromaVectorStoreAutoConfiguration; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.autoconfigure.ImportAutoConfiguration; +import org.springframework.boot.testcontainers.service.connection.ServiceConnection; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.test.context.TestPropertySource; +import org.springframework.test.context.junit.jupiter.SpringJUnitConfig; +import org.testcontainers.chromadb.ChromaDBContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +@SpringJUnitConfig +@Testcontainers +@TestPropertySource(properties = "spring.ai.vectorstore.chroma.store.collectionName=TestCollection") +class ChromaWithToken2ContainerConnectionDetailsFactoryTest { + + @Container + @ServiceConnection + static ChromaDBContainer chroma = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.24") + .withEnv("CHROMA_SERVER_AUTH_CREDENTIALS", "token") + .withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER", + "chromadb.auth.token.TokenConfigServerAuthCredentialsProvider") + .withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.token.TokenAuthServerProvider"); + + @Autowired + private VectorStore vectorStore; + + @Test + public void addAndSearchWithFilters() { + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "Bulgaria")); + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "Netherlands")); + + vectorStore.add(List.of(bgDocument, nlDocument)); + + var request = SearchRequest.query("The World").withTopK(5); + + List results = vectorStore.similaritySearch(request); + assertThat(results).hasSize(2); + + results = vectorStore + .similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Bulgaria'")); + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + results = vectorStore + .similaritySearch(request.withSimilarityThresholdAll().withFilterExpression("country == 'Netherlands'")); + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + // Remove all documents from the store + vectorStore.delete(List.of(bgDocument, nlDocument).stream().map(doc -> doc.getId()).toList()); + } + + @Configuration(proxyBeanMethods = false) + @ImportAutoConfiguration(ChromaVectorStoreAutoConfiguration.class) + static class Config { + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + } + +}