From 13972306320c84ace40a673c732536c4ae0249fc Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Fri, 6 Dec 2024 18:08:00 -0500 Subject: [PATCH 1/2] refactor(vectorstore): introduce builder pattern for MilvusVectorStore - Replace configuration object with fluent builder pattern - Move Milvus-related classes to dedicated milvus package - Deprecate MilvusVectorStoreConfig in favor of builder - Update constructor to use builder internally - Maintain backward compatibility with deprecated config - Add comprehensive builder methods with validation --- .../MilvusVectorStoreAutoConfiguration.java | 26 +- .../milvus/MilvusVectorStoreProperties.java | 2 +- .../MilvusFilterExpressionConverter.java | 2 +- .../vectorstore/MilvusVectorStore.java | 386 +++++++++++++++--- .../ai/milvus/vectorstore/package-info.java | 25 ++ .../MilvusEmbeddingDimensionsTests.java | 39 +- .../MilvusFilterExpressionConverterTests.java | 2 +- .../{ => milvus}/vectorstore/MilvusImage.java | 2 +- .../MilvusVectorStoreCustomFieldNamesIT.java | 29 +- .../vectorstore/MilvusVectorStoreIT.java | 20 +- .../MilvusVectorStoreObservationIT.java | 22 +- 11 files changed, 418 insertions(+), 137 deletions(-) rename vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusFilterExpressionConverter.java (97%) rename vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusVectorStore.java (63%) create mode 100644 vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/package-info.java rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusEmbeddingDimensionsTests.java (74%) rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusFilterExpressionConverterTests.java (99%) rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusImage.java (94%) rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java (92%) rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusVectorStoreIT.java (95%) rename vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/{ => milvus}/vectorstore/MilvusVectorStoreObservationIT.java (92%) 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 fde5c380b0e..90d48ac580c 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 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,13 @@ public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, Embedding 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) + .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 9e8457caa72..e498b627382 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 d8ded6c2e26..8e0f4b2c4f4 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 63% 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 57704b4204f..4c17d6d4532 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.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; @@ -110,44 +112,96 @@ 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).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 + */ + private MilvusVectorStore(MilvusBuilder builder) { + super(builder); + 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; + } - this.milvusClient = milvusClient; - this.embeddingModel = embeddingModel; - this.config = config; - this.batchingStrategy = batchingStrategy; + /** + * Creates a new MilvusBuilder instance with the specified Milvus client. This is the + * recommended way to instantiate a MilvusBuilder. + * @param milvusClient the Milvus service client to use for database operations + * @return a new MilvusBuilder instance + */ + public static MilvusBuilder builder(MilvusServiceClient milvusClient) { + return new MilvusBuilder(milvusClient); } @Override @@ -175,16 +229,16 @@ public void doAdd(List documents) { 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 +252,12 @@ public void doAdd(List documents) { 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 +277,20 @@ public List doSimilaritySearch(SearchRequest request) { 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 +308,11 @@ public List doSimilaritySearch(SearchRequest request) { .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 +332,7 @@ public List doSimilaritySearch(SearchRequest request) { 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 +350,16 @@ public void afterPropertiesSet() throws Exception { 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 +368,24 @@ private boolean isDatabaseCollectionExists() { 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 +395,8 @@ void createCollection() { } 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 +448,8 @@ void createCollection(String databaseName, String collectionName, String idField } 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 +459,7 @@ int embeddingDimensions() { } 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 +467,23 @@ int embeddingDimensions() { // 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 +496,222 @@ public org.springframework.ai.vectorstore.observation.VectorStoreObservationCont 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 final MilvusServiceClient milvusClient; + + private BatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(); + + /** + * Creates a new MilvusBuilder instance with the specified Milvus client. + * @param milvusClient the Milvus service client to use for database operations + * @throws IllegalArgumentException if milvusClient is null + */ + public MilvusBuilder(MilvusServiceClient milvusClient) { + Assert.notNull(milvusClient, "milvusClient must not be null"); + this.milvusClient = milvusClient; + } + + /** + * 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 +755,6 @@ private MilvusVectorStoreConfig(Builder builder) { * @return The entry point for creating a new configuration. */ public static Builder builder() { - return new Builder(); } @@ -514,6 +765,7 @@ public static MilvusVectorStoreConfig defaultConfig() { 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 00000000000..6f932739d09 --- /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 74% 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 0e0dbf5b586..dfa56d99d65 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.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,13 @@ public void explicitlySetDimensions() { final int explicitDimensions = 696; - MilvusVectorStoreConfig config = MilvusVectorStoreConfig.builder() - .withEmbeddingDimension(explicitDimensions) + MilvusVectorStore build = MilvusVectorStore.builder(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 +69,12 @@ public void explicitlySetDimensions() { 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(this.milvusClient) + .embeddingModel(this.embeddingModel) + .initializeSchema(true) + .batchingStrategy(new TokenCountBatchingStrategy()) + .build(); + var dim = build.embeddingDimensions(); assertThat(dim).isEqualTo(969); @@ -86,9 +86,12 @@ public void fallBackToDefaultDimensions() { 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(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 +101,8 @@ public void fallBackToDefaultDimensions() { @ValueSource(ints = { 0, 32769 }) public void invalidDimensionsThrowException(final int explicitDimensions) { // when - ThrowableAssert.ThrowingCallable actual = () -> MilvusVectorStoreConfig.builder() - .withEmbeddingDimension(explicitDimensions) + ThrowableAssert.ThrowingCallable actual = () -> MilvusVectorStore.builder(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 a2ef0388979..3a9f5b5a3a3 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 f06b6667790..dc82f84cfaf 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 1e5ef703d48..30a0631bd0d 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.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,20 @@ static class TestApplication { @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) + .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 757b5c1b154..f72f73004f8 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.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,15 @@ public static class TestApplication { @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) + .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 7303e4d3cda..2b8d2eac5d2 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.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,16 @@ public TestObservationRegistry observationRegistry() { @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) + .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 From b22cbb836c3ce8009be4e61d9e040ca2af405884 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 10 Dec 2024 19:07:53 -0500 Subject: [PATCH 2/2] Milvus vector store builder refactoring --- .../MilvusVectorStoreAutoConfiguration.java | 3 ++- .../milvus/vectorstore/MilvusVectorStore.java | 19 +++++++++++-------- .../MilvusEmbeddingDimensionsTests.java | 12 ++++++++---- .../MilvusVectorStoreCustomFieldNamesIT.java | 3 ++- .../vectorstore/MilvusVectorStoreIT.java | 3 ++- .../MilvusVectorStoreObservationIT.java | 3 ++- 6 files changed, 27 insertions(+), 16 deletions(-) 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 90d48ac580c..1dc3eb5dfb0 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 @@ -71,7 +71,8 @@ public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, Embedding ObjectProvider observationRegistry, ObjectProvider customObservationConvention) { - return MilvusVectorStore.builder(milvusClient) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) .embeddingModel(embeddingModel) .initializeSchema(properties.isInitializeSchema()) .batchingStrategy(batchingStrategy) diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java index 4c17d6d4532..86240c45142 100644 --- a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java @@ -165,7 +165,8 @@ public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embedd MilvusVectorStoreConfig config, boolean initializeSchema, BatchingStrategy batchingStrategy, ObservationRegistry observationRegistry, VectorStoreObservationConvention customObservationConvention) { - this(builder(milvusClient).embeddingModel(embeddingModel) + this(builder().milvusClient(milvusClient) + .embeddingModel(embeddingModel) .observationRegistry(observationRegistry) .customObservationConvention(customObservationConvention) .initializeSchema(initializeSchema) @@ -175,8 +176,11 @@ public MilvusVectorStore(MilvusServiceClient milvusClient, EmbeddingModel embedd /** * @param builder {@link Builder} for chroma vector store */ - private MilvusVectorStore(MilvusBuilder builder) { + protected MilvusVectorStore(MilvusBuilder builder) { super(builder); + + Assert.notNull(builder.milvusClient, "milvusClient must not be null"); + this.milvusClient = builder.milvusClient; this.batchingStrategy = builder.batchingStrategy; this.initializeSchema = builder.initializeSchema; @@ -197,11 +201,10 @@ private MilvusVectorStore(MilvusBuilder builder) { /** * Creates a new MilvusBuilder instance with the specified Milvus client. This is the * recommended way to instantiate a MilvusBuilder. - * @param milvusClient the Milvus service client to use for database operations * @return a new MilvusBuilder instance */ - public static MilvusBuilder builder(MilvusServiceClient milvusClient) { - return new MilvusBuilder(milvusClient); + public static MilvusBuilder builder() { + return new MilvusBuilder(); } @Override @@ -535,18 +538,18 @@ public static final class MilvusBuilder extends AbstractVectorStoreBuilder MilvusVectorStore.builder(this.milvusClient) + ThrowableAssert.ThrowingCallable actual = () -> MilvusVectorStore.builder() + .milvusClient(this.milvusClient) .embeddingDimension(explicitDimensions) .build(); diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java index 30a0631bd0d..0934fcc3c18 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java @@ -230,7 +230,8 @@ static class TestApplication { @Bean VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel) { - return MilvusVectorStore.builder(milvusClient) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) .embeddingModel(embeddingModel) .collectionName("test_vector_store_custom_fields") .databaseName("default") diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java index f72f73004f8..b3f08670aca 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java @@ -269,7 +269,8 @@ public static class TestApplication { @Bean public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel) { - return MilvusVectorStore.builder(milvusClient) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) .embeddingModel(embeddingModel) .collectionName("test_vector_store") .databaseName("default") diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java index 2b8d2eac5d2..1df3ac1aedb 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java @@ -170,7 +170,8 @@ public TestObservationRegistry observationRegistry() { @Bean public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel, ObservationRegistry observationRegistry) { - return MilvusVectorStore.builder(milvusClient) + return MilvusVectorStore.builder() + .milvusClient(milvusClient) .embeddingModel(embeddingModel) .observationRegistry(observationRegistry) .collectionName(TEST_COLLECTION_NAME)