refactor: centralize config generation

run-llama · Apr 2, 2024 · ed03236 · ed03236
1 parent b6e16ea
commit ed03236
Show file tree

Hide file tree

Showing 5 changed files with 127 additions and 110 deletions.
diff --git a/helpers/datasources.ts b/helpers/datasources.ts
@@ -1,6 +1,8 @@
+import fs from "fs/promises";
 import path from "path";
+import yaml, { Document } from "yaml";
 import { templatesDir } from "./dir";
-import { TemplateDataSource } from "./types";
+import { DbSourceConfig, TemplateDataSource, WebSourceConfig } from "./types";
 
 export const EXAMPLE_FILE: TemplateDataSource = {
   type: "file",
@@ -28,3 +30,80 @@ export function getDataSources(
   }
   return dataSources;
 }
+
+export async function writeLoadersConfig(
+  root: string,
+  dataSources: TemplateDataSource[],
+  useLlamaParse?: boolean,
+) {
+  if (dataSources.length === 0) return; // no datasources, no config needed
+  const loaderConfig = new Document({});
+  // Web loader config
+  if (dataSources.some((ds) => ds.type === "web")) {
+    const webLoaderConfig = new Document({});
+
+    // Create config for browser driver arguments
+    const driverArgNodeValue = webLoaderConfig.createNode([
+      "--no-sandbox",
+      "--disable-dev-shm-usage",
+    ]);
+    driverArgNodeValue.commentBefore =
+      " The arguments to pass to the webdriver. E.g.: add --headless to run in headless mode";
+    webLoaderConfig.set("driver_arguments", driverArgNodeValue);
+
+    // Create config for urls
+    const urlConfigs = dataSources
+      .filter((ds) => ds.type === "web")
+      .map((ds) => {
+        const dsConfig = ds.config as WebSourceConfig;
+        return {
+          base_url: dsConfig.baseUrl,
+          prefix: dsConfig.prefix,
+          depth: dsConfig.depth,
+        };
+      });
+    const urlConfigNode = webLoaderConfig.createNode(urlConfigs);
+    urlConfigNode.commentBefore = ` base_url: The URL to start crawling with
+ prefix: Only crawl URLs matching the specified prefix
+ depth: The maximum depth for BFS traversal
+ You can add more websites by adding more entries (don't forget the - prefix from YAML)`;
+    webLoaderConfig.set("urls", urlConfigNode);
+
+    // Add web config to the loaders config
+    loaderConfig.set("web", webLoaderConfig);
+  }
+
+  // File loader config
+  if (dataSources.some((ds) => ds.type === "file")) {
+    // Add documentation to web loader config
+    const node = loaderConfig.createNode({
+      use_llama_parse: useLlamaParse,
+    });
+    node.commentBefore = ` use_llama_parse: Use LlamaParse if \`true\`. Needs a \`LLAMA_CLOUD_API_KEY\` from https://cloud.llamaindex.ai set as environment variable`;
+    loaderConfig.set("file", node);
+  }
+
+  // DB loader config
+  const dbLoaders = dataSources.filter((ds) => ds.type === "db");
+  if (dbLoaders.length > 0) {
+    const dbLoaderConfig = new Document({});
+    const configEntries = dbLoaders.map((ds) => {
+      const dsConfig = ds.config as DbSourceConfig;
+      return {
+        uri: dsConfig.uri,
+        queries: [dsConfig.queries],
+      };
+    });
+
+    const node = dbLoaderConfig.createNode(configEntries);
+    node.commentBefore = ` The configuration for the database loader, only supports MySQL and PostgreSQL databases for now.
+ uri: The URI for the database. E.g.: mysql+pymysql://user:password@localhost:3306/db or postgresql+psycopg2://user:password@localhost:5432/db
+ query: The query to fetch data from the database. E.g.: SELECT * FROM table`;
+    loaderConfig.set("db", node);
+  }
+
+  // Write loaders config
+  const loaderConfigPath = path.join(root, "config", "loaders.yaml");
+  await fs.mkdir(path.join(root, "config"), { recursive: true });
+  await fs.writeFile(loaderConfigPath, yaml.stringify(loaderConfig));
+}
diff --git a/helpers/index.ts b/helpers/index.ts
@@ -4,12 +4,14 @@ import path from "path";
 import { cyan } from "picocolors";
 
 import fsExtra from "fs-extra";
