|
1 | 1 | /*
|
2 |
| - * Copyright 2023 - 2024 the original author or authors. |
| 2 | + * Copyright 2023-2024 the original author or authors. |
3 | 3 | *
|
4 | 4 | * Licensed under the Apache License, Version 2.0 (the "License");
|
5 | 5 | * you may not use this file except in compliance with the License.
|
|
37 | 37 |
|
38 | 38 | /**
|
39 | 39 | * @author Chris Smith
|
| 40 | + * @author Soby Chacko |
40 | 41 | * @since 1.0.0
|
41 | 42 | */
|
42 | 43 | public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean {
|
@@ -143,14 +144,13 @@ private org.bson.Document createSearchIndexDefinition() {
|
143 | 144 | * @param mongoDocument the mongoDocument to map to a Spring AI Document
|
144 | 145 | * @return the Spring AI Document
|
145 | 146 | */
|
146 |
| - private Document mapMongoDocument(org.bson.Document mongoDocument) { |
| 147 | + private Document mapMongoDocument(org.bson.Document mongoDocument, float[] queryEmbedding) { |
147 | 148 | String id = mongoDocument.getString(ID_FIELD_NAME);
|
148 | 149 | String content = mongoDocument.getString(CONTENT_FIELD_NAME);
|
149 | 150 | Map<String, Object> metadata = mongoDocument.get(METADATA_FIELD_NAME, org.bson.Document.class);
|
150 |
| - List<Float> embedding = mongoDocument.getList(this.config.pathName, Float.class); |
151 | 151 |
|
152 | 152 | Document document = new Document(id, content, metadata);
|
153 |
| - document.setEmbedding(EmbeddingUtils.toPrimitive(embedding)); |
| 153 | + document.setEmbedding(queryEmbedding); |
154 | 154 |
|
155 | 155 | return document;
|
156 | 156 | }
|
@@ -199,7 +199,7 @@ public List<Document> similaritySearch(SearchRequest request) {
|
199 | 199 | return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, org.bson.Document.class)
|
200 | 200 | .getMappedResults()
|
201 | 201 | .stream()
|
202 |
| - .map(this::mapMongoDocument) |
| 202 | + .map(d -> mapMongoDocument(d, queryEmbedding)) |
203 | 203 | .toList();
|
204 | 204 | }
|
205 | 205 |
|
|
0 commit comments