From 141f503341983f94a91c3b51d782e1234bbd26a8 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 19 Sep 2023 14:12:44 +0200 Subject: [PATCH] Integrate Milvus as a Vestor Store Resolves #19 * Clear and stream line. Inject search distance response as metadata * Fix similarity threshold by metric type - Fix the similary search with threshold for L2 and IP types. - Create config Builder Config (inspired by Neo4jVectorStore impelementation. - add database-name, metric-type, index-type and index-param configurations. * update vector store doc and readme --- pom.xml | 1 + .../modules/ROOT/pages/api/vectordbs.adoc | 6 +- .../spring-ai-milvus-store/README.md | 41 ++ vector-stores/spring-ai-milvus-store/pom.xml | 82 +++ .../ai/vectorstore/MilvusVectorStore.java | 526 ++++++++++++++++++ .../ai/vectorstore/MilvusVectorStoreIT.java | 233 ++++++++ .../src/test/resources/docker-compose.yml | 62 +++ 7 files changed, 949 insertions(+), 2 deletions(-) create mode 100644 vector-stores/spring-ai-milvus-store/README.md create mode 100644 vector-stores/spring-ai-milvus-store/pom.xml create mode 100644 vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java create mode 100644 vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java create mode 100644 vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml diff --git a/pom.xml b/pom.xml index 9675749f9..3cef21997 100644 --- a/pom.xml +++ b/pom.xml @@ -20,6 +20,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-azure-openai spring-ai-docs vector-stores/spring-ai-pgvector-store + vector-stores/spring-ai-milvus-store diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc index 9b5abc7f2..5290956d3 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc @@ -36,12 +36,13 @@ public interface VectorStore { List similaritySearch(String query, int k); List similaritySearch(String query, int k, double threshold); +} ``` The Spring AI library uses the `Document` class to model the `String` you want to store in the Vector Database along with the vector representing that `String`. The vector that representing the `String` is of the type `List` and is often called the string's embedding. -The `List is computed from the `String` using an implementation of the class `EmbeddingClient`. +The `List` is computed from the `String` using an implementation of the class `EmbeddingClient`. An Embedding Model is an example of an AI model, that transforms from `String` to a `List`. The embeddings are used internally by other AI Model implementations, such as models that convert `text` to `text` like ChatGPT. @@ -52,7 +53,8 @@ The `VectorStore` implementations supported by Spring AI are: * InMemoryVectorStore * SimplePersistentVectorStore -* PgVector - A Vector Store build on PostgreSQL. +* PgVector - A Vector Store build on https://github.com/pgvector/pgvector[PostgreSQL/PGVector]. +* Milvus - A Vector Store build on https://milvus.io/[Milvus] More are implementations are coming, with Pinecone being the next implementation. diff --git a/vector-stores/spring-ai-milvus-store/README.md b/vector-stores/spring-ai-milvus-store/README.md new file mode 100644 index 000000000..34bf3badf --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/README.md @@ -0,0 +1,41 @@ +# Introduction to Milvus + +[Milvus](https://milvus.io/) is an open-source vector database that has garnered significant attention in the fields of data science and machine learning. +One of its standout features lies in its robust support for vector indexing and querying. +Milvus employs cutting-edge algorithms to accelerate the search process, making it exceptionally efficient at retrieving similar vectors, even when handling extensive datasets. + +Milvus's popularity also comes from its ease of integration with popular Python based frameworks such as PyTorch and TensorFlow, allowing for seamless inclusion in existing machine learning workflows. + +is yet another open source vector database; and this one has gained popularity in the data science and machine learning fields. One of Milvus’ main advantages is its robust support for vector indexing and querying. +It uses state-of-the-art algorithms to speed up the search process, resulting in fast retrieval of similar vectors even when dealing with large-scale datasets. + +Its popularity also stems from the fact that Milvus can be easily integrated with other popular frameworks, including `PyTorch` and `TensorFlow`, enabling seamless integration into existing machine learning workflows. + +In the e-commerce industry, Milvus is used in recommendation systems, which suggest products based on user preferences. +In image and video analysis, it excels in tasks like object recognition, image similarity search, and content-based image retrieval. +Additionally, it is commonly used in natural language processing for document clustering, semantic search, and question-answering systems. + +## Starting Milvus Store + +from withing the `src/test/resources/` folder run: + +``` +docker-compose up +``` + +To clean the environment: + +``` +docker-compose down; rm -Rf ./volumes +``` + + +Then connect to the vector store on http://localhost:19530 or for management http://localhost:9001 (user: `minioadmin`, pass: `minioadmin`) + +## Throubleshooting + +If docker complains about resources: + +``` +docker system prune --all --force --volumes +``` \ No newline at end of file diff --git a/vector-stores/spring-ai-milvus-store/pom.xml b/vector-stores/spring-ai-milvus-store/pom.xml new file mode 100644 index 000000000..5574d5f2a --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/pom.xml @@ -0,0 +1,82 @@ + + + 4.0.0 + + org.springframework.experimental.ai + spring-ai + 0.2.0-SNAPSHOT + ../../pom.xml + + spring-ai-milvus-store + jar + Spring AI Milvus Vector Store + Spring AI Milvus Vector Store + https://github.com/spring-projects-experimental/spring-ai + + + https://github.com/spring-projects-experimental/spring-ai + git://github.com/spring-projects-experimental/spring-ai.git + git@github.com:spring-projects-experimental/spring-ai.git + + + + 2.3.0 + 0.2.0-SNAPSHOT + + 1.19.0 + + + + + + org.springframework.experimental.ai + spring-ai-core + ${spring-ai.version} + + + + + io.milvus + milvus-sdk-java + ${milvus.version} + + + + + org.springframework.experimental.ai + spring-ai-openai-spring-boot-starter + ${spring-ai.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + + org.testcontainers + testcontainers + ${testcontainers.version} + test + + + + org.testcontainers + junit-jupiter + ${testcontainers.version} + test + + + + + + + diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java new file mode 100644 index 000000000..c85289d8a --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java @@ -0,0 +1,526 @@ +/* + * Copyright 2023-2023 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.vectorstore; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; + +import com.alibaba.fastjson.JSONObject; +import io.milvus.client.MilvusServiceClient; +import io.milvus.common.clientenum.ConsistencyLevelEnum; +import io.milvus.grpc.DataType; +import io.milvus.grpc.DescribeIndexResponse; +import io.milvus.grpc.MutationResult; +import io.milvus.grpc.SearchResults; +import io.milvus.param.IndexType; +import io.milvus.param.MetricType; +import io.milvus.param.R; +import io.milvus.param.R.Status; +import io.milvus.param.RpcStatus; +import io.milvus.param.collection.CreateCollectionParam; +import io.milvus.param.collection.DropCollectionParam; +import io.milvus.param.collection.FieldType; +import io.milvus.param.collection.FlushParam; +import io.milvus.param.collection.HasCollectionParam; +import io.milvus.param.collection.LoadCollectionParam; +import io.milvus.param.collection.ReleaseCollectionParam; +import io.milvus.param.dml.DeleteParam; +import io.milvus.param.dml.InsertParam; +import io.milvus.param.dml.SearchParam; +import io.milvus.param.index.CreateIndexParam; +import io.milvus.param.index.DescribeIndexParam; +import io.milvus.param.index.DropIndexParam; +import io.milvus.response.SearchResultsWrapper; +import io.milvus.response.QueryResultsWrapper.RowRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.context.SmartLifecycle; +import org.springframework.util.Assert; + +/** + * @author Christian Tzolov + */ +public class MilvusVectorStore implements VectorStore, SmartLifecycle { + + private static final Logger logger = LoggerFactory.getLogger(MilvusVectorStore.class); + + public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536; + + public static final String DEFAULT_DATABASE_NAME = "default"; + + public static final String DEFAULT_COLLECTION_NAME = "vector_store"; + + public static final String DOC_ID_FIELD_NAME = "doc_id"; + + public static final String CONTENT_FIELD_NAME = "content"; + + public static final String METADATA_FIELD_NAME = "metadata"; + + public static final String EMBEDDING_FIELD_NAME = "embedding"; + + // Metadata, automatically assigned by Milvus. + public static final String DISTANCE_FIELD_NAME = "distance"; + + public static final List SEARCH_OUTPUT_FIELDS = Arrays.asList(DOC_ID_FIELD_NAME, CONTENT_FIELD_NAME, + METADATA_FIELD_NAME); + + private final MilvusServiceClient milvusClient; + + private final EmbeddingClient embeddingClient; + + private final MilvusVectorStoreConfig config; + + /** + * Configuration for the Milvus vector store. + */ + public static class MilvusVectorStoreConfig { + + private final String databaseName; + + private final String collectionName; + + private final int embeddingDimension; + + private final IndexType indexType; + + private final MetricType metricType; + + private final String indexParameters; + + /** + * Start building a new configuration. + * @return The entry point for creating a new configuration. + */ + public static Builder builder() { + + return new Builder(); + } + + /** + * {@return the default config} + */ + public static MilvusVectorStoreConfig defaultConfig() { + + return builder().build(); + } + + private MilvusVectorStoreConfig(Builder builder) { + this.databaseName = builder.databaseName; + this.collectionName = builder.collectionName; + this.embeddingDimension = builder.embeddingDimension; + this.indexType = builder.indexType; + this.metricType = builder.metricType; + this.indexParameters = builder.indexParameters; + } + + public static class Builder { + + private String databaseName = DEFAULT_DATABASE_NAME; + + private String collectionName = DEFAULT_COLLECTION_NAME; + + private int embeddingDimension = OPENAI_EMBEDDING_DIMENSION_SIZE; + + private IndexType indexType = IndexType.IVF_FLAT; + + private MetricType metricType = MetricType.L2; + + private String indexParameters = "{\"nlist\":1024}"; + + private Builder() { + } + + /** + * Configures the Milvus metric type to use. Leave {@literal null} or blank to + * use the metric metric. + * @param metricType the metric type to use + * @return this builder + */ + public Builder withMetricType(MetricType metricType) { + Assert.notNull(metricType, "Collection Name must not be empty"); + Assert.isTrue(metricType == MetricType.IP || metricType == MetricType.L2, + "Only the text metric types IP and L2 are supported"); + + this.metricType = metricType; + return this; + } + + /** + * Configures the Milvus index type to use. Leave {@literal null} or blank to + * use the default index. + * @param indexType the index type to use + * @return this builder + */ + public Builder withIndexType(IndexType indexType) { + this.indexType = indexType; + return this; + } + + /** + * Configures the Milvus index parameters to use. Leave {@literal null} or + * blank to use the default index parameters. + * @param indexParameters the index parameters to use + * @return this builder + */ + public Builder withIndexParameters(String indexParameters) { + this.indexParameters = indexParameters; + return this; + } + + /** + * Configures the Milvus database name to use. Leave {@literal null} or blank + * to use the default database. + * @param databaseName the database name to use + * @return this builder + */ + public Builder withDatabaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + /** + * Configures the Milvus collection name to use. Leave {@literal null} or + * blank to use the default collection name. + * @param collectionName the collection name to use + * @return this builder + */ + public Builder withCollectionName(String collectionName) { + this.collectionName = collectionName; + return this; + } + + /** + * Configures the size of the embedding. Defaults to {@literal 1536}, inline + * with OpenAIs embeddings. + * @param newEmbeddingDimension The dimension of the embedding + * @return this builder + */ + public Builder withEmbeddingDimension(int newEmbeddingDimension) { + + Assert.isTrue(newEmbeddingDimension >= 1 && newEmbeddingDimension <= 2048, + "Dimension has to be withing the boundaries 1 and 2048 (inclusively)"); + + this.embeddingDimension = newEmbeddingDimension; + return this; + } + + /** + * {@return the immutable configuration} + */ + public MilvusVectorStoreConfig build() { + return new MilvusVectorStoreConfig(this); + } + + } + + } + + public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingClient embeddingClient) { + this(milvusClient, embeddingClient, MilvusVectorStoreConfig.defaultConfig()); + } + + public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingClient embeddingClient, + MilvusVectorStoreConfig config) { + + Assert.notNull(milvusClient, "MilvusServiceClient must not be null"); + Assert.notNull(milvusClient, "EmbeddingClient must not be null"); + + this.milvusClient = milvusClient; + this.embeddingClient = embeddingClient; + this.config = config; + } + + @Override + public void add(List documents) { + + Assert.notNull(documents, "Documents must not be null"); + + List docIdArray = new ArrayList<>(); + List contentArray = new ArrayList<>(); + List metadataArray = new ArrayList<>(); + List> embeddingArray = new ArrayList<>(); + + for (Document document : documents) { + List embedding = this.embeddingClient.embed(document); + + docIdArray.add(document.getId()); + // Use a (future) DocumentTextLayoutFormatter instance to extract + // the content used to compute the embeddings + contentArray.add(document.getText()); + metadataArray.add(new JSONObject(document.getMetadata())); + embeddingArray.add(toFloatList(embedding)); + } + + List fields = new ArrayList<>(); + fields.add(new InsertParam.Field(DOC_ID_FIELD_NAME, docIdArray)); + fields.add(new InsertParam.Field(CONTENT_FIELD_NAME, contentArray)); + fields.add(new InsertParam.Field(METADATA_FIELD_NAME, metadataArray)); + fields.add(new InsertParam.Field(EMBEDDING_FIELD_NAME, embeddingArray)); + + InsertParam insertParam = InsertParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .withFields(fields) + .build(); + + R status = this.milvusClient.insert(insertParam); + if (status.getException() != null) { + throw new RuntimeException("Failed to insert:", status.getException()); + } + this.milvusClient.flush(FlushParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .addCollectionName(this.config.collectionName) + .build()); + } + + @Override + public Optional delete(List idList) { + Assert.notNull(idList, "Document id list must not be null"); + + String deleteExpression = String.format("%s in [%s]", DOC_ID_FIELD_NAME, + idList.stream().map(id -> "'" + id + "'").collect(Collectors.joining(","))); + + R status = this.milvusClient.delete(DeleteParam.newBuilder() + .withCollectionName(this.config.collectionName) + .withExpr(deleteExpression) + .build()); + + long deleteCount = status.getData().getDeleteCnt(); + if (deleteCount != idList.size()) { + logger.warn(String.format("Deleted only %s entries from requested %s ", deleteCount, idList.size())); + } + + return Optional.of(status.getStatus() == Status.Success.getCode()); + } + + @Override + public List similaritySearch(String query) { + return this.similaritySearch(query, 4); + } + + @Override + public List similaritySearch(String query, int topK) { + return similaritySearch(query, topK, 0.0D); + } + + @Override + public List similaritySearch(String query, int topK, double similarityThreshold) { + Assert.notNull(query, "Query string must not be null"); + + List embedding = this.embeddingClient.embed(query); + + SearchParam searchParam = SearchParam.newBuilder() + .withCollectionName(this.config.collectionName) + .withConsistencyLevel(ConsistencyLevelEnum.STRONG) + .withMetricType(this.config.metricType) + .withOutFields(SEARCH_OUTPUT_FIELDS) + .withTopK(topK) + .withVectors(List.of(toFloatList(embedding))) + .withVectorFieldName(EMBEDDING_FIELD_NAME) + .build(); + + R respSearch = milvusClient.search(searchParam); + + if (respSearch.getException() != null) { + throw new RuntimeException("Search failed!", respSearch.getException()); + } + + SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(respSearch.getData().getResults()); + + return wrapperSearch.getRowRecords() + .stream() + .filter(rowRecord -> getResultSimilarity(rowRecord) >= similarityThreshold) + .map(rowRecord -> { + String docId = (String) rowRecord.get(DOC_ID_FIELD_NAME); + String content = (String) rowRecord.get(CONTENT_FIELD_NAME); + JSONObject metadata = (JSONObject) rowRecord.get(METADATA_FIELD_NAME); + // inject the distance into the metadata. + metadata.put(DISTANCE_FIELD_NAME, 1 - getResultSimilarity(rowRecord)); + return new Document(docId, content, metadata.getInnerMap()); + }) + .collect(Collectors.toList()); + } + + private float getResultSimilarity(RowRecord rowRecord) { + Float distance = (Float) rowRecord.get(DISTANCE_FIELD_NAME); + return (this.config.metricType == MetricType.IP) ? distance : (1 - distance); + } + + private List toFloatList(List embeddingDouble) { + return embeddingDouble.stream().map(Number::floatValue).collect(Collectors.toList()); + } + + // --------------------------------------------------------------------------------- + // SmartLifecycle + // --------------------------------------------------------------------------------- + private AtomicBoolean isRunning = new AtomicBoolean(false); + + @Override + public void start() { + try { + createCollection(); + } + finally { + this.isRunning.set(true); + } + } + + @Override + public void stop() { + try { + if (isDatabaseCollectionExists()) { + this.milvusClient.releaseCollection( + ReleaseCollectionParam.newBuilder().withCollectionName(this.config.collectionName).build()); + } + } + finally { + this.isRunning.set(false); + } + } + + @Override + public boolean isRunning() { + return this.isRunning.get(); + } + + @Override + public boolean isAutoStartup() { + return true; + } + + private boolean isDatabaseCollectionExists() { + return this.milvusClient + .hasCollection(HasCollectionParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .build()) + .getData(); + } + + // used by the test as well + void createCollection() { + + if (!isDatabaseCollectionExists()) { + + FieldType docIdFieldType = FieldType.newBuilder() + .withName(DOC_ID_FIELD_NAME) + .withDataType(DataType.VarChar) + .withMaxLength(36) + .withPrimaryKey(true) + .withAutoID(false) + .build(); + FieldType contentFieldType = FieldType.newBuilder() + .withName(CONTENT_FIELD_NAME) + .withDataType(DataType.VarChar) + .withMaxLength(65535) + .build(); + FieldType metadataFieldType = FieldType.newBuilder() + .withName(METADATA_FIELD_NAME) + .withDataType(DataType.JSON) + .build(); + FieldType embeddingFieldType = FieldType.newBuilder() + .withName(EMBEDDING_FIELD_NAME) + .withDataType(DataType.FloatVector) + .withDimension(this.config.embeddingDimension) + .build(); + + CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + // .withDatabaseName(this.collectionName) + .withDescription("Spring AI Vector Store") + .withConsistencyLevel(ConsistencyLevelEnum.STRONG) + .withShardsNum(2) + .addFieldType(docIdFieldType) + .addFieldType(contentFieldType) + .addFieldType(metadataFieldType) + .addFieldType(embeddingFieldType) + .build(); + + R collectionStatus = this.milvusClient.createCollection(createCollectionReq); + if (collectionStatus.getException() != null) { + throw new RuntimeException("Failed to create collection", collectionStatus.getException()); + } + } + + R indexDescriptionResponse = this.milvusClient + .describeIndex(DescribeIndexParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .build()); + + if (indexDescriptionResponse.getData() == null) { + R indexStatus = this.milvusClient.createIndex(CreateIndexParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .withFieldName(EMBEDDING_FIELD_NAME) + .withIndexType(this.config.indexType) + .withMetricType(this.config.metricType) + .withExtraParam(this.config.indexParameters) + .withSyncMode(Boolean.FALSE) + .build()); + + if (indexStatus.getException() != null) { + throw new RuntimeException("Failed to create Index", indexStatus.getException()); + } + } + + R loadCollectionStatus = this.milvusClient.loadCollection(LoadCollectionParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .build()); + + if (loadCollectionStatus.getException() != null) { + throw new RuntimeException("Collection loading failed!", loadCollectionStatus.getException()); + } + } + + // used by the test as well + void dropCollection() { + + R status = this.milvusClient.releaseCollection( + ReleaseCollectionParam.newBuilder().withCollectionName(this.config.collectionName).build()); + + if (status.getException() != null) { + throw new RuntimeException("Release collection failed!", status.getException()); + } + + status = this.milvusClient + .dropIndex(DropIndexParam.newBuilder().withCollectionName(this.config.collectionName).build()); + + if (status.getException() != null) { + throw new RuntimeException("Drop Index failed!", status.getException()); + } + + status = this.milvusClient.dropCollection(DropCollectionParam.newBuilder() + .withDatabaseName(this.config.databaseName) + .withCollectionName(this.config.collectionName) + .build()); + + if (status.getException() != null) { + throw new RuntimeException("Drop Collection failed!", status.getException()); + } + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java new file mode 100644 index 000000000..9b33b7012 --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java @@ -0,0 +1,233 @@ +/* + * Copyright 2023-2023 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.vectorstore; + +import java.io.File; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +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.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.autoconfigure.openai.OpenAiAutoConfiguration; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.util.FileSystemUtils; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +@Testcontainers +public class MilvusVectorStoreIT { + + @Container + public static DockerComposeContainer milvusContainer = new DockerComposeContainer( + new File("src/test/resources/docker-compose.yml")) + .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))); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class) + .withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY")); + + List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")), + new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"), + new Document( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression", + Collections.singletonMap("meta2", "meta2"))); + + @AfterAll + public static void afterAll() { + FileSystemUtils.deleteRecursively(new File("src/test/resources/volumes")); + } + + private void resetCollection(VectorStore vectorStore) { + ((MilvusVectorStore) vectorStore).dropCollection(); + ((MilvusVectorStore) vectorStore).createCollection(); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "L2", "IP" }) + public void addAndSearchTest(String metricType) { + + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)) + .withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType) + .run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + vectorStore.add(documents); + + List results = vectorStore.similaritySearch("Great", 1); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getText()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).hasSize(2); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(doc -> doc.getId()).collect(Collectors.toList())); + + List results2 = vectorStore.similaritySearch("Hello", 1); + assertThat(results2).hasSize(0); + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "L2", "IP" }) + public void documentUpdateTest(String metricType) { + + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)) + .withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType) + .run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + + List results = vectorStore.similaritySearch("Spring", 5); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsKey("meta1"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + results = vectorStore.similaritySearch("FooBar", 5); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getText()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + vectorStore.delete(List.of(document.getId())); + + }); + } + + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "L2", "IP" }) + public void searchThresholdTest(String metricType) { + + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)) + .withPropertyValues("spring.ai.vectorstore.milvus.metricType=" + metricType) + .run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + resetCollection(vectorStore); + + vectorStore.add(documents); + + List fullResult = vectorStore.similaritySearch("Great", 5, 0.0); + + List distances = fullResult.stream() + .map(doc -> (Float) doc.getMetadata().get("distance")) + .collect(Collectors.toList()); + + assertThat(distances).hasSize(3); + + List results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.01))); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getText()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).containsKey("meta2"); + assertThat(resultDoc.getMetadata()).containsKey("distance"); + + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) + public static class TestApplication { + + @Value("${spring.ai.vectorstore.milvus.metricType}") + private MetricType metricType; + + @Bean + public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingClient embeddingClient) { + MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() + .withCollectionName("test_vector_store") + .withDatabaseName("default") + .withIndexType(IndexType.IVF_FLAT) + .withMetricType(metricType) + .withEmbeddingDimension(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE) + .build(); + return new MilvusVectorStore(milvusClient, embeddingClient, config); + } + + @Bean + public MilvusServiceClient milvusClient() { + return new MilvusServiceClient(ConnectParam.newBuilder() + .withHost("localhost") + .withPort(milvusContainer.getServicePort("standalone", 19530)) + .build()); + } + + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml b/vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml new file mode 100644 index 000000000..aec1e5238 --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/src/test/resources/docker-compose.yml @@ -0,0 +1,62 @@ +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-09-23T03-47-50Z + 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.1 + command: ["milvus", "run", "standalone"] + 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