diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java index fde5c380b..1dc3eb5df 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java @@ -27,8 +27,8 @@ import io.milvus.param.MetricType; import org.springframework.ai.embedding.BatchingStrategy; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.MilvusVectorStore; -import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.milvus.vectorstore.MilvusVectorStore; +import org.springframework.ai.milvus.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -71,23 +71,14 @@ public class MilvusVectorStoreAutoConfiguration { ObjectProvider observationRegistry, ObjectProvider customObservationConvention) { - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withCollectionName(properties.getCollectionName()) - .withDatabaseName(properties.getDatabaseName()) - .withIndexType(IndexType.valueOf(properties.getIndexType().name())) - .withMetricType(MetricType.valueOf(properties.getMetricType().name())) - .withIndexParameters(properties.getIndexParameters()) - .withEmbeddingDimension(properties.getEmbeddingDimension()) - .withIDFieldName(properties.getIdFieldName()) - .withAutoId(properties.isAutoId()) - .withContentFieldName(properties.getContentFieldName()) - .withMetadataFieldName(properties.getMetadataFieldName()) - .withEmbeddingFieldName(properties.getEmbeddingFieldName()) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) + .embeddingModel(embeddingModel) + .initializeSchema(properties.isInitializeSchema()) + .batchingStrategy(batchingStrategy) + .observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)) + .customObservationConvention(customObservationConvention.getIfAvailable(() -> null)) .build(); - - return new MilvusVectorStore(milvusClient, embeddingModel, config, properties.isInitializeSchema(), - batchingStrategy, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP), - customObservationConvention.getIfAvailable(() -> null)); } @Bean diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java index 9e8457caa..e498b6273 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java @@ -17,7 +17,7 @@ package org.springframework.ai.autoconfigure.vectorstore.milvus; import org.springframework.ai.autoconfigure.vectorstore.CommonVectorStoreProperties; -import org.springframework.ai.vectorstore.MilvusVectorStore; +import org.springframework.ai.milvus.vectorstore.MilvusVectorStore; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.util.Assert; diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverter.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverter.java similarity index 97% rename from vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverter.java rename to vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverter.java index d8ded6c2e..8e0f4b2c4 100644 --- a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverter.java +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverter.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import org.springframework.ai.vectorstore.filter.Filter.Expression; import org.springframework.ai.vectorstore.filter.Filter.ExpressionType; 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/milvus/vectorstore/MilvusVectorStore.java similarity index 58% rename from vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java rename to vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java index 57704b420..3b6d1d08a 100644 --- 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/milvus/vectorstore/MilvusVectorStore.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import java.util.ArrayList; import java.util.List; @@ -62,6 +62,8 @@ import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.model.EmbeddingUtils; import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; +import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder; +import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext; @@ -71,13 +73,78 @@ import org.springframework.util.Assert; import org.springframework.util.StringUtils; /** - * Milvus based implementation of the - * {@link org.springframework.ai.vectorstore.VectorStore}. + * Milvus implementation of the {@link org.springframework.ai.vectorstore.VectorStore} + * interface. This implementation supports storing and searching document embeddings using + * Milvus, an open-source vector database optimized for similarity search and AI + * applications. + * + *

+ * Key features include: + *

+ * + *

+ * Example usage: 

{@code
+ * // Create a basic Milvus vector store
+ * MilvusVectorStore vectorStore = MilvusVectorStore.builder()
+ *     .milvusClient(milvusClient)
+ *     .embeddingModel(embeddingModel)
+ *     .initializeSchema(true)
+ *     .build();
+ *
+ * // Create a customized Milvus vector store
+ * MilvusVectorStore customVectorStore = MilvusVectorStore.builder()
+ *     .milvusClient(milvusClient)
+ *     .embeddingModel(embeddingModel)
+ *     .databaseName("my_database")
+ *     .collectionName("my_collection")
+ *     .metricType(MetricType.COSINE)
+ *     .indexType(IndexType.IVF_FLAT)
+ *     .indexParameters("{\"nlist\":1024}")
+ *     .embeddingDimension(1536)
+ *     .batchingStrategy(new TokenCountBatchingStrategy())
+ *     .initializeSchema(true)
+ *     .build();
+ *
+ * // Add documents to the store
+ * List documents = List.of(
+ *     new Document("content1", Map.of("meta1", "value1")),
+ *     new Document("content2", Map.of("meta2", "value2"))
+ * );
+ * vectorStore.add(documents);
+ *
+ * // Perform similarity search
+ * List results = vectorStore.similaritySearch(
+ *     SearchRequest.query("search text")
+ *         .withTopK(5)
+ *         .withSimilarityThreshold(0.7)
+ *         .withFilterExpression("meta1 == 'value1'")
+ * );
+ * }
+ * + *

+ * The vector store supports various configuration options through its builder: + *

* * @author Christian Tzolov * @author Soby Chacko * @author Thomas Vitale * @author Ilayaperumal Gopinathan + * @see org.springframework.ai.vectorstore.VectorStore + * @see io.milvus.client.MilvusServiceClient */ public class MilvusVectorStore extends AbstractObservationVectorStore implements InitializingBean { @@ -110,44 +177,99 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements private final MilvusServiceClient milvusClient; - private final EmbeddingModel embeddingModel; - + @Deprecated(forRemoval = true, since = "1.0.0-M5") private final MilvusVectorStoreConfig config; private final boolean initializeSchema; private final BatchingStrategy batchingStrategy; + private final String databaseName; + + private final String collectionName; + + private final int embeddingDimension; + + private final IndexType indexType; + + private final MetricType metricType; + + private final String indexParameters; + + private final String idFieldName; + + private final boolean isAutoId; + + private final String contentFieldName; + + private final String metadataFieldName; + + private final String embeddingFieldName; + + @Deprecated(forRemoval = true, since = "1.0.0-M5") public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, boolean initializeSchema) { this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema, new TokenCountBatchingStrategy()); } + @Deprecated(forRemoval = true, since = "1.0.0-M5") public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, boolean initializeSchema, BatchingStrategy batchingStrategy) { this(milvusClient, embeddingModel, MilvusVectorStoreConfig.defaultConfig(), initializeSchema, batchingStrategy); } + @Deprecated(forRemoval = true, since = "1.0.0-M5") public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy) { this(milvusClient, embeddingModel, config, initializeSchema, batchingStrategy, ObservationRegistry.NOOP, null); } + @Deprecated(forRemoval = true, since = "1.0.0-M5") public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention) { - super(observationRegistry, customObservationConvention); - this.initializeSchema = initializeSchema; + this(builder().milvusClient(milvusClient) + .embeddingModel(embeddingModel) + .observationRegistry(observationRegistry) + .customObservationConvention(customObservationConvention) + .initializeSchema(initializeSchema) + .batchingStrategy(batchingStrategy)); + } - Assert.notNull(milvusClient, "MilvusServiceClient must not be null"); - Assert.notNull(milvusClient, "EmbeddingModel must not be null"); + /** + * @param builder {@link Builder} for chroma vector store + */ + protected MilvusVectorStore(MilvusBuilder builder) { + super(builder); - this.milvusClient = milvusClient; - this.embeddingModel = embeddingModel; - this.config = config; - this.batchingStrategy = batchingStrategy; + Assert.notNull(builder.milvusClient, "milvusClient must not be null"); + + this.milvusClient = builder.milvusClient; + this.batchingStrategy = builder.batchingStrategy; + this.initializeSchema = builder.initializeSchema; + this.config = null; + this.databaseName = builder.databaseName; + this.collectionName = builder.collectionName; + this.embeddingDimension = builder.embeddingDimension; + this.indexType = builder.indexType; + this.metricType = builder.metricType; + this.indexParameters = builder.indexParameters; + this.idFieldName = builder.idFieldName; + this.isAutoId = builder.isAutoId; + this.contentFieldName = builder.contentFieldName; + this.metadataFieldName = builder.metadataFieldName; + this.embeddingFieldName = builder.embeddingFieldName; + } + + /** + * Creates a new MilvusBuilder instance with the specified Milvus client. This is the + * recommended way to instantiate a MilvusBuilder. + * @return a new MilvusBuilder instance + */ + public static MilvusBuilder builder() { + return new MilvusBuilder(); } @Override @@ -175,16 +297,16 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements List fields = new ArrayList<>(); // Insert ID field only if it is not auto ID - if (!this.config.isAutoId) { - fields.add(new InsertParam.Field(this.config.idFieldName, docIdArray)); + if (!this.isAutoId) { + fields.add(new InsertParam.Field(this.idFieldName, docIdArray)); } - fields.add(new InsertParam.Field(this.config.contentFieldName, contentArray)); - fields.add(new InsertParam.Field(this.config.metadataFieldName, metadataArray)); - fields.add(new InsertParam.Field(this.config.embeddingFieldName, embeddingArray)); + fields.add(new InsertParam.Field(this.contentFieldName, contentArray)); + fields.add(new InsertParam.Field(this.metadataFieldName, metadataArray)); + fields.add(new InsertParam.Field(this.embeddingFieldName, embeddingArray)); InsertParam insertParam = InsertParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .withFields(fields) .build(); @@ -198,12 +320,12 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements public Optional doDelete(List idList) { Assert.notNull(idList, "Document id list must not be null"); - String deleteExpression = String.format("%s in [%s]", this.config.idFieldName, + String deleteExpression = String.format("%s in [%s]", this.idFieldName, idList.stream().map(id -> "'" + id + "'").collect(Collectors.joining(","))); R status = this.milvusClient.delete(DeleteParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .withExpr(deleteExpression) .build()); @@ -223,20 +345,20 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements Assert.notNull(request.getQuery(), "Query string must not be null"); List outFieldNames = new ArrayList<>(); - outFieldNames.add(this.config.idFieldName); - outFieldNames.add(this.config.contentFieldName); - outFieldNames.add(this.config.metadataFieldName); + outFieldNames.add(this.idFieldName); + outFieldNames.add(this.contentFieldName); + outFieldNames.add(this.metadataFieldName); float[] embedding = this.embeddingModel.embed(request.getQuery()); var searchParamBuilder = SearchParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .withConsistencyLevel(ConsistencyLevelEnum.STRONG) - .withMetricType(this.config.metricType) + .withMetricType(this.metricType) .withOutFields(outFieldNames) .withTopK(request.getTopK()) .withVectors(List.of(EmbeddingUtils.toList(embedding))) - .withVectorFieldName(this.config.embeddingFieldName); + .withVectorFieldName(this.embeddingFieldName); if (StringUtils.hasText(nativeFilterExpressions)) { searchParamBuilder.withExpr(nativeFilterExpressions); @@ -254,11 +376,11 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements .stream() .filter(rowRecord -> getResultSimilarity(rowRecord) >= request.getSimilarityThreshold()) .map(rowRecord -> { - String docId = String.valueOf(rowRecord.get(this.config.idFieldName)); - String content = (String) rowRecord.get(this.config.contentFieldName); + String docId = String.valueOf(rowRecord.get(this.idFieldName)); + String content = (String) rowRecord.get(this.contentFieldName); JSONObject metadata = null; try { - metadata = (JSONObject) rowRecord.get(this.config.metadataFieldName); + metadata = (JSONObject) rowRecord.get(this.metadataFieldName); // inject the distance into the metadata. metadata.put(DocumentMetadata.DISTANCE.value(), 1 - getResultSimilarity(rowRecord)); } @@ -278,8 +400,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements private float getResultSimilarity(RowRecord rowRecord) { Float distance = (Float) rowRecord.get(DISTANCE_FIELD_NAME); - return (this.config.metricType == MetricType.IP || this.config.metricType == MetricType.COSINE) ? distance - : (1 - distance); + return (this.metricType == MetricType.IP || this.metricType == MetricType.COSINE) ? distance : (1 - distance); } // --------------------------------------------------------------------------------- @@ -297,16 +418,16 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements void releaseCollection() { if (isDatabaseCollectionExists()) { - this.milvusClient.releaseCollection( - ReleaseCollectionParam.newBuilder().withCollectionName(this.config.collectionName).build()); + this.milvusClient + .releaseCollection(ReleaseCollectionParam.newBuilder().withCollectionName(this.collectionName).build()); } } private boolean isDatabaseCollectionExists() { return this.milvusClient .hasCollection(HasCollectionParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .build()) .getData(); } @@ -315,25 +436,24 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements void createCollection() { if (!isDatabaseCollectionExists()) { - createCollection(this.config.databaseName, this.config.collectionName, this.config.idFieldName, - this.config.isAutoId, this.config.contentFieldName, this.config.metadataFieldName, - this.config.embeddingFieldName); + createCollection(this.databaseName, this.collectionName, this.idFieldName, this.isAutoId, + this.contentFieldName, this.metadataFieldName, this.embeddingFieldName); } R indexDescriptionResponse = this.milvusClient .describeIndex(DescribeIndexParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .build()); if (indexDescriptionResponse.getData() == null) { R indexStatus = this.milvusClient.createIndex(CreateIndexParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) - .withFieldName(this.config.embeddingFieldName) - .withIndexType(this.config.indexType) - .withMetricType(this.config.metricType) - .withExtraParam(this.config.indexParameters) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) + .withFieldName(this.embeddingFieldName) + .withIndexType(this.indexType) + .withMetricType(this.metricType) + .withExtraParam(this.indexParameters) .withSyncMode(Boolean.FALSE) .build()); @@ -343,8 +463,8 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements } R loadCollectionStatus = this.milvusClient.loadCollection(LoadCollectionParam.newBuilder() - .withDatabaseName(this.config.databaseName) - .withCollectionName(this.config.collectionName) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .build()); if (loadCollectionStatus.getException() != null) { @@ -396,8 +516,8 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements } int embeddingDimensions() { - if (this.config.embeddingDimension != INVALID_EMBEDDING_DIMENSION) { - return this.config.embeddingDimension; + if (this.embeddingDimension != INVALID_EMBEDDING_DIMENSION) { + return this.embeddingDimension; } try { int embeddingDimensions = this.embeddingModel.dimensions(); @@ -407,7 +527,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements } catch (Exception e) { logger.warn("Failed to obtain the embedding dimensions from the embedding model and fall backs to default:" - + this.config.embeddingDimension, e); + + this.embeddingDimension, e); } return OPENAI_EMBEDDING_DIMENSION_SIZE; } @@ -415,23 +535,23 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements // used by the test as well void dropCollection() { - R status = this.milvusClient.releaseCollection( - ReleaseCollectionParam.newBuilder().withCollectionName(this.config.collectionName).build()); + R status = this.milvusClient + .releaseCollection(ReleaseCollectionParam.newBuilder().withCollectionName(this.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()); + .dropIndex(DropIndexParam.newBuilder().withCollectionName(this.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) + .withDatabaseName(this.databaseName) + .withCollectionName(this.collectionName) .build()); if (status.getException() != null) { @@ -444,22 +564,222 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements String operationName) { return VectorStoreObservationContext.builder(VectorStoreProvider.MILVUS.value(), operationName) - .withCollectionName(this.config.collectionName) + .withCollectionName(this.collectionName) .withDimensions(this.embeddingModel.dimensions()) .withSimilarityMetric(getSimilarityMetric()) - .withNamespace(this.config.databaseName); + .withNamespace(this.databaseName); } private String getSimilarityMetric() { - if (!SIMILARITY_TYPE_MAPPING.containsKey(this.config.metricType)) { - return this.config.metricType.name(); + if (!SIMILARITY_TYPE_MAPPING.containsKey(this.metricType)) { + return this.metricType.name(); } - return SIMILARITY_TYPE_MAPPING.get(this.config.metricType).value(); + return SIMILARITY_TYPE_MAPPING.get(this.metricType).value(); } - /** - * Configuration for the Milvus vector store. - */ + public static final class MilvusBuilder extends AbstractVectorStoreBuilder { + + private String databaseName = DEFAULT_DATABASE_NAME; + + private String collectionName = DEFAULT_COLLECTION_NAME; + + private int embeddingDimension = INVALID_EMBEDDING_DIMENSION; + + private IndexType indexType = IndexType.IVF_FLAT; + + private MetricType metricType = MetricType.COSINE; + + private String indexParameters = "{\"nlist\":1024}"; + + private String idFieldName = DOC_ID_FIELD_NAME; + + private boolean isAutoId = false; + + private String contentFieldName = CONTENT_FIELD_NAME; + + private String metadataFieldName = METADATA_FIELD_NAME; + + private String embeddingFieldName = EMBEDDING_FIELD_NAME; + + private boolean initializeSchema = false; + + private MilvusServiceClient milvusClient; + + private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); + + /** + * @param milvusClient the Milvus service client to use for database operations + * @throws IllegalArgumentException if milvusClient is null + */ + public MilvusBuilder milvusClient(MilvusServiceClient milvusClient) { + Assert.notNull(milvusClient, "milvusClient must not be null"); + this.milvusClient = milvusClient; + return this; + } + + /** + * Configures the Milvus metric type to use for similarity calculations. See: + * https://milvus.io/docs/metric.md#floating for details on metric types. + * @param metricType the metric type to use (IP, L2, or COSINE) + * @return this builder instance + * @throws IllegalArgumentException if metricType is null or not one of IP, L2, or + * COSINE + */ + public MilvusBuilder metricType(MetricType metricType) { + Assert.notNull(metricType, "Collection Name must not be empty"); + Assert.isTrue(metricType == MetricType.IP || metricType == MetricType.L2 || metricType == MetricType.COSINE, + "Only the text metric types IP and L2 are supported"); + this.metricType = metricType; + return this; + } + + /** + * Configures the Milvus index type to use for vector search optimization. + * @param indexType the index type to use (defaults to IVF_FLAT if not specified) + * @return this builder instance + */ + public MilvusBuilder indexType(IndexType indexType) { + this.indexType = indexType; + return this; + } + + /** + * Configures the Milvus index parameters as a JSON string. + * @param indexParameters the index parameters to use (defaults to {"nlist":1024} + * if not specified) + * @return this builder instance + */ + public MilvusBuilder indexParameters(String indexParameters) { + this.indexParameters = indexParameters; + return this; + } + + /** + * Configures the Milvus database name. + * @param databaseName the database name to use (defaults to DEFAULT_DATABASE_NAME + * if not specified) + * @return this builder instance + */ + public MilvusBuilder databaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + /** + * Configures the Milvus collection name. + * @param collectionName the collection name to use (defaults to + * DEFAULT_COLLECTION_NAME if not specified) + * @return this builder instance + */ + public MilvusBuilder collectionName(String collectionName) { + this.collectionName = collectionName; + return this; + } + + /** + * Configures the dimension size of the embedding vectors. + * @param newEmbeddingDimension The dimension of the embedding (must be between 1 + * and 32768) + * @return this builder instance + * @throws IllegalArgumentException if dimension is not between 1 and 32768 + */ + public MilvusBuilder embeddingDimension(int newEmbeddingDimension) { + Assert.isTrue(newEmbeddingDimension >= 1 && newEmbeddingDimension <= 32768, + "Dimension has to be withing the boundaries 1 and 32768 (inclusively)"); + this.embeddingDimension = newEmbeddingDimension; + return this; + } + + /** + * Configures the name of the field used for document IDs. + * @param idFieldName The name for the ID field (defaults to DOC_ID_FIELD_NAME) + * @return this builder instance + */ + public MilvusBuilder iDFieldName(String idFieldName) { + this.idFieldName = idFieldName; + return this; + } + + /** + * Configures whether to use auto-generated IDs for documents. + * @param isAutoId true to enable auto-generated IDs, false to use provided IDs + * @return this builder instance + */ + public MilvusBuilder autoId(boolean isAutoId) { + this.isAutoId = isAutoId; + return this; + } + + /** + * Configures the name of the field used for document content. + * @param contentFieldName The name for the content field (defaults to + * CONTENT_FIELD_NAME) + * @return this builder instance + */ + public MilvusBuilder contentFieldName(String contentFieldName) { + this.contentFieldName = contentFieldName; + return this; + } + + /** + * Configures the name of the field used for document metadata. + * @param metadataFieldName The name for the metadata field (defaults to + * METADATA_FIELD_NAME) + * @return this builder instance + */ + public MilvusBuilder metadataFieldName(String metadataFieldName) { + this.metadataFieldName = metadataFieldName; + return this; + } + + /** + * Configures the name of the field used for embedding vectors. + * @param embeddingFieldName The name for the embedding field (defaults to + * EMBEDDING_FIELD_NAME) + * @return this builder instance + */ + public MilvusBuilder embeddingFieldName(String embeddingFieldName) { + this.embeddingFieldName = embeddingFieldName; + return this; + } + + /** + * Configures whether to initialize the collection schema automatically. + * @param initializeSchema true to initialize schema automatically, false to use + * existing schema + * @return this builder instance + */ + public MilvusBuilder initializeSchema(boolean initializeSchema) { + this.initializeSchema = initializeSchema; + return this; + } + + /** + * Configures the strategy for batching operations. + * @param batchingStrategy the batching strategy to use for grouping operations + * @return this builder instance + * @throws IllegalArgumentException if batchingStrategy is null + */ + public MilvusBuilder batchingStrategy(BatchingStrategy batchingStrategy) { + Assert.notNull(batchingStrategy, "batchingStrategy must not be null"); + this.batchingStrategy = batchingStrategy; + return this; + } + + /** + * Builds and returns a new MilvusVectorStore instance with the configured + * settings. + * @return a new MilvusVectorStore instance + * @throws IllegalStateException if the builder configuration is invalid + */ + public MilvusVectorStore build() { + validate(); + return new MilvusVectorStore(this); + } + + } + + @Deprecated(forRemoval = true, since = "1.0.0-M5") public static final class MilvusVectorStoreConfig { private final String databaseName; @@ -503,7 +823,6 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements * @return The entry point for creating a new configuration. */ public static Builder builder() { - return new Builder(); } @@ -514,6 +833,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements return builder().build(); } + @Deprecated(forRemoval = true, since = "1.0.0-M5") public static final class Builder { private String databaseName = DEFAULT_DATABASE_NAME; diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/package-info.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/package-info.java new file mode 100644 index 000000000..6f932739d --- /dev/null +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/package-info.java @@ -0,0 +1,25 @@ +/* + * Copyright 2023-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * Provides the API for embedding observations. + */ +@NonNullApi +@NonNullFields +package org.springframework.ai.milvus.vectorstore; + +import org.springframework.lang.NonNullApi; +import org.springframework.lang.NonNullFields; diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusEmbeddingDimensionsTests.java similarity index 72% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusEmbeddingDimensionsTests.java index 0e0dbf5b5..32cb665e9 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusEmbeddingDimensionsTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import io.milvus.client.MilvusServiceClient; import org.assertj.core.api.ThrowableAssert; @@ -27,7 +27,6 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; -import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -54,13 +53,14 @@ public class MilvusEmbeddingDimensionsTests { final int explicitDimensions = 696; - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withEmbeddingDimension(explicitDimensions) + MilvusVectorStore build = MilvusVectorStore.builder() + .milvusClient(this.milvusClient) + .embeddingModel(this.embeddingModel) + .initializeSchema(true) + .batchingStrategy(new TokenCountBatchingStrategy()) + .embeddingDimension(explicitDimensions) .build(); - - var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel, config, true, - new TokenCountBatchingStrategy()) - .embeddingDimensions(); + var dim = build.embeddingDimensions(); assertThat(dim).isEqualTo(explicitDimensions); verify(this.embeddingModel, never()).dimensions(); @@ -70,11 +70,13 @@ public class MilvusEmbeddingDimensionsTests { public void embeddingModelDimensions() { given(this.embeddingModel.dimensions()).willReturn(969); - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder().build(); - - var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel, config, true, - new TokenCountBatchingStrategy()) - .embeddingDimensions(); + MilvusVectorStore build = MilvusVectorStore.builder() + .milvusClient(this.milvusClient) + .embeddingModel(this.embeddingModel) + .initializeSchema(true) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); + var dim = build.embeddingDimensions(); assertThat(dim).isEqualTo(969); @@ -86,9 +88,13 @@ public class MilvusEmbeddingDimensionsTests { given(this.embeddingModel.dimensions()).willThrow(new RuntimeException()); - var dim = new MilvusVectorStore(this.milvusClient, this.embeddingModel, - MilvusVectorStoreConfig.builder().build(), true, new TokenCountBatchingStrategy()) - .embeddingDimensions(); + MilvusVectorStore build = MilvusVectorStore.builder() + .milvusClient(this.milvusClient) + .embeddingModel(this.embeddingModel) + .initializeSchema(true) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); + var dim = build.embeddingDimensions(); assertThat(dim).isEqualTo(MilvusVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE); verify(this.embeddingModel, only()).dimensions(); @@ -98,8 +104,9 @@ public class MilvusEmbeddingDimensionsTests { @ValueSource(ints = { 0, 32769 }) public void invalidDimensionsThrowException(final int explicitDimensions) { // when - ThrowableAssert.ThrowingCallable actual = () -> MilvusVectorStoreConfig.builder() - .withEmbeddingDimension(explicitDimensions) + ThrowableAssert.ThrowingCallable actual = () -> MilvusVectorStore.builder() + .milvusClient(this.milvusClient) + .embeddingDimension(explicitDimensions) .build(); // then diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverterTests.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverterTests.java similarity index 99% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverterTests.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverterTests.java index a2ef03889..3a9f5b5a3 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverterTests.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverterTests.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import java.util.List; diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusImage.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusImage.java similarity index 94% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusImage.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusImage.java index f06b66677..dc82f84cf 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusImage.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusImage.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import org.testcontainers.utility.DockerImageName; diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java similarity index 92% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java index 1e5ef703d..0934fcc3c 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -37,7 +37,8 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; @@ -229,18 +230,21 @@ class MilvusVectorStoreCustomFieldNamesIT { @Bean VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel) { - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withCollectionName("test_vector_store_custom_fields") - .withDatabaseName("default") - .withIndexType(IndexType.IVF_FLAT) - .withMetricType(this.metricType) - .withIDFieldName(this.idFieldName) - .withAutoId(this.isAutoId) - .withContentFieldName(this.contentFieldName) - .withEmbeddingFieldName(this.embeddingFieldName) - .withMetadataFieldName(this.metadataFieldName) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) + .embeddingModel(embeddingModel) + .collectionName("test_vector_store_custom_fields") + .databaseName("default") + .indexType(IndexType.IVF_FLAT) + .metricType(this.metricType) + .iDFieldName(this.idFieldName) + .autoId(this.isAutoId) + .contentFieldName(this.contentFieldName) + .embeddingFieldName(this.embeddingFieldName) + .metadataFieldName(this.metadataFieldName) + .batchingStrategy(new TokenCountBatchingStrategy()) + .initializeSchema(true) .build(); - return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy()); } @Bean 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/milvus/vectorstore/MilvusVectorStoreIT.java similarity index 95% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java index 757b5c1b1..b3f08670a 100644 --- 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/milvus/vectorstore/MilvusVectorStoreIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -40,7 +40,9 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.milvus.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; @@ -267,13 +269,16 @@ public class MilvusVectorStoreIT { @Bean public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel) { - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withCollectionName("test_vector_store") - .withDatabaseName("default") - .withIndexType(IndexType.IVF_FLAT) - .withMetricType(this.metricType) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) + .embeddingModel(embeddingModel) + .collectionName("test_vector_store") + .databaseName("default") + .indexType(IndexType.IVF_FLAT) + .metricType(this.metricType) + .batchingStrategy(new TokenCountBatchingStrategy()) + .initializeSchema(true) .build(); - return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy()); } @Bean diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java similarity index 92% rename from vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java rename to vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java index 7303e4d3c..1df3ac1ae 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.vectorstore; +package org.springframework.ai.milvus.vectorstore; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -42,7 +42,9 @@ import org.springframework.ai.observation.conventions.VectorStoreProvider; import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.milvus.vectorstore.MilvusVectorStore.MilvusVectorStoreConfig; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames; import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames; @@ -168,14 +170,17 @@ public class MilvusVectorStoreObservationIT { @Bean public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, ObservationRegistry observationRegistry) { - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withCollectionName(TEST_COLLECTION_NAME) - .withDatabaseName("default") - .withIndexType(IndexType.IVF_FLAT) - .withMetricType(MetricType.COSINE) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) + .embeddingModel(embeddingModel) + .observationRegistry(observationRegistry) + .collectionName(TEST_COLLECTION_NAME) + .databaseName("default") + .indexType(IndexType.IVF_FLAT) + .metricType(MetricType.COSINE) + .batchingStrategy(new TokenCountBatchingStrategy()) + .initializeSchema(true) .build(); - return new MilvusVectorStore(milvusClient, embeddingModel, config, true, new TokenCountBatchingStrategy(), - observationRegistry, null); } @Bean