feat: add top k to asQueryEngine (#801)

Co-authored-by: Alex Yang <[email protected]>
run-llama · May 3, 2024 · 9e74a43 · 9e74a43
1 parent 5e61934
commit 9e74a43
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/.changeset/slow-adults-nail.md b/.changeset/slow-adults-nail.md
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: add top k to `asQueryEngine`
diff --git a/packages/core/src/indices/vectorStore/index.ts b/packages/core/src/indices/vectorStore/index.ts
@@ -279,18 +279,29 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     return new VectorIndexRetriever({ index: this, ...options });
   }
 
+  /**
+   * Create a RetrieverQueryEngine.
+   * similarityTopK is only used if no existing retriever is provided.
+   */
   asQueryEngine(options?: {
     retriever?: BaseRetriever;
     responseSynthesizer?: BaseSynthesizer;
     preFilters?: MetadataFilters;
     nodePostprocessors?: BaseNodePostprocessor[];
+    similarityTopK?: number;
   }): QueryEngine & RetrieverQueryEngine {
-    const { retriever, responseSynthesizer } = options ?? {};
+    const {
+      retriever,
+      responseSynthesizer,
+      preFilters,
+      nodePostprocessors,
+      similarityTopK,
+    } = options ?? {};
     return new RetrieverQueryEngine(
-      retriever ?? this.asRetriever(),
+      retriever ?? this.asRetriever({ similarityTopK }),
       responseSynthesizer,
-      options?.preFilters,
-      options?.nodePostprocessors,
+      preFilters,
+      nodePostprocessors,
     );
   }