Skip to content

Commit 3f70ed3

Browse files
committed
GH-513 fix: Support custom field names for Milvus VectorStore collection
- Add configuration properties to override the default field names for doc_id, content, metadata and embedding - Add support to allow auto-id when enabled - Add tests Resolves #513
1 parent c02d136 commit 3f70ed3

File tree

5 files changed

+542
-55
lines changed

5 files changed

+542
-55
lines changed

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
* @author Christian Tzolov
4343
* @author Eddú Meléndez
4444
* @author Soby Chacko
45+
* @author Ilayaperumal Gopinathan
4546
*/
4647
@AutoConfiguration
4748
@ConditionalOnClass({ MilvusVectorStore.class, EmbeddingModel.class })
@@ -75,6 +76,11 @@ public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, Embedding
7576
.withMetricType(MetricType.valueOf(properties.getMetricType().name()))
7677
.withIndexParameters(properties.getIndexParameters())
7778
.withEmbeddingDimension(properties.getEmbeddingDimension())
79+
.withIDFieldName(properties.getIdFieldName())
80+
.withAutoId(properties.isAutoId())
81+
.withContentFieldName(properties.getContentFieldName())
82+
.withMetadataFieldName(properties.getMetadataFieldName())
83+
.withEmbeddingFieldName(properties.getEmbeddingFieldName())
7884
.build();
7985

8086
return new MilvusVectorStore(milvusClient, embeddingModel, config, properties.isInitializeSchema(),

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,11 @@
2121
import org.springframework.boot.context.properties.ConfigurationProperties;
2222
import org.springframework.util.Assert;
2323

24+
import java.util.List;
25+
2426
/**
2527
* @author Christian Tzolov
28+
* @author Ilayaperumal Gopinathan
2629
*/
2730
@ConfigurationProperties(MilvusVectorStoreProperties.CONFIG_PREFIX)
2831
public class MilvusVectorStoreProperties extends CommonVectorStoreProperties {
@@ -59,6 +62,31 @@ public class MilvusVectorStoreProperties extends CommonVectorStoreProperties {
5962
*/
6063
private String indexParameters = "{\"nlist\":1024}";
6164

65+
/**
66+
* The ID field name for the collection.
67+
*/
68+
private String idFieldName = MilvusVectorStore.DOC_ID_FIELD_NAME;
69+
70+
/**
71+
* Boolean flag to indicate if the auto-id is used.
72+
*/
73+
private boolean isAutoId = false;
74+
75+
/**
76+
* The content field name for the collection.
77+
*/
78+
private String contentFieldName = MilvusVectorStore.CONTENT_FIELD_NAME;
79+
80+
/**
81+
* The metadata field name for the collection.
82+
*/
83+
private String metadataFieldName = MilvusVectorStore.METADATA_FIELD_NAME;
84+
85+
/**
86+
* The embedding field name for the collection.
87+
*/
88+
private String embeddingFieldName = MilvusVectorStore.EMBEDDING_FIELD_NAME;
89+
6290
public String getDatabaseName() {
6391
return this.databaseName;
6492
}
@@ -113,6 +141,50 @@ public void setIndexParameters(String indexParameters) {
113141
this.indexParameters = indexParameters;
114142
}
115143

144+
public String getIdFieldName() {
145+
return this.idFieldName;
146+
}
147+
148+
public void setIdFieldName(String idFieldName) {
149+
Assert.notNull(idFieldName, "idFieldName can not be null");
150+
this.idFieldName = idFieldName;
151+
}
152+
153+
public boolean isAutoId() {
154+
return this.isAutoId;
155+
}
156+
157+
public void setAutoId(boolean autoId) {
158+
this.isAutoId = autoId;
159+
}
160+
161+
public String getContentFieldName() {
162+
return this.contentFieldName;
163+
}
164+
165+
public void setContentFieldName(String contentFieldName) {
166+
Assert.notNull(contentFieldName, "contentFieldName can not be null");
167+
this.contentFieldName = contentFieldName;
168+
}
169+
170+
public String getMetadataFieldName() {
171+
return this.metadataFieldName;
172+
}
173+
174+
public void setMetadataFieldName(String metadataFieldName) {
175+
Assert.notNull(metadataFieldName, "metadataFieldName can not be null");
176+
this.metadataFieldName = metadataFieldName;
177+
}
178+
179+
public String getEmbeddingFieldName() {
180+
return this.embeddingFieldName;
181+
}
182+
183+
public void setEmbeddingFieldName(String embeddingFieldName) {
184+
Assert.notNull(embeddingFieldName, "embeddingFieldName can not be null");
185+
this.embeddingFieldName = embeddingFieldName;
186+
}
187+
116188
public enum MilvusMetricType {
117189

118190
/**

spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/milvus/MilvusVectorStoreAutoConfigurationIT.java

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,57 @@ public void addAndSearch() {
109109
});
110110
}
111111

112+
@Test
113+
public void searchWithCustomFields() {
114+
contextRunner
115+
.withPropertyValues("spring.ai.vectorstore.milvus.metricType=COSINE",
116+
"spring.ai.vectorstore.milvus.indexType=IVF_FLAT",
117+
"spring.ai.vectorstore.milvus.embeddingDimension=384",
118+
"spring.ai.vectorstore.milvus.collectionName=myCustomCollection",
119+
"spring.ai.vectorstore.milvus.idFieldName=identity",
120+
"spring.ai.vectorstore.milvus.contentFieldName=text",
121+
"spring.ai.vectorstore.milvus.embeddingFieldName=vectors",
122+
"spring.ai.vectorstore.milvus.metadataFieldName=meta",
123+
"spring.ai.vectorstore.milvus.initializeSchema=true",
124+
"spring.ai.vectorstore.milvus.client.host=" + milvus.getHost(),
125+
"spring.ai.vectorstore.milvus.client.port=" + milvus.getMappedPort(19530))
126+
.run(context -> {
127+
VectorStore vectorStore = context.getBean(VectorStore.class);
128+
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
129+
130+
vectorStore.add(documents);
131+
132+
assertObservationRegistry(observationRegistry, VectorStoreProvider.MILVUS,
133+
VectorStoreObservationContext.Operation.ADD);
134+
observationRegistry.clear();
135+
136+
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
137+
138+
assertThat(results).hasSize(1);
139+
Document resultDoc = results.get(0);
140+
assertThat(resultDoc.getId()).isEqualTo(documents.get(0).getId());
141+
assertThat(resultDoc.getContent()).contains(
142+
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
143+
assertThat(resultDoc.getMetadata()).hasSize(2);
144+
assertThat(resultDoc.getMetadata()).containsKeys("spring", "distance");
145+
146+
assertObservationRegistry(observationRegistry, VectorStoreProvider.MILVUS,
147+
VectorStoreObservationContext.Operation.QUERY);
148+
observationRegistry.clear();
149+
150+
// Remove all documents from the store
151+
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
152+
153+
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
154+
assertThat(results).hasSize(0);
155+
156+
assertObservationRegistry(observationRegistry, VectorStoreProvider.MILVUS,
157+
VectorStoreObservationContext.Operation.DELETE);
158+
observationRegistry.clear();
159+
160+
});
161+
}
162+
112163
@Configuration(proxyBeanMethods = false)
113164
static class Config {
114165

0 commit comments

Comments
 (0)