fix: simplify agents selection for typescript

run-llama · Apr 3, 2024 · 3c8a09f · 3c8a09f
1 parent ed03236
commit 3c8a09f
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 58 deletions.
diff --git a/helpers/python.ts b/helpers/python.ts
@@ -243,8 +243,6 @@ export const installPythonTemplate = async ({
   } else {
     engine = "agent";
   }
-
-  // Copy engine code
   await copy("**", enginePath, {
     parents: true,
     cwd: path.join(compPath, "engines", "python", engine),

diff --git a/helpers/typescript.ts b/helpers/typescript.ts
@@ -104,46 +104,33 @@ export const installTSTemplate = async ({
       : path.join("src", "controllers");
   const enginePath = path.join(root, relativeEngineDestPath, "engine");
 
-  if (dataSources.length > 0) {
-    if (vectorDb) {
-      // copy vector db component
-      console.log("\nUsing vector DB:", vectorDb, "\n");
-      await copy("**", enginePath, {
-        parents: true,
-        cwd: path.join(compPath, "vectordbs", "typescript", vectorDb),
-      });
-    }
-    // copy loader component (TS only supports llama_parse and file for now)
-    const loaderFolder = useLlamaParse ? "llama_parse" : "file";
-    await copy("**", enginePath, {
-      parents: true,
-      cwd: path.join(compPath, "loaders", "typescript", loaderFolder),
-    });
-  }
+  // copy vector db component
+  console.log("\nUsing vector DB:", vectorDb, "\n");
+  await copy("**", enginePath, {
+    parents: true,
+    cwd: path.join(compPath, "vectordbs", "typescript", vectorDb ?? "none"),
+  });
 
-  /**
-   * Selected chat engine and copy necessary files
-   */
-  // copy engine
-  if (tools && tools.length > 0) {
-    // use agent chat engine if user selects tools
-    console.log("\nUsing agent chat engine\n");
-    await copy("**", enginePath, {
-      parents: true,
-      cwd: path.join(compPath, "engines", "typescript", "agent"),
-    });
-  } else if (dataSources.length > 0) {
-    // use context chat engine if user does not select tools
-    console.log("\nUsing context chat engine\n");
-    await copy("**", enginePath, {
-      parents: true,
-      cwd: path.join(compPath, "engines", "typescript", "chat"),
-    });
+  // copy loader component (TS only supports llama_parse and file for now)
+  const loaderFolder = useLlamaParse ? "llama_parse" : "file";
+  await copy("**", enginePath, {
+    parents: true,
+    cwd: path.join(compPath, "loaders", "typescript", loaderFolder),
+  });
+
+  // Select and copy engine code based on data sources and tools
+  let engine;
+  tools = tools ?? [];
+  if (dataSources.length > 0 && tools.length === 0) {
+    console.log("\nNo tools selected - use optimized context chat engine\n");
+    engine = "chat";
   } else {
-    console.log(
-      "\nUsing simple chat as neither a datasource nor tools are selected\n",
-    );
+    engine = "agent";
   }
+  await copy("**", enginePath, {
+    parents: true,
+    cwd: path.join(compPath, "engines", "typescript", engine),
+  });
 
   /**
    * Copy the selected UI files to the target directory and reference it.

diff --git a/templates/components/engines/typescript/agent/chat.ts b/templates/components/engines/typescript/agent/chat.ts
@@ -1,26 +1,44 @@
-import config from "@/config/tools.json";
-import { OpenAI, OpenAIAgent, QueryEngineTool, ToolFactory } from "llamaindex";
+import {
+  BaseTool,
+  OpenAI,
+  OpenAIAgent,
+  QueryEngineTool,
+  ToolFactory,
+} from "llamaindex";
+import fs from "node:fs/promises";
+import path from "node:path";
 import { STORAGE_CACHE_DIR } from "./constants.mjs";
 import { getDataSource } from "./index";
 
 export async function createChatEngine(llm: OpenAI) {
+  let tools: BaseTool[] = [];
+
+  // Add a query engine tool if we have a data source
+  // Delete this code if you don't have a data source
   const index = await getDataSource(llm);
-  const queryEngine = index.asQueryEngine();
-  const queryEngineTool = new QueryEngineTool({
-    queryEngine: queryEngine,
-    metadata: {
-      name: "data_query_engine",
-      description: `A query engine for documents in storage folder: ${STORAGE_CACHE_DIR}`,
-    },
-  });
+  if (index) {
+    tools.push(
+      new QueryEngineTool({
+        queryEngine: index.asQueryEngine(),
+        metadata: {
+          name: "data_query_engine",
+          description: `A query engine for documents in storage folder: ${STORAGE_CACHE_DIR}`,
+        },
+      }),
+    );
+  }
 
-  const externalTools = await ToolFactory.createTools(config);
+  try {
+    // add tools from config file if it exists
+    const config = JSON.parse(
+      await fs.readFile(path.join("config", "tools.json"), "utf8"),
+    );
+    tools = tools.concat(await ToolFactory.createTools(config));
+  } catch {}
 
-  const agent = new OpenAIAgent({
-    tools: [queryEngineTool, ...externalTools],
-    verbose: true,
+  return new OpenAIAgent({
+    tools,
     llm,
+    verbose: true,
   });
-
-  return agent;
 }
diff --git a/templates/components/engines/typescript/chat/chat.ts b/templates/components/engines/typescript/chat/chat.ts
@@ -3,6 +3,11 @@ import { getDataSource } from "./index";
 
 export async function createChatEngine(llm: LLM) {
   const index = await getDataSource(llm);
+  if (!index) {
+    throw new Error(
+      `StorageContext is empty - call 'npm run generate' to generate the storage first`,
+    );
+  }
   const retriever = index.asRetriever();
   retriever.similarityTopK = 3;
 

diff --git a/templates/components/vectordbs/typescript/none/index.ts b/templates/components/vectordbs/typescript/none/index.ts
@@ -21,9 +21,7 @@ export async function getDataSource(llm: LLM) {
     (storageContext.docStore as SimpleDocumentStore).toDict(),
   ).length;
   if (numberOfDocs === 0) {
-    throw new Error(
-      `StorageContext is empty - call 'npm run generate' to generate the storage first`,
-    );
+    return null;
   }
   return await VectorStoreIndex.init({
     storageContext,