Skip to content

Commit

Permalink
Fixing float[] type embedding on MongoDBAtlas vector store
Browse files Browse the repository at this point in the history
  • Loading branch information
sobychacko committed Aug 14, 2024
1 parent d538e00 commit a051978
Showing 1 changed file with 5 additions and 5 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -37,6 +37,7 @@

/**
* @author Chris Smith
* @author Soby Chacko
* @since 1.0.0
*/
public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean {
Expand Down Expand Up @@ -143,14 +144,13 @@ private org.bson.Document createSearchIndexDefinition() {
* @param mongoDocument the mongoDocument to map to a Spring AI Document
* @return the Spring AI Document
*/
private Document mapMongoDocument(org.bson.Document mongoDocument) {
private Document mapMongoDocument(org.bson.Document mongoDocument, float[] queryEmbedding) {
String id = mongoDocument.getString(ID_FIELD_NAME);
String content = mongoDocument.getString(CONTENT_FIELD_NAME);
Map<String, Object> metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class);
List<Float> embedding = mongoDocument.getList(this.config.pathName, Float.class);

Document document = new Document(id, content, metadata);
document.setEmbedding(EmbeddingUtils.toPrimitive(embedding));
document.setEmbedding(queryEmbedding);

return document;
}
Expand Down Expand Up @@ -199,7 +199,7 @@ public List<Document> similaritySearch(SearchRequest request) {
return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, org.bson.Document.class)
.getMappedResults()
.stream()
.map(this::mapMongoDocument)
.map(d -> mapMongoDocument(d, queryEmbedding))
.toList();
}

Expand Down

0 comments on commit a051978

Please sign in to comment.