Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use Testcontainers Vector Database modules

Browse Source 
- Update testcontainers version to 1.19.6
 - Use ChromaDBContainer
 - Use MilvusContainer
 - Use WeaviateContainer
 - Fix format
This commit is contained in:
Eddú Meléndez
2024-02-21 17:32:17 -06:00
committed by Christian Tzolov
parent 21fc676fba
commit 6de57d5f37
10 changed files with 38 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
vector-stores/spring-ai-chroma/pom.xml
View File

@@ -48,7 +48,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<artifactId>chromadb</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

13
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -21,7 +21,7 @@ import java.util.Map;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -39,14 +39,14 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@SpringBootTest
@Testcontainers
public class ChromaApiIT {
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.22.dev44")
.withExposedPorts(8000);
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22.dev44");
@Autowired
ChromaApi chroma;
@@ -186,10 +186,7 @@ public class ChromaApiIT {
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
String host = chromaContainer.getHost();
int port = chromaContainer.getMappedPort(8000);
String baseUrl = "http://%s:%d".formatted(host, port);
return new ChromaApi(baseUrl, restTemplate);
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
}
}

16
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/BasicAuthChromaWhereIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -20,7 +20,7 @@ import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.images.builder.Transferable;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -46,6 +46,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* <code>htpasswd -Bbn admin admin > server.htpasswd</code>
*
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@Testcontainers
public class BasicAuthChromaWhereIT {
@@ -55,13 +56,12 @@ public class BasicAuthChromaWhereIT {
* https://docs.trychroma.com/usage-guide#basic-authentication
*/
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.22")
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_FILE", "server.htpasswd")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider")
.withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.basic.BasicAuthServerProvider")
.withCopyToContainer(Transferable.of("src/test/resources/server.htpasswd"), "server.htpasswd")
.withExposedPorts(8000);
.withCopyToContainer(Transferable.of("src/test/resources/server.htpasswd"), "server.htpasswd");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
@@ -103,10 +103,8 @@ public class BasicAuthChromaWhereIT {
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
String host = chromaContainer.getHost();
int port = chromaContainer.getMappedPort(8000);
String baseUrl = "http://%s:%d".formatted(host, port);
return new ChromaApi(baseUrl, restTemplate).withBasicAuthCredentials("admin", "admin");
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate).withBasicAuthCredentials("admin",
"admin");
}
@Bean

