feat: add filter for query in ts templates

run-llama · Jul 18, 2024 · 769717e · 769717e
1 parent 8a5f098
commit 769717e
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 10 deletions.
diff --git a/templates/components/engines/typescript/agent/chat.ts b/templates/components/engines/typescript/agent/chat.ts
@@ -1,4 +1,9 @@
-import { BaseToolWithCall, OpenAIAgent, QueryEngineTool } from "llamaindex";
+import {
+  BaseToolWithCall,
+  MetadataFilters,
+  OpenAIAgent,
+  QueryEngineTool,
+} from "llamaindex";
 import fs from "node:fs/promises";
 import path from "node:path";
 import { getDataSource } from "./index";
@@ -14,7 +19,7 @@ export async function createChatEngine(documentIds?: string[]) {
     tools.push(
       new QueryEngineTool({
         queryEngine: index.asQueryEngine({
-          preFilters: undefined, // TODO: Add filters once LITS supports it (getQueryFilters)
+          preFilters: generateFilters(documentIds || []),
         }),
         metadata: {
           name: "data_query_engine",
@@ -41,3 +46,32 @@ export async function createChatEngine(documentIds?: string[]) {
     systemPrompt: process.env.SYSTEM_PROMPT,
   });
 }
+
+function generateFilters(documentIds: string[]): MetadataFilters | undefined {
+  if (!documentIds.length) {
+    return {
+      filters: [
+        {
+          key: "private",
+          value: ["true"],
+          operator: "nin",
+        },
+      ],
+    };
+  }
+  return {
+    filters: [
+      {
+        key: "private",
+        value: "true",
+        operator: "!=",
+      },
+      {
+        key: "doc_id",
+        value: documentIds,
+        operator: "in",
+      },
+    ],
+    condition: "or",
+  };
+}
diff --git a/templates/components/llamaindex/typescript/documents/pipeline.ts b/templates/components/llamaindex/typescript/documents/pipeline.ts
@@ -18,6 +18,7 @@ export async function runPipeline(documents: Document[], filename: string) {
   for (const document of documents) {
     document.metadata = {
       ...document.metadata,
+      doc_id: document.id_,
       file_name: filename,
       private: "true", // to separate from other public documents
       is_local_file: "true",

diff --git a/templates/components/llamaindex/typescript/streaming/annotations.ts b/templates/components/llamaindex/typescript/streaming/annotations.ts
@@ -35,9 +35,9 @@ export function retrieveDocumentIds(annotations?: JSONValue[]): string[] {
     ) {
       const files = data.files as DocumentFile[];
       for (const file of files) {
-        if (Array.isArray(file.content)) {
+        if (Array.isArray(file.content.value)) {
           // it's an array, so it's an array of doc IDs
-          for (const id of file.content) {
+          for (const id of file.content.value) {
             ids.push(id);
           }
         }

diff --git a/templates/components/llamaindex/typescript/streaming/events.ts b/templates/components/llamaindex/typescript/streaming/events.ts
@@ -77,7 +77,7 @@ export function createCallbackManager(stream: StreamData) {
   const callbackManager = new CallbackManager();
 
   callbackManager.on("retrieve-end", async (data) => {
-    const { nodes, query } = data.detail.payload;
+    const { nodes, query } = data.detail;
     await appendSourceData(stream, nodes);
     appendEventData(stream, `Retrieving context for query: '${query}'`);
     appendEventData(
@@ -87,7 +87,7 @@ export function createCallbackManager(stream: StreamData) {
   });
 
   callbackManager.on("llm-tool-call", (event) => {
-    const { name, input } = event.detail.payload.toolCall;
+    const { name, input } = event.detail.toolCall;
     const inputString = Object.entries(input)
       .map(([key, value]) => `${key}: ${value}`)
       .join(", ");
@@ -98,7 +98,7 @@ export function createCallbackManager(stream: StreamData) {
   });
 
   callbackManager.on("llm-tool-result", (event) => {
-    const { toolCall, toolResult } = event.detail.payload;
+    const { toolCall, toolResult } = event.detail;
     appendToolData(stream, toolCall, toolResult);
   });
 

diff --git a/templates/types/streaming/express/package.json b/templates/types/streaming/express/package.json
@@ -20,7 +20,7 @@
     "dotenv": "^16.3.1",
     "duck-duck-scrape": "^2.2.5",
     "express": "^4.18.2",
-    "llamaindex": "0.4.14",
+    "llamaindex": "0.5.5",
     "pdf2json": "3.0.5",
     "ajv": "^8.12.0",
     "@e2b/code-interpreter": "^0.0.5",

diff --git a/templates/types/streaming/nextjs/next.config.mjs b/templates/types/streaming/nextjs/next.config.mjs
@@ -1,8 +1,9 @@
 /** @type {import('next').NextConfig} */
 import fs from "fs";
+import withLlamaIndex from "llamaindex/next";
 import webpack from "./webpack.config.mjs";
 
 const nextConfig = JSON.parse(fs.readFileSync("./next.config.json", "utf-8"));
 nextConfig.webpack = webpack;
 
-export default nextConfig;
+export default withLlamaIndex(nextConfig);
diff --git a/templates/types/streaming/nextjs/package.json b/templates/types/streaming/nextjs/package.json
@@ -24,7 +24,7 @@
     "duck-duck-scrape": "^2.2.5",
     "formdata-node": "^6.0.3",
     "got": "^14.4.1",
-    "llamaindex": "0.4.14",
+    "llamaindex": "0.5.5",
     "lucide-react": "^0.294.0",
     "next": "^14.2.4",
     "react": "^18.2.0",