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

@@ -20,6 +20,7 @@
 		<module>spring-ai-spring-boot-starters/spring-ai-starter-azure-openai</module>
 		<module>spring-ai-docs</module>
 		<module>vector-stores/spring-ai-pgvector-store</module>
+		<module>vector-stores/spring-ai-milvus-store</module>
 	</modules>
 
     <organization>

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc

@@ -36,12 +36,13 @@ public interface VectorStore {
 	List<Document> similaritySearch(String query, int k);
 
 	List<Document> 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<Double>` and is often called the string's embedding.
 
-The `List<Double> is computed from the `String` using an implementation of the class `EmbeddingClient`.
+The `List<Double>` 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<Double>`.
 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.
 

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
+```

vector-stores/spring-ai-milvus-store/pom.xml

@@ -0,0 +1,82 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+	<modelVersion>4.0.0</modelVersion>
+	<parent>
+		<groupId>org.springframework.experimental.ai</groupId>
+		<artifactId>spring-ai</artifactId>
+		<version>0.2.0-SNAPSHOT</version>
+		<relativePath>../../pom.xml</relativePath>
+	</parent>
+	<artifactId>spring-ai-milvus-store</artifactId>
+	<packaging>jar</packaging>
+	<name>Spring AI Milvus Vector Store</name>
+	<description>Spring AI Milvus Vector Store</description>
+	<url>https://github.com/spring-projects-experimental/spring-ai</url>
+
+	<scm>
+		<url>https://github.com/spring-projects-experimental/spring-ai</url>
+		<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
+		<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
+	</scm>
+
+	<properties>
+		<milvus.version>2.3.0</milvus.version>
+		<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
+		<!-- testing -->
+		<testcontainers.version>1.19.0</testcontainers.version>
+		<!-- <hikari-cp.version>4.0.3</hikari-cp.version> -->
+	</properties>
+
+	<dependencies>
+		<dependency>
+			<groupId>org.springframework.experimental.ai</groupId>
+			<artifactId>spring-ai-core</artifactId>
+			<version>${spring-ai.version}</version>
+		</dependency>
+
+
+		<dependency>
+			<groupId>io.milvus</groupId>
+			<artifactId>milvus-sdk-java</artifactId>
+			<version>${milvus.version}</version>
+		</dependency>
+
+		<!-- TESTING -->
+		<dependency>
+			<groupId>org.springframework.experimental.ai</groupId>
+			<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
+			<version>${spring-ai.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-starter-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>testcontainers</artifactId>
+			<version>${testcontainers.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>junit-jupiter</artifactId>
+			<version>${testcontainers.version}</version>
+			<scope>test</scope>
+		</dependency>
+
+		<!-- <dependency>
+			<groupId>com.zaxxer</groupId>
+			<artifactId>HikariCP</artifactId>
+			<version>${hikari-cp.version}</version>
+			<scope>test</scope>
+		</dependency> -->
+
+	</dependencies>
+
+</project>

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<String> 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<Document> documents) {
+
+		Assert.notNull(documents, "Documents must not be null");
+
+		List<String> docIdArray = new ArrayList<>();
+		List<String> contentArray = new ArrayList<>();
+		List<JSONObject> metadataArray = new ArrayList<>();
+		List<List<Float>> embeddingArray = new ArrayList<>();
+
+		for (Document document : documents) {
+			List<Double> 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<InsertParam.Field> 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<MutationResult> 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<Boolean> delete(List<String> 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<MutationResult> 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<Document> similaritySearch(String query) {
+		return this.similaritySearch(query, 4);
+	}
+
+	@Override
+	public List<Document> similaritySearch(String query, int topK) {
+		return similaritySearch(query, topK, 0.0D);
+	}
+
+	@Override
+	public List<Document> similaritySearch(String query, int topK, double similarityThreshold) {
+		Assert.notNull(query, "Query string must not be null");
+
+		List<Double> 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<SearchResults> 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<Float> toFloatList(List<Double> 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<RpcStatus> collectionStatus = this.milvusClient.createCollection(createCollectionReq);
+			if (collectionStatus.getException() != null) {
+				throw new RuntimeException("Failed to create collection", collectionStatus.getException());
+			}
+		}
+
+		R<DescribeIndexResponse> indexDescriptionResponse = this.milvusClient
+			.describeIndex(DescribeIndexParam.newBuilder()
+				.withDatabaseName(this.config.databaseName)
+				.withCollectionName(this.config.collectionName)
+				.build());
+
+		if (indexDescriptionResponse.getData() == null) {
+			R<RpcStatus> 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<RpcStatus> 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<RpcStatus> 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());
+		}
+	}
+
+}

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<Document> 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<Document> 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<Document> 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<Document> 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<Document> fullResult = vectorStore.similaritySearch("Great", 5, 0.0);
+
+				List<Float> distances = fullResult.stream()
+					.map(doc -> (Float) doc.getMetadata().get("distance"))
+					.collect(Collectors.toList());
+
+				assertThat(distances).hasSize(3);
+
+				List<Document> 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());
+		}
+
+	}
+
+}

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