Skip to content

Commit 7672ae9

Browse files
ilayaperumalgilayaperumalg
committed
Upgrade Milvus java sdk to 2.5.4
 - Upgrade Milvus java sdk to 2.5.4 - Remove use of fastjson - Use com.google.gson.JsonObject as Milvus SDK replaces fastjson with gson - Change SearchResult RowRecord similarity metric name to score - Milvus SDK 2.5.4 uses "score" instead of "distance" Signed-off-by: Ilayaperumal Gopinathan <[email protected]>
1 parent b29076f commit 7672ae9

File tree

7 files changed

+37
-34
lines changed

7 files changed

+37
-34
lines changed

‎pom.xml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,10 +216,10 @@
216216
<pgvector.version>0.1.6</pgvector.version>
217217
<sap.hanadb.version>2.20.11</sap.hanadb.version>
218218
<coherence.version>24.09</coherence.version>
219-
<milvus.version>2.3.5</milvus.version>
219+
<milvus.version>2.5.4</milvus.version>
220220
<gemfire.testcontainers.version>2.3.0</gemfire.testcontainers.version>
221221
<pinecone.version>0.8.0</pinecone.version>
222-
<fastjson.version>2.0.46</fastjson.version>
222+
<fastjson2.version>2.0.46</fastjson2.version>
223223
<azure-core.version>1.53.0</azure-core.version>
224224
<azure-json.version>1.3.0</azure-json.version>
225225
<azure-identity.version>1.14.0</azure-identity.version>

‎spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,8 @@ protected void deleteById() {
110110
assertThat(results.get(0).getId()).isEqualTo(documents.get(2).getId());
111111
Map<String, Object> metadata = results.get(0).getMetadata();
112112
assertThat(normalizeValue(metadata.get("country"))).isEqualTo("BG");
113-
assertThat(normalizeValue(metadata.get("year"))).isEqualTo("2023");
113+
// the values are converted into Double
114+
assertThat(normalizeValue(metadata.get("year"))).isEqualTo("2023.0");
114115

115116
vectorStore.delete(List.of(documents.get(2).getId()));
116117
});

‎vector-stores/spring-ai-azure-store/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@
8080
<dependency>
8181
<groupId>com.alibaba.fastjson2</groupId>
8282
<artifactId>fastjson2</artifactId>
83-
<version>${fastjson.version}</version>
83+
<version>${fastjson2.version}</version>
8484
</dependency>
8585

8686
<!-- TESTING -->

‎vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,16 @@
1616

1717
package org.springframework.ai.vectorstore.milvus;
1818

19+
import java.lang.reflect.Type;
1920
import java.util.ArrayList;
2021
import java.util.List;
2122
import java.util.Map;
2223
import java.util.Optional;
2324
import java.util.stream.Collectors;
2425

25-
import com.alibaba.fastjson.JSONObject;
26+
import com.google.gson.Gson;
27+
import com.google.gson.JsonObject;
28+
import com.google.gson.reflect.TypeToken;
2629
import io.milvus.client.MilvusServiceClient;
2730
import io.milvus.common.clientenum.ConsistencyLevelEnum;
2831
import io.milvus.exception.ParamException;
@@ -159,7 +162,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
159162
public static final String EMBEDDING_FIELD_NAME = "embedding";
160163

161164
// Metadata, automatically assigned by Milvus.
162-
private static final String DISTANCE_FIELD_NAME = "distance";
165+
public static final String SIMILARITY_FIELD_NAME = "score";
163166

164167
private static final Logger logger = LoggerFactory.getLogger(MilvusVectorStore.class);
165168

@@ -234,7 +237,7 @@ public void doAdd(List<Document> documents) {
234237

235238
List<String> docIdArray = new ArrayList<>();
236239
List<String> contentArray = new ArrayList<>();
237-
List<JSONObject> metadataArray = new ArrayList<>();
240+
List<JsonObject> metadataArray = new ArrayList<>();
238241
List<List<Float>> embeddingArray = new ArrayList<>();
239242

240243
// TODO: Need to customize how we pass the embedding options
@@ -246,7 +249,9 @@ public void doAdd(List<Document> documents) {
246249
// Use a (future) DocumentTextLayoutFormatter instance to extract
247250
// the content used to compute the embeddings
248251
contentArray.add(document.getText());
249-
metadataArray.add(new JSONObject(document.getMetadata()));
252+
Gson gson = new Gson();
253+
String jsonString = gson.toJson(document.getMetadata());
254+
metadataArray.add(gson.fromJson(jsonString, JsonObject.class));
250255
embeddingArray.add(EmbeddingUtils.toList(embeddings.get(documents.indexOf(document))));
251256
}
252257

@@ -357,29 +362,32 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
357362
.map(rowRecord -> {
358363
String docId = String.valueOf(rowRecord.get(this.idFieldName));
359364
String content = (String) rowRecord.get(this.contentFieldName);
360-
JSONObject metadata = null;
365+
JsonObject metadata = new JsonObject();
361366
try {
362-
metadata = (JSONObject) rowRecord.get(this.metadataFieldName);
367+
metadata = (JsonObject) rowRecord.get(this.metadataFieldName);
363368
// inject the distance into the metadata.
364-
metadata.put(DocumentMetadata.DISTANCE.value(), 1 - getResultSimilarity(rowRecord));
369+
metadata.addProperty(DocumentMetadata.DISTANCE.value(), 1 - getResultSimilarity(rowRecord));
365370
}
366371
catch (ParamException e) {
367372
// skip the ParamException if metadata doesn't exist for the custom
368373
// collection
369374
}
375+
Gson gson = new Gson();
376+
Type type = new TypeToken<Map<String, Object>>() {
377+
}.getType();
370378
return Document.builder()
371379
.id(docId)
372380
.text(content)
373-
.metadata((metadata != null) ? metadata.getInnerMap() : Map.of())
381+
.metadata((metadata != null) ? gson.fromJson(metadata, type) : Map.of())
374382
.score((double) getResultSimilarity(rowRecord))
375383
.build();
376384
})
377385
.toList();
378386
}
379387

