From 20b560f6b26de0bdfbe5a0db5bcafab9caf9f0a3 Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Fri, 9 May 2025 16:22:51 -0400 Subject: [PATCH] Fix failing Coherence test Fixes #2850 Signed-off-by: Mark Pollack --- .../coherence/CoherenceVectorStore.java | 5 ++-- .../coherence/CoherenceVectorStoreIT.java | 23 ++++++++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java b/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java index 5caa0a380..ef73947bb 100644 --- a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java +++ b/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java @@ -209,11 +209,12 @@ public class CoherenceVectorStore extends AbstractObservationVectorStore impleme if (this.distanceType != DistanceType.COSINE || (1 - r.getDistance()) >= request.getSimilarityThreshold()) { DocumentChunk.Id id = r.getKey(); DocumentChunk chunk = r.getValue(); - chunk.metadata().put(DocumentMetadata.DISTANCE.value(), r.getDistance()); + Map mergedMetadata = new HashMap<>(chunk.metadata()); + mergedMetadata.put(DocumentMetadata.DISTANCE.value(), r.getDistance()); documents.add(Document.builder() .id(id.docId()) .text(chunk.text()) - .metadata(chunk.metadata()) + .metadata(mergedMetadata) .score(1 - r.getDistance()) .build()); } diff --git a/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java b/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java index 526fb361a..2cb74d1bb 100644 --- a/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java +++ b/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java @@ -281,10 +281,15 @@ public class CoherenceVectorStoreIT { .similarityThreshold(similarityThreshold) .build()); + // Debug: print all returned document IDs and metadata + for (Document doc : results) { + System.out.println("Returned doc ID: " + doc.getId() + ", metadata: " + doc.getMetadata()); + } + assertThat(results).hasSize(1); Document resultDoc = results.get(0); assertThat(resultDoc.getId()).isEqualTo(this.documents.get(1).getId()); - assertThat(resultDoc.getMetadata()).containsKeys("meta1", DocumentMetadata.DISTANCE.value()); + assertThat(resultDoc.getMetadata()).containsKey(DocumentMetadata.DISTANCE.value()); assertThat(resultDoc.getScore()).isGreaterThanOrEqualTo(similarityThreshold); truncateMap(context, ((CoherenceVectorStore) vectorStore).getMapName()); @@ -300,6 +305,22 @@ public class CoherenceVectorStoreIT { }); } + @Test + public void similaritySearchReturnsMetadata() { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + vectorStore.add(this.documents); + + // Query that matches the first document, which has meta1 + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("spring ai").topK(1).build()); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getMetadata()).containsKeys("meta1", DocumentMetadata.DISTANCE.value()); + }); + } + private static boolean isSortedByDistance(final List documents) { final List distances = documents.stream() .map(doc -> (Double) doc.getMetadata().get(DocumentMetadata.DISTANCE.value()))