add top K option

package.json

@@ -20,7 +20,7 @@
     "clsx": "^2.0.0",
     "eslint": "8.45.0",
     "eslint-config-next": "13.4.12",
-    "llamaindex": "^0.0.10",
+    "llamaindex": "^0.0.11",
     "lucide-react": "^0.263.0",
     "next": "13.4.12",
     "postcss": "8.4.27",

pages/api/retrieveandquery.ts

@@ -1,10 +1,16 @@
 // Next.js API route support: https://nextjs.org/docs/api-routes/introduction
 import type { NextApiRequest, NextApiResponse } from "next";
 
-import { IndexDict, TextNode, VectorStoreIndex } from "llamaindex";
+import {
+  IndexDict,
+  RetrieverQueryEngine,
+  TextNode,
+  VectorStoreIndex,
+} from "llamaindex";
 
 type Input = {
   query: string;
+  topK?: number;
   nodesWithEmbedding: {
     text: string;
     embedding: number[];
@@ -27,7 +33,7 @@ export default async function handler(
     return;
   }
 
-  const { query, nodesWithEmbedding }: Input = req.body;
+  const { query, topK, nodesWithEmbedding }: Input = req.body;
 
   const embeddingResults = nodesWithEmbedding.map((config) => {
     return {
@@ -52,7 +58,9 @@ export default async function handler(
   await index.indexStore?.addIndexStruct(indexDict);
   index.indexStruct = indexDict;
 
-  const queryEngine = index.asQueryEngine();
+  const retriever = index.asRetriever();
+  retriever.similarityTopK = topK ?? 2;
+  const queryEngine = new RetrieverQueryEngine(retriever);
   const result = await queryEngine.query(query);
 
   res.status(200).json({ payload: { response: result.response } });

pages/index.tsx

@@ -10,6 +10,7 @@ import { Slider } from "@/components/ui/slider";
 
 const DEFAULT_CHUNK_SIZE = 1024;
 const DEFAULT_CHUNK_OVERLAP = 20;
+const DEFAULT_TOP_K = 2;
 
 export default function Home() {
   const answerId = useId();
@@ -17,6 +18,7 @@ export default function Home() {
   const chunkOverlapId = useId();
   const queryId = useId();
   const sourceId = useId();
+  const topKId = useId();
   const [text, setText] = useState(essay);
   const [query, setQuery] = useState("");
   const [needsNewIndex, setNeedsNewIndex] = useState(true);
@@ -25,6 +27,7 @@ export default function Home() {
   const [nodesWithEmbedding, setNodesWithEmbedding] = useState([]);
   const [chunkSize, setChunkSize] = useState(DEFAULT_CHUNK_SIZE);
   const [chunkOverlap, setChunkOverlap] = useState(DEFAULT_CHUNK_OVERLAP);
+  const [topK, setTopK] = useState(DEFAULT_TOP_K);
   const [answer, setAnswer] = useState("");
 
   return (
@@ -132,6 +135,29 @@ export default function Home() {
         >
           Build Vector Index
         </Button>
+        <div className="my-2">
+          <Label htmlFor={topKId}>Top K:</Label>
+          <div className="flex flex-row space-x-2">
+            <Slider
+              defaultValue={[DEFAULT_TOP_K]}
+              value={[topK]}
+              min={1}
+              max={15}
+              step={1}
+              onValueChange={(values: number[]) => {
+                setTopK(values[0]);
+              }}
+            />
+            <Input
+              id={topKId}
+              value={topK}
+              onChange={(e: ChangeEvent<HTMLInputElement>) => {
+                setTopK(parseInt(e.target.value));
+              }}
+              className="max-w-[100px]"
+            />
+          </div>
+        </div>
         <div className="my-2 space-y-2">
           <Label htmlFor={queryId}>Query:</Label>
           <div className="flex w-full space-x-2">
@@ -156,6 +182,7 @@ export default function Home() {
                   },
                   body: JSON.stringify({
                     query,
+                    topK,
                     nodesWithEmbedding,
                   }),
                 });

pnpm-lock.yaml

@@ -39,8 +39,8 @@ dependencies:
     specifier: 13.4.12
     version: 13.4.12(eslint@8.45.0)(typescript@5.1.6)
   llamaindex:
-    specifier: ^0.0.10
-    version: 0.0.10
+    specifier: ^0.0.11
+    version: 0.0.11
   lucide-react:
     specifier: ^0.263.0
     version: 0.263.0(react@18.2.0)
@@ -2064,8 +2064,8 @@ packages:
     resolution: {integrity: sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg==}
     dev: false
 
-  /llamaindex@0.0.10:
-    resolution: {integrity: sha512-K9C+6+QR6O6pwHPVnDihOo3PiR9n3vRFrcGvP2hQineMXP3mNgqV7ZN0i1k94lJM2MEruXwFSzNlbdncCLBZLA==}
+  /llamaindex@0.0.11:
+    resolution: {integrity: sha512-bdkN6bEgUGsNEy2mn7tIDYtT8+dC/ot9Nq6Wr3S5VK3cqdlJwS87rmqubLO5cqDVc7VBV6cdv1EeYccsvd3tzQ==}
     engines: {node: '>=18.0.0'}
     dependencies:
       js-tiktoken: 1.0.7