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