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

6
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java
View File

@@ -66,12 +66,12 @@ class EmbeddingIT extends AbstractIT {
@Test
void embeddingBatchDocuments() throws Exception {
assertThat(this.embeddingModel).isNotNull();
List<float[]> embedded = this.embeddingModel.embed(
List<float[]> embeddings = this.embeddingModel.embed(
List.of(new Document("Hello world"), new Document("Hello Spring"), new Document("Hello Spring AI!")),
OpenAiEmbeddingOptions.builder().withModel(OpenAiApi.DEFAULT_EMBEDDING_MODEL).build(),
new TokenCountBatchingStrategy());
assertThat(embedded.size()).isEqualTo(3);
embedded.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions()));
assertThat(embeddings.size()).isEqualTo(3);
embeddings.forEach(embedding -> assertThat(embedding.length).isEqualTo(this.embeddingModel.dimensions()));
}
@Test

4
spring-ai-core/src/main/java/org/springframework/ai/embedding/BatchingStrategy.java
View File

@@ -31,7 +31,9 @@ public interface BatchingStrategy {
/**
* {@link EmbeddingModel} implementations can call this method to optimize embedding
* tokens. The incoming collection of {@link Document}s are split into su-batches.
* tokens. The incoming collection of {@link Document}s are split into sub-batches. It
* is important to preserve the order of the list of {@link Document}s when batching
* as they are mapped to their corresponding embeddings by their order.
* @param documents to batch
* @return a list of sub-batches that contain {@link Document}s.
*/

14
spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingModel.java
View File

@@ -78,25 +78,23 @@ public interface EmbeddingModel extends Model<EmbeddingRequest, EmbeddingRespons
* @param options {@link EmbeddingOptions}.
* @param batchingStrategy {@link BatchingStrategy}.
* @return a list of float[] that represents the vectors for the incoming
* {@link Document}s.
* {@link Document}s. The returned list is expected to be in the same order of the
* {@link Document} list.
*/
default List<float[]> embed(List<Document> documents, EmbeddingOptions options, BatchingStrategy batchingStrategy) {
Assert.notNull(documents, "Documents must not be null");
List<float[]> embeddings = new ArrayList<>();
List<float[]> embeddings = new ArrayList<>(documents.size());
List<List<Document>> batch = batchingStrategy.batch(documents);
for (List<Document> subBatch : batch) {
List<String> texts = subBatch.stream().map(Document::getContent).toList();
EmbeddingRequest request = new EmbeddingRequest(texts, options);
EmbeddingResponse response = this.call(request);
for (int i = 0; i < subBatch.size(); i++) {
Document document = subBatch.get(i);
float[] output = response.getResults().get(i).getOutput();
embeddings.add(output);
document.setEmbedding(output);
embeddings.add(response.getResults().get(i).getOutput());
}
}
Assert.isTrue(embeddings.size() == documents.size(),
"Embeddings must have the same number as that of the documents");
return embeddings;
}

6
spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
View File

@@ -17,7 +17,7 @@
package org.springframework.ai.embedding;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -139,7 +139,9 @@ public class TokenCountBatchingStrategy implements BatchingStrategy {
List<List<Document>> batches = new ArrayList<>();
int currentSize = 0;
List<Document> currentBatch = new ArrayList<>();
Map<Document, Integer> documentTokens = new HashMap<>();
// Make sure the documentTokens' entry order is preserved by making it a
// LinkedHashMap.
Map<Document, Integer> documentTokens = new LinkedHashMap<>();
for (Document document : documents) {
int tokenCount = this.tokenCountEstimator

7
vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java
View File

@@ -204,13 +204,14 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
public void doAdd(List<Document> documents) {
// Batch the documents based on the batching strategy
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
// Create a list to hold both the CosmosItemOperation and the corresponding
// document ID
List<ImmutablePair<String, CosmosItemOperation>> itemOperationsWithIds = documents.stream().map(doc -> {
CosmosItemOperation operation = CosmosBulkOperations
.getCreateItemOperation(mapCosmosDocument(doc, doc.getEmbedding()), new PartitionKey(doc.getId()));
CosmosItemOperation operation = CosmosBulkOperations.getCreateItemOperation(
mapCosmosDocument(doc, embeddings.get(documents.indexOf(doc))), new PartitionKey(doc.getId()));
return new ImmutablePair<>(doc.getId(), operation); // Pair the document ID
// with the operation
}).toList();

6
vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java
View File

@@ -223,12 +223,13 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
return; // nothing to do;
}
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
final var searchDocuments = documents.stream().map(document -> {
SearchDocument searchDocument = new SearchDocument();
searchDocument.put(ID_FIELD_NAME, document.getId());
searchDocument.put(EMBEDDING_FIELD_NAME, document.getEmbedding());
searchDocument.put(EMBEDDING_FIELD_NAME, embeddings.get(documents.indexOf(document)));
searchDocument.put(CONTENT_FIELD_NAME, document.getContent());
searchDocument.put(METADATA_FIELD_NAME, new JSONObject(document.getMetadata()).toJSONString());
@@ -327,7 +328,6 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
.content(entry.content)
.metadata(metadata)
.score(result.getScore())
.embedding(EmbeddingUtils.toPrimitive(entry.embedding))
.build();
})
.collect(Collectors.toList());

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.");

8
vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java
View File

@@ -145,14 +145,14 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
List<String> contents = new ArrayList<>();
List<float[]> embeddings = new ArrayList<>();
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> documentEmbeddings = this.embeddingModel.embed(documents,
EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
for (Document document : documents) {
ids.add(document.getId());
metadatas.add(document.getMetadata());
contents.add(document.getContent());
document.setEmbedding(document.getEmbedding());
embeddings.add(document.getEmbedding());
embeddings.add(documentEmbeddings.get(documents.indexOf(document)));
}
this.chromaApi.upsertEmbeddings(this.collectionId,
@@ -192,12 +192,12 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
if (metadata == null) {
metadata = new HashMap<>();
}
metadata.put(DocumentMetadata.DISTANCE.value(), distance);
Document document = Document.builder()
.id(id)
.content(content)
.metadata(metadata)
.embedding(chromaEmbedding.embedding())
.score(1.0 - distance)
.build();
responseDocuments.add(document);

4
vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/CoherenceVectorStore.java
View File

@@ -167,9 +167,9 @@ public class CoherenceVectorStore implements VectorStore, InitializingBean {
public void add(final List<Document> documents) {
Map<DocumentChunk.Id, DocumentChunk> chunks = new HashMap<>((int) Math.ceil(documents.size() / 0.75f));
for (Document doc : documents) {
doc.setEmbedding(this.embeddingModel.embed(doc));
var id = toChunkId(doc.getId());
var chunk = new DocumentChunk(doc.getContent(), doc.getMetadata(), toFloat32Vector(doc.getEmbedding()));
var chunk = new DocumentChunk(doc.getContent(), doc.getMetadata(),
toFloat32Vector(this.embeddingModel.embed(doc)));
chunks.put(id, chunk);
}
this.documentChunks.putAll(chunks);

21
vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/ElasticsearchVectorStore.java
View File

@@ -72,6 +72,7 @@ import org.springframework.util.Assert;
* @author Soby Chacko
* @author Christian Tzolov
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
* @since 1.0.0
*/
public class ElasticsearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -133,11 +134,14 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
}
BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
bulkRequestBuilder.operations(op -> op
.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(document)));
ElasticSearchDocument doc = new ElasticSearchDocument(document.getId(), document.getContent(),
document.getMetadata(), embeddings.get(documents.indexOf(document)));
bulkRequestBuilder.operations(
op -> op.index(idx -> idx.index(this.options.getIndexName()).id(document.getId()).document(doc)));
}
BulkResponse bulkRequest = bulkRequest(bulkRequestBuilder.build());
if (bulkRequest.errors()) {
@@ -282,4 +286,15 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
return SIMILARITY_TYPE_MAPPING.get(this.options.getSimilarity()).value();
}
/**
* The representation of {@link Document} along with its embedding.
*
* @param id The id of the document
* @param content The content of the document
* @param metadata The metadata of the document
* @param embedding The vectors representing the content of the document
*/
public record ElasticSearchDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
}
}

7
vector-stores/spring-ai-gemfire-store/src/main/java/org/springframework/ai/vectorstore/GemFireVectorStore.java
View File

@@ -208,10 +208,11 @@ public class GemFireVectorStore extends AbstractObservationVectorStore implement
@Override
public void doAdd(List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
UploadRequest upload = new UploadRequest(documents.stream()
.map(document -> new UploadRequest.Embedding(document.getId(), document.getEmbedding(), DOCUMENT_FIELD,
document.getContent(), document.getMetadata()))
.map(document -> new UploadRequest.Embedding(document.getId(), embeddings.get(documents.indexOf(document)),
DOCUMENT_FIELD, document.getContent(), document.getMetadata()))
.toList());
String embeddingsJson = null;

5
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
View File

@@ -161,7 +161,8 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
List<List<Float>> embeddingArray = new ArrayList<>();
// TODO: Need to customize how we pass the embedding options
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
docIdArray.add(document.getId());
@@ -169,7 +170,7 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
// the content used to compute the embeddings
contentArray.add(document.getContent());
metadataArray.add(new JSONObject(document.getMetadata()));
embeddingArray.add(EmbeddingUtils.toList(document.getEmbedding()));
embeddingArray.add(EmbeddingUtils.toList(embeddings.get(documents.indexOf(document))));
}
List<InsertParam.Field> fields = new ArrayList<>();

24
vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java
View File

@@ -51,6 +51,7 @@ import org.springframework.util.Assert;
* @author Soby Chacko
* @author Christian Tzolov
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
* @since 1.0.0
*/
public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -175,22 +176,26 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
String content = mongoDocument.getString(CONTENT_FIELD_NAME);
double score = mongoDocument.getDouble(SCORE_FIELD_NAME);
Map<String, Object> metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class);
metadata.put(DocumentMetadata.DISTANCE.value(), 1 - score);
// @formatter:off
return Document.builder()
.id(id)
.content(content)
.metadata(metadata)
.score(score)
.embedding(queryEmbedding)
.build();
.build(); // @formatter:on
}
@Override
public void doAdd(List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
this.mongoTemplate.save(document, this.config.collectionName);
MongoDBDocument mdbDocument = new MongoDBDocument(document.getId(), document.getContent(),
document.getMetadata(), embeddings.get(documents.indexOf(document)));
this.mongoTemplate.save(mdbDocument, this.config.collectionName);
}
}
@@ -338,4 +343,15 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
}
/**
* The representation of {@link Document} along with its embedding.
*
* @param id The id of the document
* @param content The content of the document
* @param metadata The metadata of the document
* @param embedding The vectors representing the content of the document
*/
public record MongoDBDocument(String id, String content, Map<String, Object> metadata, float[] embedding) {
}
}

