From 9e74a4327fc5b27be3e02d2f6b6d561d9b4991b4 Mon Sep 17 00:00:00 2001
From: Yi Ding <yi.s.ding@gmail.com>
Date: Thu, 2 May 2024 17:59:36 -0700
Subject: [PATCH] feat: add top k to `asQueryEngine` (#801)
Co-authored-by: Alex Yang <himself65@outlook.com>
---
.changeset/slow-adults-nail.md | 5 +++++
.../core/src/indices/vectorStore/index.ts | 19 +++++++++++++++----
2 files changed, 20 insertions(+), 4 deletions(-)
create mode 100644 .changeset/slow-adults-nail.md
diff --git a/.changeset/slow-adults-nail.md b/.changeset/slow-adults-nail.md
new file mode 100644
index 000000000..b48f661eb
--- /dev/null
+++ 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
index 7450c233d..b81d897af 100644
--- 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,
);
}
--
GitLab