Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Support similarity scores in Document API

Browse Source 
Document
* Introduced “score” attribute in Document API. It stores the similarity score.
* Consolidate “distance” metadata for Documents. It stores the distance measurement.
* Adopted prefix-less naming convention in Document.Builder and deprecated old methods.
* Deprecated the many overloaded Document constructors in favour of Document.Builder.

Vector Stores
* Every vector store implementation now configures a “score” attribute with the similarity score of the Document embedding. It also includes the “distance” metadata with the distance measurement.
* Fixed error in Elasticsearch where distance and similarity were mixed up.
* Added missing integration tests for SimpleVectorStore.
* The Azure Vector Store and HanaDB Vector Store do not include those measurements because the product documentation do not include information about how the similarity score is returned, and without access to the cloud products I could not verify that via debugging.
* Improved tests to actually assert the result of the similarity search based on the returned score.

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-11-26 08:03:40 +01:00
committed by Mark Pollack
parent 50223d20e3
commit fe58fd30eb
59 changed files with 1096 additions and 434 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -35,6 +35,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
@@ -502,12 +503,15 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
Float distance = rs.getFloat(COLUMN_DISTANCE);
Map<String, Object> metadata = toMap(pgMetadata);
metadata.put(COLUMN_DISTANCE, distance);
metadata.put(DocumentMetadata.DISTANCE.value(), distance);
Document document = new Document(id, content, metadata);
document.setEmbedding(toFloatArray(embedding));
return document;
return Document.builder()
.id(id)
.content(content)
.metadata(metadata)
.score(1.0 - distance)
.embedding(toFloatArray(embedding))
.build();
}
private float[] toFloatArray(PGobject embedding) throws SQLException {

32
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -34,6 +34,7 @@ import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.ai.document.DocumentMetadata;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -113,20 +114,20 @@ public class PgVectorStoreIT {
);
}
private static boolean isSortedByDistance(List<Document> docs) {
private static boolean isSortedBySimilarity(List<Document> docs) {
List<Float> distances = docs.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
List<Double> scores = docs.stream().map(Document::getScore).toList();
if (CollectionUtils.isEmpty(distances) || distances.size() == 1) {
if (CollectionUtils.isEmpty(scores) || scores.size() == 1) {
return true;
}
Iterator<Float> iter = distances.iterator();
Float current;
Float previous = iter.next();
Iterator<Double> iter = scores.iterator();
Double current;
Double previous = iter.next();
while (iter.hasNext()) {
current = iter.next();
if (previous > current) {
if (previous < current) {
return false;
}
previous = current;
@@ -150,7 +151,7 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", DocumentMetadata.DISTANCE.value());
// Remove all documents from the store
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
@@ -289,7 +290,7 @@ public class PgVectorStoreIT {
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", "distance");
assertThat(resultDoc.getMetadata()).containsKeys("meta1", DocumentMetadata.DISTANCE.value());
Document sameIdDocument = new Document(document.getId(),
"The World is Big and Salvation Lurks Around the Corner",
@@ -303,7 +304,7 @@ public class PgVectorStoreIT {
resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(document.getId());
assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
assertThat(resultDoc.getMetadata()).containsKeys("meta2", DocumentMetadata.DISTANCE.value());
dropTable(context);
});
@@ -326,20 +327,19 @@ public class PgVectorStoreIT {
assertThat(fullResult).hasSize(3);
assertThat(isSortedByDistance(fullResult)).isTrue();
assertThat(isSortedBySimilarity(fullResult)).isTrue();
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.toList();
List<Double> scores = fullResult.stream().map(Document::getScore).toList();
float threshold = (distances.get(0) + distances.get(1)) / 2;
double similarityThreshold = (scores.get(0) + scores.get(1)) / 2;
List<Document> results = vectorStore.similaritySearch(
SearchRequest.query("Time Shelter").withTopK(5).withSimilarityThreshold(1 - threshold));
SearchRequest.query("Time Shelter").withTopK(5).withSimilarityThreshold(similarityThreshold));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(1).getId());
assertThat(resultDoc.getScore()).isGreaterThanOrEqualTo(similarityThreshold);
dropTable(context);
});