Merge pull request #14429 from jkorach/gh-14428

Add TypeScript interface for the new PipelineStage - Vector Search - solving issue #14428
Automattic · Mar 14, 2024 · 06067de · 06067de
2 parents 3ba9778 + 580bc8c
commit 06067de
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/test/types/PipelineStage.test.ts b/test/types/PipelineStage.test.ts
@@ -520,3 +520,20 @@ function gh12269() {
     }
   };
 }
+const vectorSearchStages: PipelineStage[] = [
+  {
+    $vectorSearch: {
+      index: 'title_vector_index',
+      path: 'embedding',
+      queryVector: [0.522, 0.123, 0.487],
+      limit: 5,
+      numCandidates: 100
+    }
+  },
+  {
+    $project: {
+      title: 1,
+      score: { $meta: 'searchScore' }
+    }
+  }
+];
diff --git a/types/pipelinestage.d.ts b/types/pipelinestage.d.ts
@@ -36,7 +36,8 @@ declare module 'mongoose' {
     | PipelineStage.SortByCount
     | PipelineStage.UnionWith
     | PipelineStage.Unset
-    | PipelineStage.Unwind;
+    | PipelineStage.Unwind
+    | PipelineStage.VectorSearch;
 
   export namespace PipelineStage {
     export interface AddFields {
@@ -308,5 +309,17 @@ declare module 'mongoose' {
       /** [`$unwind` reference](https://www.mongodb.com/docs/manual/reference/operator/aggregation/unwind/) */
       $unwind: string | { path: string; includeArrayIndex?: string; preserveNullAndEmptyArrays?: boolean }
     }
+    export interface VectorSearch {
+      /** [`$vectorSearch` reference](https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/) */
+      $vectorSearch: {
+        index: string,
+        path: string,
+        queryVector: number[],
+        numCandidates: number,
+        limit: number,
+        filter?: Expression,
+      }
+    }
+
   }
 }