Skip to content

Commit a051978

Browse files
committed
Fixing float[] type embedding on MongoDBAtlas vector store
1 parent d538e00 commit a051978

File tree

1 file changed

+5
-5
lines changed

1 file changed

+5
-5
lines changed

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 - 2024 the original author or authors.
2+
* Copyright 2023-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -37,6 +37,7 @@
3737

3838
/**
3939
* @author Chris Smith
40+
* @author Soby Chacko
4041
* @since 1.0.0
4142
*/
4243
public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean {
@@ -143,14 +144,13 @@ private org.bson.Document createSearchIndexDefinition() {
143144
* @param mongoDocument the mongoDocument to map to a Spring AI Document
144145
* @return the Spring AI Document
145146
*/
146-
private Document mapMongoDocument(org.bson.Document mongoDocument) {
147+
private Document mapMongoDocument(org.bson.Document mongoDocument, float[] queryEmbedding) {
147148
String id = mongoDocument.getString(ID_FIELD_NAME);
148149
String content = mongoDocument.getString(CONTENT_FIELD_NAME);
149150
Map<String, Object> metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class);
150-
List<Float> embedding = mongoDocument.getList(this.config.pathName, Float.class);
151151

152152
Document document = new Document(id, content, metadata);
153-
document.setEmbedding(EmbeddingUtils.toPrimitive(embedding));
153+
document.setEmbedding(queryEmbedding);
154154

155155
return document;
156156
}
@@ -199,7 +199,7 @@ public List<Document> similaritySearch(SearchRequest request) {
199199
return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, org.bson.Document.class)
200200
.getMappedResults()
201201
.stream()
202-
.map(this::mapMongoDocument)
202+
.map(d -> mapMongoDocument(d, queryEmbedding))
203203
.toList();
204204
}
205205

0 commit comments

Comments
 (0)