+import { writeLoadersConfig } from "./datasources";
 import { createBackendEnvFile, createFrontendEnvFile } from "./env-variables";
 import { PackageManager } from "./get-pkg-manager";
 import { installLlamapackProject } from "./llama-pack";
 import { isHavingPoetryLockFile, tryPoetryRun } from "./poetry";
 import { installPythonTemplate } from "./python";
 import { downloadAndExtractRepo } from "./repo";
+import { ConfigFileType, writeToolsConfig } from "./tools";
 import {
   FileSourceConfig,
   InstallTemplateArgs,
@@ -121,6 +123,14 @@ export const installTemplate = async (
     await installTSTemplate(props);
   }
 
+  // write configuration files
+  await writeLoadersConfig(props.root, props.dataSources, props.useLlamaParse);
+  await writeToolsConfig(
+    props.root,
+    props.tools,
+    props.framework === "fastapi" ? ConfigFileType.YAML : ConfigFileType.JSON,
+  );
+
   if (props.backend) {
     // This is a backend, so we need to copy the test data and create the env file.
 

diff --git a/helpers/python.ts b/helpers/python.ts
@@ -3,17 +3,15 @@ import path from "path";
 import { cyan, red } from "picocolors";
 import { parse, stringify } from "smol-toml";
 import terminalLink from "terminal-link";
-import yaml, { Document } from "yaml";
+
 import { assetRelocator, copy } from "./copy";
 import { templatesDir } from "./dir";
 import { isPoetryAvailable, tryPoetryInstall } from "./poetry";
 import { Tool } from "./tools";
 import {
-  DbSourceConfig,
   InstallTemplateArgs,
   TemplateDataSource,
   TemplateVectorDB,
-  WebSourceConfig,
 } from "./types";
 
 interface Dependency {
@@ -236,18 +234,14 @@ export const installPythonTemplate = async ({
     cwd: path.join(compPath, "loaders", "python"),
   });
 
-  // write configuration for  loaders
-  await writeLoadersConfig(root, dataSources, useLlamaParse);
-
-  // Select engine code based on data sources and tools
+  // 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 {
     engine = "agent";
-    await writeToolsConfig(root, tools);
   }
 
   // Copy engine code
@@ -270,91 +264,3 @@ export const installPythonTemplate = async ({
     cwd: path.join(compPath, "deployments", "python"),
   });
 };
-
-async function writeToolsConfig(root: string, tools: Tool[]) {
-  if (tools.length === 0) return; // no tools selected, no config need
-  const configContent: Record<string, any> = {};
-  tools.forEach((tool) => {
-    configContent[tool.name] = tool.config ?? {};
-  });
-  const configFilePath = path.join(root, "config", "tools.yaml");
-  await fs.mkdir(path.join(root, "config"), { recursive: true });
-  await fs.writeFile(configFilePath, yaml.stringify(configContent));
-}
-
-async function writeLoadersConfig(
-  root: string,
-  dataSources: TemplateDataSource[],
-  useLlamaParse?: boolean,
-) {
-  if (dataSources.length === 0) return; // no datasources, no config needed
-  const loaderConfig = new Document({});
-  // Web loader config
-  if (dataSources.some((ds) => ds.type === "web")) {
-    const webLoaderConfig = new Document({});
-
-    // Create config for browser driver arguments
-    const driverArgNodeValue = webLoaderConfig.createNode([
-      "--no-sandbox",
-      "--disable-dev-shm-usage",
-    ]);
-    driverArgNodeValue.commentBefore =
-      " The arguments to pass to the webdriver. E.g.: add --headless to run in headless mode";
-    webLoaderConfig.set("driver_arguments", driverArgNodeValue);
-
-    // Create config for urls
-    const urlConfigs = dataSources
-      .filter((ds) => ds.type === "web")
-      .map((ds) => {
-        const dsConfig = ds.config as WebSourceConfig;
-        return {
-          base_url: dsConfig.baseUrl,
-          prefix: dsConfig.prefix,
-          depth: dsConfig.depth,
-        };
-      });
-    const urlConfigNode = webLoaderConfig.createNode(urlConfigs);
-    urlConfigNode.commentBefore = ` base_url: The URL to start crawling with
- prefix: Only crawl URLs matching the specified prefix
- depth: The maximum depth for BFS traversal
- You can add more websites by adding more entries (don't forget the - prefix from YAML)`;
-    webLoaderConfig.set("urls", urlConfigNode);
-
-    // Add web config to the loaders config
-    loaderConfig.set("web", webLoaderConfig);
-  }
-
-  // File loader config
-  if (dataSources.some((ds) => ds.type === "file")) {
-    // Add documentation to web loader config
-    const node = loaderConfig.createNode({
-      use_llama_parse: useLlamaParse,
-    });
-    node.commentBefore = ` use_llama_parse: Use LlamaParse if \`true\`. Needs a \`LLAMA_CLOUD_API_KEY\` from https://cloud.llamaindex.ai set as environment variable`;
-    loaderConfig.set("file", node);
-  }
-
-  // DB loader config
-  const dbLoaders = dataSources.filter((ds) => ds.type === "db");
-  if (dbLoaders.length > 0) {
-    const dbLoaderConfig = new Document({});
-    const configEntries = dbLoaders.map((ds) => {
-      const dsConfig = ds.config as DbSourceConfig;
-      return {
-        uri: dsConfig.uri,
-        queries: [dsConfig.queries],
-      };
-    });
-
-    const node = dbLoaderConfig.createNode(configEntries);
-    node.commentBefore = ` The configuration for the database loader, only supports MySQL and PostgreSQL databases for now.
- uri: The URI for the database. E.g.: mysql+pymysql://user:password@localhost:3306/db or postgresql+psycopg2://user:password@localhost:5432/db
- query: The query to fetch data from the database. E.g.: SELECT * FROM table`;
-    loaderConfig.set("db", node);
-  }
-
-  // Write loaders config
-  const loaderConfigPath = path.join(root, "config/loaders.yaml");
-  await fs.mkdir(path.join(root, "config"), { recursive: true });
-  await fs.writeFile(loaderConfigPath, yaml.stringify(loaderConfig));
-}
diff --git a/helpers/tools.ts b/helpers/tools.ts
@@ -1,4 +1,8 @@
+import fs from "fs/promises";
+import path from "path";
 import { red } from "picocolors";
+import yaml from "yaml";
+import { makeDir } from "./make-dir";
 import { TemplateFramework } from "./types";
 
 export type Tool = {
@@ -8,6 +12,7 @@ export type Tool = {
   dependencies?: ToolDependencies[];
   supportedFrameworks?: Array<TemplateFramework>;
 };
+
 export type ToolDependencies = {
   name: string;
   version?: string;
@@ -73,3 +78,33 @@ export const toolsRequireConfig = (tools?: Tool[]): boolean => {
   }
   return false;
 };
+
+export enum ConfigFileType {
+  YAML = "yaml",
+  JSON = "json",
+}
+
+export const writeToolsConfig = async (
+  root: string,
+  tools: Tool[] = [],
+  type: ConfigFileType = ConfigFileType.YAML,
+) => {
+  if (tools.length === 0) return; // no tools selected, no config need
+  const configContent: Record<string, any> = {};
+  tools.forEach((tool) => {
+    configContent[tool.name] = tool.config ?? {};
+  });
+  const configPath = path.join(root, "config");
+  await makeDir(configPath);
+  if (type === ConfigFileType.YAML) {
+    await fs.writeFile(
+      path.join(configPath, "tools.yaml"),
+      yaml.stringify(configContent),
+    );
+  } else {
+    await fs.writeFile(
+      path.join(configPath, "tools.json"),
+      JSON.stringify(configContent, null, 2),
+    );
+  }
+};
diff --git a/helpers/typescript.ts b/helpers/typescript.ts
@@ -6,7 +6,6 @@ import { assetRelocator, copy } from "../helpers/copy";
 import { callPackageManager } from "../helpers/install";
 import { templatesDir } from "./dir";
 import { PackageManager } from "./get-pkg-manager";
-import { makeDir } from "./make-dir";
 import { InstallTemplateArgs } from "./types";
 
 /**
@@ -133,18 +132,6 @@ export const installTSTemplate = async ({
       parents: true,
       cwd: path.join(compPath, "engines", "typescript", "agent"),
     });
-
-    // Write config/tools.json
-    const configContent: Record<string, any> = {};
-    tools.forEach((tool) => {
-      configContent[tool.name] = tool.config ?? {};
-    });
-    const configPath = path.join(root, "config");
-    await makeDir(configPath);
-    await fs.writeFile(
-      path.join(configPath, "tools.json"),
-      JSON.stringify(configContent, null, 2),
-    );
   } else if (dataSources.length > 0) {
     // use context chat engine if user does not select tools
     console.log("\nUsing context chat engine\n");