Skip to content

Commit af1bb56

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 af1bb56

File tree

5 files changed

+540
-55
lines changed

5 files changed

+540
-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: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
/**
2525
* @author Christian Tzolov
26+
* @author Ilayaperumal Gopinathan
2627
*/
2728
@ConfigurationProperties(MilvusVectorStoreProperties.CONFIG_PREFIX)
2829
public class MilvusVectorStoreProperties extends CommonVectorStoreProperties {
@@ -59,6 +60,31 @@ public class MilvusVectorStoreProperties extends CommonVectorStoreProperties {
5960
*/
6061
private String indexParameters = "{\"nlist\":1024}";
6162

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

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

118188
/**

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)