380388
private float getResultSimilarity(RowRecord rowRecord) {
381-
Float distance = (Float) rowRecord.get(DISTANCE_FIELD_NAME);
382-
return (this.metricType == MetricType.IP || this.metricType == MetricType.COSINE) ? distance : (1 - distance);
389+
Float score = (Float) rowRecord.get(SIMILARITY_FIELD_NAME);
390+
return (this.metricType == MetricType.IP || this.metricType == MetricType.COSINE) ? score : (1 - score);
383391
}
384392

385393
// ---------------------------------------------------------------------------------

‎vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusImage.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
*/
2424
public final class MilvusImage {
2525

26-
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("milvusdb/milvus:v2.4.9");
26+
public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("milvusdb/milvus:v2.5.4");
2727

2828
private MilvusImage() {
2929

‎vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreCustomFieldNamesIT.java‎

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -103,16 +103,14 @@ void searchWithCustomFieldNames(String metricType) {
103103
List<Document> fullResult = vectorStore
104104
.similaritySearch(SearchRequest.builder().query("Spring").build());
105105

106-
List<Float> distances = fullResult.stream()
107-
.map(doc -> (Float) doc.getMetadata().get("distance"))
108-
.toList();
106+
List<Double> scores = fullResult.stream().map(doc -> doc.getScore()).toList();
109107

110-
assertThat(distances).hasSize(3);
108+
assertThat(scores).hasSize(3);
111109

112-
float threshold = (distances.get(0) + distances.get(1)) / 2;
110+
double threshold = (scores.get(0) + scores.get(1)) / 2;
113111

114112
List<Document> results = vectorStore.similaritySearch(
115-
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(1 - threshold).build());
113+
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(threshold).build());
116114

117115
assertThat(results).hasSize(1);
118116
Document resultDoc = results.get(0);
@@ -144,16 +142,14 @@ void searchWithoutMetadataFieldOverride(String metricType) {
144142
List<Document> fullResult = vectorStore
145143
.similaritySearch(SearchRequest.builder().query("Spring").build());
146144

147-
List<Float> distances = fullResult.stream()
148-
.map(doc -> (Float) doc.getMetadata().get("distance"))
149-
.toList();
145+
List<Double> scores = fullResult.stream().map(doc -> doc.getScore()).toList();
150146

151-
assertThat(distances).hasSize(3);
147+
assertThat(scores).hasSize(3);
152148

153-
float threshold = (distances.get(0) + distances.get(1)) / 2;
149+
double threshold = (scores.get(0) + scores.get(1)) / 2;
154150

155151
List<Document> results = vectorStore.similaritySearch(
156-
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(1 - threshold).build());
152+
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(threshold).build());
157153

158154
assertThat(results).hasSize(1);
159155
Document resultDoc = results.get(0);
@@ -187,16 +183,14 @@ void searchWithAutoIdEnabled(String metricType) {
187183
List<Document> fullResult = vectorStore
188184
.similaritySearch(SearchRequest.builder().query("Spring").build());
189185

190-
List<Float> distances = fullResult.stream()
191-
.map(doc -> (Float) doc.getMetadata().get("distance"))
192-
.toList();
186+
List<Double> scores = fullResult.stream().map(doc -> doc.getScore()).toList();
193187

194-
assertThat(distances).hasSize(3);
188+
assertThat(scores).hasSize(3);
195189

196-
float threshold = (distances.get(0) + distances.get(1)) / 2;
190+
double threshold = (scores.get(0) + scores.get(1)) / 2;
197191

198192
List<Document> results = vectorStore.similaritySearch(
199-
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(1 - threshold).build());
193+
SearchRequest.builder().query("Spring").topK(5).similarityThreshold(threshold).build());
200194

201195
assertThat(results).hasSize(1);
202196
Document resultDoc = results.get(0);

‎vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,7 @@ public void deleteWithComplexFilterExpression() {
319319
assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList()))
320320
.containsExactlyInAnyOrder("A", "B");
321321
assertThat(results.stream().map(doc -> doc.getMetadata().get("priority")).collect(Collectors.toList()))
322-
.containsExactlyInAnyOrder(1, 1);
322+
.containsExactlyInAnyOrder(1.0, 1.0);
323323
});
324324
}
325325

0 commit comments

Comments
 (0)