13
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -22,7 +22,7 @@ import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -41,13 +41,13 @@ import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@Testcontainers
public class ChromaVectorStoreIT {
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.22")
.withExposedPorts(8000);
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22");
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
@@ -210,10 +210,7 @@ public class ChromaVectorStoreIT {
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
String host = chromaContainer.getHost();
int port = chromaContainer.getMappedPort(8000);
String baseUrl = "http://%s:%d".formatted(host, port);
return new ChromaApi(baseUrl, restTemplate);
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
}
@Bean

16
vector-stores/spring-ai-chroma/src/test/java/org/springframework/ai/vectorstore/TokenSecuredChromaWhereIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -20,7 +20,7 @@ import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -46,6 +46,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* https://github.com/chroma-core/chroma/blob/main/examples/basic_functionality/in_not_in_filtering.ipynb
*
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@Testcontainers
public class TokenSecuredChromaWhereIT {
@@ -57,13 +58,11 @@ public class TokenSecuredChromaWhereIT {
* https://docs.trychroma.com/usage-guide#static-api-token-authentication
*/
@Container
static GenericContainer<?> chromaContainer = new GenericContainer<>("ghcr.io/chroma-core/chroma:0.4.22")
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS", CHROMA_SERVER_AUTH_CREDENTIALS)
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.token.TokenConfigServerAuthCredentialsProvider")
.withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.token.TokenAuthServerProvider")
.withExposedPorts(8000);
.withEnv("CHROMA_SERVER_AUTH_PROVIDER", "chromadb.auth.token.TokenAuthServerProvider");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
@@ -135,10 +134,7 @@ public class TokenSecuredChromaWhereIT {
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
String host = chromaContainer.getHost();
int port = chromaContainer.getMappedPort(8000);
String baseurl = "http://%s:%d".formatted(host, port);
var chromaApi = new ChromaApi(baseurl, restTemplate);
var chromaApi = new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
chromaApi.withKeyToken(CHROMA_SERVER_AUTH_CREDENTIALS);
return chromaApi;
}

2
vector-stores/spring-ai-milvus-store/pom.xml
View File

@@ -64,7 +64,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<artifactId>milvus</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

45
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -16,10 +16,8 @@
package org.springframework.ai.vectorstore;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.Map;
@@ -29,20 +27,15 @@ import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig;
import org.springframework.beans.factory.annotation.Value;
@@ -52,21 +45,20 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.util.FileSystemUtils;
import org.testcontainers.milvus.MilvusContainer;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class MilvusVectorStoreIT {
@SuppressWarnings("rawtypes")
private static DockerComposeContainer milvusContainer;
private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
@Container
private static MilvusContainer milvusContainer = new MilvusContainer("milvusdb/milvus:v2.3.8");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
@@ -86,28 +78,6 @@ public class MilvusVectorStoreIT {
}
}
@SuppressWarnings({ "rawtypes", "resource" })
@BeforeAll
public static void beforeAll() {
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
TEMP_FOLDER.mkdirs();
milvusContainer = new DockerComposeContainer(new File("src/test/resources/docker-compose.yml"))
.withEnv("DOCKER_VOLUME_DIRECTORY", TEMP_FOLDER.getAbsolutePath())
.withExposedService("standalone", 19530)
.withExposedService("standalone", 9091,
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
.waitingFor("standalone", Wait.forLogMessage(".*Proxy successfully started.*\\s", 1)
.withStartupTimeout(Duration.ofSeconds(100)));
milvusContainer.start();
}
@AfterAll
public static void afterAll() {
milvusContainer.stop();
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
}
private void resetCollection(VectorStore vectorStore) {
((MilvusVectorStore) vectorStore).dropCollection();
((MilvusVectorStore) vectorStore).createCollection();
@@ -303,8 +273,7 @@ public class MilvusVectorStoreIT {
public MilvusServiceClient milvusClient() {
return new MilvusServiceClient(ConnectParam.newBuilder()
.withAuthorization("minioadmin", "minioadmin")
.withHost(milvusContainer.getServiceHost("standalone", 19530))
.withPort(milvusContainer.getServicePort("standalone", 19530))
.withUri(milvusContainer.getEndpoint())
.build());
}

64
vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml
View File

@@ -1,64 +0,0 @@
version: '3.5'
services:
etcd:
image: quay.io/coreos/etcd:v3.5.5
ports:
- "2379:2379"
environment:
- ETCD_AUTO_COMPACTION_MODE=revision
- ETCD_AUTO_COMPACTION_RETENTION=1000
- ETCD_QUOTA_BACKEND_BYTES=4294967296
- ETCD_SNAPSHOT_COUNT=50000
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
healthcheck:
test: ["CMD", "etcdctl", "endpoint", "health"]
interval: 30s
timeout: 20s
retries: 3
minio:
image: minio/minio:RELEASE.2023-03-20T20-16-18Z
environment:
MINIO_ACCESS_KEY: minioadmin
MINIO_SECRET_KEY: minioadmin
ports:
- "9001:9001"
- "9000:9000"
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
command: minio server /minio_data --console-address ":9001"
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
interval: 30s
timeout: 20s
retries: 3
standalone:
image: milvusdb/milvus:v2.3.8
command: ["milvus", "run", "standalone"]
security_opt:
- seccomp:unconfined
environment:
ETCD_ENDPOINTS: etcd:2379
MINIO_ADDRESS: minio:9000
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
interval: 30s
start_period: 90s
timeout: 20s
retries: 3
ports:
- "19530:19530"
- "9091:9091"
depends_on:
- "etcd"
- "minio"
networks:
default:
name: milvus

2
vector-stores/spring-ai-weaviate/pom.xml
View File

@@ -67,7 +67,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<artifactId>weaviate</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

15
vector-stores/spring-ai-weaviate/src/test/java/org/springframework/ai/vectorstore/WeaviateVectorStoreIT.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2023 the original author or authors.
* 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.
@@ -24,7 +24,6 @@ import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -38,23 +37,19 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.weaviate.WeaviateContainer;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@Testcontainers
public class WeaviateVectorStoreIT {
@Container
static GenericContainer<?> weaviateContainer = new GenericContainer<>("semitechnologies/weaviate:1.22.4")
.withEnv("AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED", "true")
.withEnv("PERSISTENCE_DATA_PATH", "/var/lib/weaviate")
.withEnv("QUERY_DEFAULTS_LIMIT", "25")
.withEnv("DEFAULT_VECTORIZER_MODULE", "none")
.withEnv("CLUSTER_HOSTNAME", "node1")
.withExposedPorts(8080);
static WeaviateContainer weaviateContainer = new WeaviateContainer("semitechnologies/weaviate:1.22.4");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class);
@@ -250,7 +245,7 @@ public class WeaviateVectorStoreIT {
public VectorStore vectorStore(EmbeddingClient embeddingClient) {
WeaviateVectorStoreConfig config = WeaviateVectorStore.WeaviateVectorStoreConfig.builder()
.withScheme("http")
.withHost(String.format("%s:%s", weaviateContainer.getHost(), weaviateContainer.getMappedPort(8080)))
.withHost(weaviateContainer.getHttpHostAddress())
.withFilterableMetadataFields(List.of(MetadataField.text("country"), MetadataField.number("year")))
.withConsistencyLevel(WeaviateVectorStoreConfig.ConsistentLevel.ONE)
.build();