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

GH-1826 Fix EmbeddingModel's usage on Document#embedding

Browse Source 
- Since the Document object's reference to the `embedding` is deprecated and will be removed, the VectorStore implementations require a way to store the embedding of the corresponding Document objects

     - One way to fix this is, to have the EmbeddingModel#embed to return the embeddings in the same order as that of the Documents passed to it.

       - Since both the Document and embedding collections use the List object, their iteration operation will make sure to keep them in line with the same order.

       - A fix is required to preserve the order when batching strategy is applied.
	  - Updated the Javadoc for BatchingStrategy
          - Fixed the Document List order in TokenCountBatchingStrategy

    - Refactored the vector store implementations to update this change

Resolves #GH-1826
This commit is contained in:
Ilayaperumal Gopinathan
2024-11-25 22:21:11 +00:00
committed by Christian Tzolov
parent 6cfe5e79e8
commit ebd29e0959
24 changed files with 118 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java
View File

@@ -181,7 +181,8 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
public void doAdd(List<Document> documents) {
var futures = new CompletableFuture[documents.size()];
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
int i = 0;
for (Document d : documents) {
@@ -196,7 +197,8 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
builder = builder.setString(this.conf.schema.content(), d.getContent())
.setVector(this.conf.schema.embedding(),
CqlVector.newInstance(EmbeddingUtils.toList(d.getEmbedding())), Float.class);
CqlVector.newInstance(EmbeddingUtils.toList(embeddings.get(documents.indexOf(d)))),
Float.class);
for (var metadataColumn : this.conf.schema.metadataColumns()
.stream()
@@ -265,10 +267,6 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
.score((double) score)
.build();
if (this.conf.returnEmbeddings) {
doc.setEmbedding(EmbeddingUtils
.toPrimitive(row.getVector(this.conf.schema.embedding(), Float.class).stream().toList()));
}
documents.add(doc);
}
return documents;

2
vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java
View File

@@ -90,6 +90,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
final boolean disallowSchemaChanges;
// TODO: Remove this flag as the document no longer holds embeddings.
@Deprecated(since = "1.0.0-M5", forRemoval = true)
final boolean returnEmbeddings;
final DocumentIdTranslator documentIdTranslator;

11
vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java
View File

@@ -122,18 +122,12 @@ class CassandraVectorStoreIT {
List<Document> documents = documents();
store.add(documents);
for (Document d : documents) {
assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(),
e -> assertThat(e).isNotEmpty());
}
List<Document> results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents().get(0).getId());
assertThat(resultDoc.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNull(),
e -> assertThat(e).isEmpty());
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
@@ -159,17 +153,12 @@ class CassandraVectorStoreIT {
try (CassandraVectorStore store = createTestStore(context, builder)) {
List<Document> documents = documents();
store.add(documents);
for (Document d : documents) {
assertThat(d.getEmbedding()).satisfiesAnyOf(e -> assertThat(e).isNotNull(),
e -> assertThat(e).isNotEmpty());
}
List<Document> results = store.similaritySearch(SearchRequest.query("Spring").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents().get(0).getId());
assertThat(resultDoc.getEmbedding()).isNotEmpty();
assertThat(resultDoc.getContent()).contains(
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");