Fix failing Coherence test
Fixes #2850 Signed-off-by: Mark Pollack <mark.pollack@broadcom.com>
This commit is contained in:
@@ -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());
|
||||
}
|
||||
|
||||
@@ -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()))
|
||||
|
||||
Reference in New Issue
Block a user