From 7c4a53b86baf0ca31528b6a5ea50b9f1ed55e2a3 Mon Sep 17 00:00:00 2001
From: Bailey Pearson <bailey.pearson@mongodb.com>
Date: Fri, 6 Oct 2023 16:45:46 -0600
Subject: [PATCH] misc cleanups

---
 langchain/src/vectorstores/mongodb_atlas.ts | 47 +++++++--------------
 1 file changed, 15 insertions(+), 32 deletions(-)

diff --git a/langchain/src/vectorstores/mongodb_atlas.ts b/langchain/src/vectorstores/mongodb_atlas.ts
index 4138dd21a0b6..2052e2c3c096 100755
--- a/langchain/src/vectorstores/mongodb_atlas.ts
+++ b/langchain/src/vectorstores/mongodb_atlas.ts
@@ -102,19 +102,15 @@ export class MongoDBAtlasVectorSearch extends VectorStore {
     k: number,
     filter?: MongoDBAtlasFilter
   ): Promise<[Document, number][]> {
-    let preFilter: MongoDBDocument | undefined;
-    let postFilterPipeline: MongoDBDocument[] | undefined;
-    let includeEmbeddings: boolean | undefined;
-    if (
-      filter?.preFilter ||
-      filter?.postFilterPipeline ||
-      filter?.includeEmbeddings
-    ) {
-      preFilter = filter.preFilter;
-      postFilterPipeline = filter.postFilterPipeline;
-      includeEmbeddings = filter.includeEmbeddings || false;
-    } else preFilter = filter;
-
+    const postFilterPipeline = filter?.postFilterPipeline ?? [];
+    const preFilter: MongoDBDocument | undefined = filter?.preFilter || filter?.postFilterPipeline || filter?.includeEmbeddings
+      ? filter.preFilter
+      : filter;
+    const removeEmbeddingsPipeline = !filter?.includeEmbeddings ? [{
+      $project: {
+        [this.embeddingKey]: 0,
+      },
+    }] : []
 
     const pipeline: MongoDBDocument[] = [
       {
@@ -132,29 +128,16 @@ export class MongoDBAtlasVectorSearch extends VectorStore {
           score: { $meta: "vectorSearchScore" },
         },
       },
+      ...removeEmbeddingsPipeline,
+      ...postFilterPipeline
     ];
 
-    if (!includeEmbeddings) {
-      const removeEmbeddingsStage = {
-        $project: {
-          [this.embeddingKey]: 0,
-        },
-      };
-      pipeline.push(removeEmbeddingsStage);
-    }
-
-    if (postFilterPipeline) {
-      pipeline.push(...postFilterPipeline);
-    }
-    const results = this.collection.aggregate(pipeline);
-
-    const ret: [Document, number][] = [];
-    for await (const result of results) {
+    const results = this.collection.aggregate(pipeline).map<[Document, number]>((result) => {
       const { score, [this.textKey]: text, ...metadata } = result;
-      ret.push([new Document({ pageContent: text, metadata }), score]);
-    }
+      return [new Document({ pageContent: text, metadata }), score];
+    });
 
-    return ret;
+    return results.toArray();
   }
 
   /**