Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add builder pattern for MilvusVectorStore

Browse Source 
Introduces a fluent builder API to improve configuration readability and
type safety when creating MilvusVectorStore instances. This replaces the
existing configuration object approach which was less intuitive and harder
to maintain.

The builder pattern provides better encapsulation of configuration logic
and validation, while maintaining backward compatibility through a
deprecated config class. This change makes the codebase more maintainable
and the API more discoverable for users.

Key changes:
- 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
This commit is contained in:
Soby Chacko
2024-12-06 18:08:00 -05:00
committed by Mark Pollack
parent 5b11501cbe
commit 218c9677f4
11 changed files with 496 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfiguration.java
View File

@@ -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> observationRegistry,
ObjectProvider<VectorStoreObservationConvention> 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

2
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreProperties.java
View File

@@ -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;

2
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverter.java → vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverter.java
View File

@@ -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;

458
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java → vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStore.java
View File

@@ -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.
*
* <p>
* Key features include:
* <ul>
* <li>Support for different similarity metrics (Cosine, L2, Inner Product)</li>
* <li>Configurable index types for performance optimization</li>
* <li>Metadata filtering capabilities</li>
* <li>Automatic schema initialization</li>
* <li>Batching strategy support for efficient operations</li>
* </ul>
*
* <p>
* Example usage: <pre>{@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<Document> documents = List.of(
* new Document("content1", Map.of("meta1", "value1")),
* new Document("content2", Map.of("meta2", "value2"))
* );
* vectorStore.add(documents);
*
* // Perform similarity search
* List<Document> results = vectorStore.similaritySearch(
* SearchRequest.query("search text")
* .withTopK(5)
* .withSimilarityThreshold(0.7)
* .withFilterExpression("meta1 == 'value1'")
* );
* }</pre>
*
* <p>
* The vector store supports various configuration options through its builder:
* <ul>
* <li>{@code milvusClient}: Required Milvus service client for database operations</li>
* <li>{@code embeddingModel}: Required model for generating embeddings</li>
* <li>{@code metricType}: Similarity metric (COSINE, L2, IP)</li>
* <li>{@code indexType}: Type of index for search optimization</li>
* <li>{@code databaseName}: Name of the Milvus database (default: "default")</li>
* <li>{@code collectionName}: Name of the collection (default: "vector_store")</li>
* <li>{@code initializeSchema}: Whether to automatically create the schema</li>
* </ul>
*
* @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<InsertParam.Field> 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<Boolean> doDelete(List<String> 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<MutationResult> 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<String> 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<DescribeIndexResponse> 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<RpcStatus> 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<RpcStatus> 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<RpcStatus> status = this.milvusClient.releaseCollection(
ReleaseCollectionParam.newBuilder().withCollectionName(this.config.collectionName).build());
R<RpcStatus> 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<MilvusBuilder> {
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;

25
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/milvus/vectorstore/package-info.java Normal file
View File

@@ -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;

43
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusEmbeddingDimensionsTests.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusEmbeddingDimensionsTests.java
View File

@@ -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

2
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusFilterExpressionConverterTests.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusFilterExpressionConverterTests.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.milvus.vectorstore;
import java.util.List;

2
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusImage.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusImage.java
View File

@@ -14,7 +14,7 @@
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
package org.springframework.ai.milvus.vectorstore;
import org.testcontainers.utility.DockerImageName;

30
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreCustomFieldNamesIT.java
View File

@@ -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

21
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreIT.java
View File

@@ -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

23
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreObservationIT.java → vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/milvus/vectorstore/MilvusVectorStoreObservationIT.java
View File

@@ -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