support $vectorSearch

langchain-ai · Oct 6, 2023 · 43dd522 · 43dd522
1 parent 1a1c75d
commit 43dd522
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 14 deletions.
diff --git a/langchain/src/vectorstores/mongodb_atlas.ts b/langchain/src/vectorstores/mongodb_atlas.ts
@@ -102,12 +102,6 @@ export class MongoDBAtlasVectorSearch extends VectorStore {
     k: number,
     filter?: MongoDBAtlasFilter
   ): Promise<[Document, number][]> {
-    const knnBeta: MongoDBDocument = {
-      vector: query,
-      path: this.embeddingKey,
-      k,
-    };
-
     let preFilter: MongoDBDocument | undefined;
     let postFilterPipeline: MongoDBDocument[] | undefined;
     let includeEmbeddings: boolean | undefined;
@@ -121,19 +115,21 @@ export class MongoDBAtlasVectorSearch extends VectorStore {
       includeEmbeddings = filter.includeEmbeddings || false;
     } else preFilter = filter;
 
-    if (preFilter) {
-      knnBeta.filter = preFilter;
-    }
+
     const pipeline: MongoDBDocument[] = [
       {
-        $search: {
+        $vectorSearch: {
+          queryVector: query,
           index: this.indexName,
-          knnBeta,
-        },
+          path: this.embeddingKey,
+          limit: k,
+          numCandidates: 10 * k,
+          ...(preFilter && { filter: preFilter })
+        }
       },
       {
         $set: {
-          score: { $meta: "searchScore" },
+          score: { $meta: "vectorSearchScore" },
         },
       },
     ];

diff --git a/langchain/src/vectorstores/tests/mongodb_atlas.int.test.ts b/langchain/src/vectorstores/tests/mongodb_atlas.int.test.ts
@@ -16,6 +16,7 @@ import { OpenAIEmbeddings } from "../../embeddings/openai.js";
 {
   "mappings": {
     "fields": {
+      "e": { "type": "number" },
       "embedding": {
         "dimensions": 1536,
         "similarity": "euclidean",
@@ -67,7 +68,7 @@ test("MongoDBAtlasVectorSearch with external ids", async () => {
 
     // we can pre filter the search
     const preFilter = {
-      range: { lte: 1, path: "e" },
+      e: { $lte: 1 },
     };
 
     const filteredResults = await vectorStore.similaritySearch(