Fix failing Coherence test

Fixes #2850

Signed-off-by: Mark Pollack <mark.pollack@broadcom.com>
This commit is contained in:
Mark Pollack
2025-05-09 16:22:51 -04:00
parent 08958af670
commit 20b560f6b2
2 changed files with 25 additions and 3 deletions

View File

@@ -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<String, Object> 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());
}

View File

@@ -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<Document> 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<Document> documents) {
final List<Double> distances = documents.stream()
.map(doc -> (Double) doc.getMetadata().get(DocumentMetadata.DISTANCE.value()))