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

Updated MongoDBAtlasVectorStore to use org.bson.Documents (#1026)

Browse Source 
org.bson.Documents have been the preferred document structure
since the 3.0 release of the mongo-java-driver
This commit is contained in:
Ross Lawley
2024-07-24 22:03:42 +01:00
committed by GitHub
parent 8936aac1e2
commit 04191a77e0
1 changed files with 10 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -21,8 +21,6 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import com.mongodb.BasicDBObject;
import com.mongodb.MongoCommandException;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -140,16 +138,15 @@ public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean {
}
/**
* Maps a BasicDBObject to a Spring AI Document
* @param basicDBObject the basicDBObject to map to a spring ai document
* @return the spring ai document
* Maps a Bson Document to a Spring AI Document
* @param mongoDocument the mongoDocument to map to a Spring AI Document
* @return the Spring AI Document
*/
@SuppressWarnings("unchecked")
private Document mapBasicDbObject(BasicDBObject basicDBObject) {
String id = basicDBObject.getString(ID_FIELD_NAME);
String content = basicDBObject.getString(CONTENT_FIELD_NAME);
Map<String, Object> metadata = (Map<String, Object>) basicDBObject.get(METADATA_FIELD_NAME);
List<Double> embedding = (List<Double>) basicDBObject.get(this.config.pathName);
private Document mapMongoDocument(org.bson.Document mongoDocument) {
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<Double> embedding = mongoDocument.getList(this.config.pathName, Double.class);
Document document = new Document(id, content, metadata);
document.setEmbedding(embedding);
@@ -198,10 +195,10 @@ public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean {
.build(),
Aggregation.match(new Criteria(SCORE_FIELD_NAME).gte(request.getSimilarityThreshold())));
return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, BasicDBObject.class)
return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, org.bson.Document.class)
.getMappedResults()
.stream()
.map(this::mapBasicDbObject)
.map(this::mapMongoDocument)
.toList();
}