12
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java
View File

@@ -108,9 +108,12 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
@Override
public void doAdd(List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
var rows = documents.stream().map(this::documentToRecord).toList();
var rows = documents.stream()
.map(document -> documentToRecord(document, embeddings.get(documents.indexOf(document))))
.toList();
try (var session = this.driver.session()) {
var statement = """
@@ -204,8 +207,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
}
}
private Map<String, Object> documentToRecord(Document document) {
document.setEmbedding(document.getEmbedding());
private Map<String, Object> documentToRecord(Document document, float[] embedding) {
var row = new HashMap<String, Object>();
@@ -217,7 +219,7 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
document.getMetadata().forEach((k, v) -> properties.put("metadata." + k, Values.value(v)));
row.put("properties", properties);
row.put(this.config.embeddingProperty, Values.value(document.getEmbedding()));
row.put(this.config.embeddingProperty, Values.value(embedding));
return row;
}

7
vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/OracleVectorStore.java
View File

@@ -206,7 +206,8 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
@Override
public void doAdd(final List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
this.jdbcTemplate.batchUpdate(getIngestStatement(), new BatchPreparedStatementSetter() {
@Override
@@ -214,7 +215,7 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
final Document document = documents.get(i);
final String content = document.getContent();
final byte[] json = toJson(document.getMetadata());
final VECTOR embeddingVector = toVECTOR(document.getEmbedding());
final VECTOR embeddingVector = toVECTOR(embeddings.get(documents.indexOf(document)));
org.springframework.jdbc.core.StatementCreatorUtils.setParameterValue(ps, 1, Types.VARCHAR,
document.getId());
@@ -653,13 +654,11 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
final Map<String, Object> metadata = getMap(rs.getObject(3, OracleJsonValue.class));
metadata.put(DocumentMetadata.DISTANCE.value(), rs.getDouble(5));
final float[] embedding = rs.getObject(4, float[].class);
return Document.builder()
.id(rs.getString(1))
.content(rs.getString(2))
.metadata(metadata)
.score(1 - rs.getDouble(5))
.embedding(embedding)
.build();
}

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

@@ -196,10 +196,11 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
@Override
public void doAdd(List<Document> documents) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
List<List<Document>> batchedDocuments = batchDocuments(documents);
batchedDocuments.forEach(this::insertOrUpdateBatch);
batchedDocuments.forEach(batchDocument -> insertOrUpdateBatch(batchDocument, documents, embeddings));
}
private List<List<Document>> batchDocuments(List<Document> documents) {
@@ -210,7 +211,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return batches;
}
private void insertOrUpdateBatch(List<Document> batch) {
private void insertOrUpdateBatch(List<Document> batch, List<Document> documents, List<float[]> embeddings) {
String sql = "INSERT INTO " + getFullyQualifiedTableName()
+ " (id, content, metadata, embedding) VALUES (?, ?, ?::jsonb, ?) " + "ON CONFLICT (id) DO "
+ "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? ";
@@ -223,7 +224,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
var document = batch.get(i);
var content = document.getContent();
var json = toJson(document.getMetadata());
var embedding = document.getEmbedding();
var embedding = embeddings.get(documents.indexOf(document));
var pGvector = new PGvector(embedding);
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN,
@@ -499,23 +500,18 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
String id = rs.getString(COLUMN_ID);
String content = rs.getString(COLUMN_CONTENT);
PGobject pgMetadata = rs.getObject(COLUMN_METADATA, PGobject.class);
PGobject embedding = rs.getObject(COLUMN_EMBEDDING, PGobject.class);
Float distance = rs.getFloat(COLUMN_DISTANCE);
Map<String, Object> metadata = toMap(pgMetadata);
metadata.put(DocumentMetadata.DISTANCE.value(), distance);
// @formatter:off
return Document.builder()
.id(id)
.content(content)
.metadata(metadata)
.score(1.0 - distance)
.embedding(toFloatArray(embedding))
.build();
}
private float[] toFloatArray(PGobject embedding) throws SQLException {
return new PGvector(embedding.getValue()).toArray();
.build(); // @formatter:on
}
private Map<String, Object> toMap(PGobject pgObject) {

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

@@ -146,9 +146,6 @@ class PgVectorStoreWithChatMemoryAdvisorIT {
EmbeddingModel embeddingModel = mock(EmbeddingModel.class);
Mockito.doAnswer(invocationOnMock -> {
Object[] arguments = invocationOnMock.getArguments();
List<Document> documents = (List<Document>) arguments[0];
documents.forEach(d -> d.setEmbedding(this.embed));
return List.of(this.embed, this.embed);
}).when(embeddingModel).embed(ArgumentMatchers.any(), any(), any());
given(embeddingModel.embed(any(String.class))).willReturn(this.embed);

5
vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/PineconeVectorStore.java
View File

@@ -124,11 +124,12 @@ public class PineconeVectorStore extends AbstractObservationVectorStore {
* @param namespace The namespace to add the documents to
*/
public void add(List<Document> documents, String namespace) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
List<Vector> upsertVectors = documents.stream()
.map(document -> Vector.newBuilder()
.setId(document.getId())
.addAllValues(EmbeddingUtils.toList(document.getEmbedding()))
.addAllValues(EmbeddingUtils.toList(embeddings.get(documents.indexOf(document))))
.setMetadata(metadataToStruct(document))
.build())
.toList();

5
vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java
View File

@@ -127,12 +127,13 @@ public class QdrantVectorStore extends AbstractObservationVectorStore implements
try {
// Compute and assign an embedding to the document.
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
List<PointStruct> points = documents.stream()
.map(document -> PointStruct.newBuilder()
.setId(io.qdrant.client.PointIdFactory.id(UUID.fromString(document.getId())))
.setVectors(io.qdrant.client.VectorsFactory.vectors(document.getEmbedding()))
.setVectors(io.qdrant.client.VectorsFactory.vectors(embeddings.get(documents.indexOf(document))))
.putAllPayload(toPayload(document))
.build())
.toList();

6
vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java
View File

@@ -163,12 +163,12 @@ public class RedisVectorStore extends AbstractObservationVectorStore implements
public void doAdd(List<Document> documents) {
try (Pipeline pipeline = this.jedis.pipelined()) {
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
for (Document document : documents) {
document.setEmbedding(document.getEmbedding());
var fields = new HashMap<String, Object>();
fields.put(this.config.embeddingFieldName, document.getEmbedding());
fields.put(this.config.embeddingFieldName, embeddings.get(documents.indexOf(document)));
fields.put(this.config.contentFieldName, document.getContent());
fields.putAll(document.getMetadata());
pipeline.jsonSetWithEscape(key(document.getId()), JSON_SET_PATH, fields);

5
vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/TypesenseVectorStore.java
View File

@@ -125,14 +125,15 @@ public class TypesenseVectorStore extends AbstractObservationVectorStore impleme
public void doAdd(List<Document> documents) {
Assert.notNull(documents, "Documents must not be null");
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
List<HashMap<String, Object>> documentList = documents.stream().map(document -> {
HashMap<String, Object> typesenseDoc = new HashMap<>();
typesenseDoc.put(DOC_ID_FIELD_NAME, document.getId());
typesenseDoc.put(CONTENT_FIELD_NAME, document.getContent());
typesenseDoc.put(METADATA_FIELD_NAME, document.getMetadata());
typesenseDoc.put(EMBEDDING_FIELD_NAME, document.getEmbedding());
typesenseDoc.put(EMBEDDING_FIELD_NAME, embeddings.get(documents.indexOf(document)));
return typesenseDoc;
}).toList();

16
vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/WeaviateVectorStore.java
View File

@@ -197,9 +197,12 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
return;
}
this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(), this.batchingStrategy);
List<float[]> embeddings = this.embeddingModel.embed(documents, EmbeddingOptionsBuilder.builder().build(),
this.batchingStrategy);
List<WeaviateObject> weaviateObjects = documents.stream().map(this::toWeaviateObject).toList();
List<WeaviateObject> weaviateObjects = documents.stream()
.map(document -> toWeaviateObject(document, documents, embeddings))
.toList();
Result<ObjectGetResponse[]> response = this.weaviateClient.batch()
.objectsBatcher()
@@ -235,7 +238,7 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
}
}
private WeaviateObject toWeaviateObject(Document document) {
private WeaviateObject toWeaviateObject(Document document, List<Document> documents, List<float[]> embeddings) {
// https://weaviate.io/developers/weaviate/config-refs/datatypes
Map<String, Object> fields = new HashMap<>();
@@ -259,7 +262,7 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
return WeaviateObject.builder()
.className(this.weaviateObjectClass)
.id(document.getId())
.vector(EmbeddingUtils.toFloatArray(document.getEmbedding()))
.vector(EmbeddingUtils.toFloatArray(embeddings.get(documents.indexOf(document))))
.properties(fields)
.build();
}
@@ -363,7 +366,6 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
Map<String, ?> additional = (Map<String, ?>) item.get(ADDITIONAL_FIELD_NAME);
double certainty = (Double) additional.get(ADDITIONAL_CERTAINTY_FIELD_NAME);
String id = (String) additional.get(ADDITIONAL_ID_FIELD_NAME);
List<Double> embedding = ((List<Double>) additional.get(ADDITIONAL_VECTOR_FIELD_NAME)).stream().toList();
// Metadata
Map<String, Object> metadata = new HashMap<>();
@@ -382,13 +384,13 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
// Content
String content = (String) item.get(CONTENT_FIELD_NAME);
// @formatter:off
return Document.builder()
.id(id)
.content(content)
.metadata(metadata)
.embedding(EmbeddingUtils.toPrimitive(EmbeddingUtils.doubleToFloat(embedding)))
.score(certainty)
.build();
.build(); // @formatter:on
}
@Override