feat: Add support for llamahub tools (#517)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

create-app.ts

 /* eslint-disable import/no-extraneous-dependencies */
 import path from "path";
-import { green } from "picocolors";
+import { green, yellow } from "picocolors";
 import { tryGitInit } from "./helpers/git";
 import { isFolderEmpty } from "./helpers/is-folder-empty";
 import { getOnline } from "./helpers/is-online";
@@ -12,6 +12,7 @@ import terminalLink from "terminal-link";
 import type { InstallTemplateArgs } from "./helpers";
 import { installTemplate } from "./helpers";
 import { templatesDir } from "./helpers/dir";
+import { toolsRequireConfig } from "./helpers/tools";
 
 export type InstallAppArgs = Omit<
   InstallTemplateArgs,
@@ -38,6 +39,7 @@ export async function createApp({
   externalPort,
   postInstallAction,
   dataSource,
+  tools,
 }: InstallAppArgs): Promise<void> {
   const root = path.resolve(appPath);
 
@@ -82,6 +84,7 @@ export async function createApp({
     externalPort,
     postInstallAction,
     dataSource,
+    tools,
   };
 
   if (frontend) {
@@ -114,6 +117,17 @@ export async function createApp({
     console.log();
   }
 
+  if (toolsRequireConfig(tools)) {
+    console.log(
+      yellow(
+        `You have selected tools that require configuration. Please configure them in the ${terminalLink(
+          "tools_config.json",
+          `file://${root}/tools_config.json`,
+        )} file.`,
+      ),
+    );
+  }
+  console.log("");
   console.log(`${green("Success!")} Created ${appName} at ${appPath}`);
 
   console.log(

e2e/utils.ts

@@ -117,6 +117,8 @@ export async function runCreateLlama(
     externalPort,
     "--post-install-action",
     postInstallAction,
+    "--tools",
+    "none",
   ].join(" ");
   console.log(`running command '${command}' in ${cwd}`);
   let appProcess = exec(command, {

helpers/python.ts

@@ -6,6 +6,7 @@ import terminalLink from "terminal-link";
 import { copy } from "./copy";
 import { templatesDir } from "./dir";
 import { isPoetryAvailable, tryPoetryInstall } from "./poetry";
+import { getToolConfig } from "./tools";
 import { InstallTemplateArgs, TemplateVectorDB } from "./types";
 
 interface Dependency {
@@ -128,6 +129,7 @@ export const installPythonTemplate = async ({
   engine,
   vectorDb,
   dataSource,
+  tools,
   postInstallAction,
 }: Pick<
   InstallTemplateArgs,
@@ -137,6 +139,7 @@ export const installPythonTemplate = async ({
   | "engine"
   | "vectorDb"
   | "dataSource"
+  | "tools"
   | "postInstallAction"
 >) => {
   console.log("\nInitializing Python project with template:", template, "\n");
@@ -162,20 +165,44 @@ export const installPythonTemplate = async ({
   });
 
   if (engine === "context") {
+    const enginePath = path.join(root, "app", "engine");
     const compPath = path.join(templatesDir, "components");
-    let vectorDbDirName = vectorDb ?? "none";
+
+    const vectorDbDirName = vectorDb ?? "none";
     const VectorDBPath = path.join(
       compPath,
       "vectordbs",
       "python",
       vectorDbDirName,
     );
-    const enginePath = path.join(root, "app", "engine");
-    await copy("**", path.join(root, "app", "engine"), {
+    await copy("**", enginePath, {
       parents: true,
       cwd: VectorDBPath,
     });
 
+    // Copy engine code
+    if (tools !== undefined && tools.length > 0) {
+      await copy("**", enginePath, {
+        parents: true,
+        cwd: path.join(compPath, "engines", "python", "agent"),
+      });
+      // Write tools_config.json
+      const configContent: Record<string, any> = {};
+      tools.forEach((tool) => {
+        configContent[tool] = getToolConfig(tool) ?? {};
+      });
+      const configFilePath = path.join(root, "tools_config.json");
+      await fs.writeFile(
+        configFilePath,
+        JSON.stringify(configContent, null, 2),
+      );
+    } else {
+      await copy("**", enginePath, {
+        parents: true,
+        cwd: path.join(compPath, "engines", "python", "chat"),
+      });
+    }
+
     const dataSourceType = dataSource?.type;
     if (dataSourceType !== undefined && dataSourceType !== "none") {
       let loaderPath =

helpers/tools.ts

+export type Tool = {
+  display: string;
+  name: string;
+  config?: Record<string, any>;
+};
+
+export const supportedTools: Tool[] = [
+  {
+    display: "Google Search (configuration required)",
+    name: "google_search",
+    config: {
+      engine: "Your search engine id",
+      key: "Your search api key",
+      num: 2,
+    },
+  },
+  {
+    display: "Wikipedia",
+    name: "wikipedia",
+  },
+];
+
+export const getToolConfig = (name: string) => {
+  return supportedTools.find((tool) => tool.name === name)?.config;
+};
+
+export const toolsRequireConfig = (tools?: string[]): boolean => {
+  if (tools) {
+    return tools.some((tool) => getToolConfig(tool));
+  }
+  return false;
+};

helpers/types.ts

@@ -41,4 +41,5 @@ export interface InstallTemplateArgs {
   vectorDb?: TemplateVectorDB;
   externalPort?: number;
   postInstallAction?: TemplatePostInstallAction;
+  tools?: string[];
 }

index.ts

@@ -11,6 +11,7 @@ import { createApp } from "./create-app";
 import { getPkgManager } from "./helpers/get-pkg-manager";
 import { isFolderEmpty } from "./helpers/is-folder-empty";
 import { runApp } from "./helpers/run-app";
+import { supportedTools } from "./helpers/tools";
 import { validateNpmName } from "./helpers/validate-pkg";
 import packageJson from "./package.json";
 import { QuestionArgs, askQuestions, onPromptState } from "./questions";
@@ -146,6 +147,13 @@ const program = new Commander.Command(packageJson.name)
     `
 
   Select which vector database you would like to use, such as 'none', 'pg' or 'mongo'. The default option is not to use a vector database and use the local filesystem instead ('none').
+`,
+  )
+  .option(
+    "--tools <tools>",
+    `
+
+  Specify the tools you want to use by providing a comma-separated list. For example, 'google_search,wikipedia'. Use 'none' to not using any tools.
 `,
   )
   .allowUnknownOption()
@@ -156,6 +164,25 @@ if (process.argv.includes("--no-frontend")) {
 if (process.argv.includes("--no-eslint")) {
   program.eslint = false;
 }
+if (process.argv.includes("--tools")) {
+  if (program.tools === "none") {
+    program.tools = [];
+  } else {
+    program.tools = program.tools.split(",");
+    // Check if tools are available
+    const toolsName = supportedTools.map((tool) => tool.name);
+    program.tools.forEach((tool: string) => {
+      if (!toolsName.includes(tool)) {
+        console.error(
+          `Error: Tool '${tool}' is not supported. Supported tools are: ${toolsName.join(
+            ", ",
+          )}`,
+        );
+        process.exit(1);
+      }
+    });
+  }
+}
 
 const packageManager = !!program.useNpm
   ? "npm"
@@ -256,6 +283,7 @@ async function run(): Promise<void> {
     externalPort: program.externalPort,
     postInstallAction: program.postInstallAction,
     dataSource: program.dataSource,
+    tools: program.tools,
   });
   conf.set("preferences", preferences);
 

questions.ts

@@ -10,6 +10,7 @@ import { COMMUNITY_OWNER, COMMUNITY_REPO } from "./helpers/constant";
 import { templatesDir } from "./helpers/dir";
 import { getAvailableLlamapackOptions } from "./helpers/llama-pack";
 import { getRepoRootFolders } from "./helpers/repo";
+import { supportedTools, toolsRequireConfig } from "./helpers/tools";
 
 export type QuestionArgs = Omit<
   InstallAppArgs,
@@ -70,6 +71,7 @@ const defaults: QuestionArgs = {
     type: "none",
     config: {},
   },
+  tools: [],
 };
 
 const handlers = {
@@ -214,7 +216,12 @@ export const askQuestions = async (
 
         const hasOpenAiKey = program.openAiKey || process.env["OPENAI_API_KEY"];
         const hasVectorDb = program.vectorDb && program.vectorDb !== "none";
-        if (!hasVectorDb && hasOpenAiKey) {
+        // Can run the app if all tools do not require configuration
+        if (
+          !hasVectorDb &&
+          hasOpenAiKey &&
+          !toolsRequireConfig(program.tools)
+        ) {
           actionChoices.push({
             title:
               "Generate code, install dependencies, and run the app (~2 min)",
@@ -563,6 +570,29 @@ export const askQuestions = async (
     }
   }
 
+  if (
+    !program.tools &&
+    program.framework === "fastapi" &&
+    program.engine === "context"
+  ) {
+    if (ciInfo.isCI) {
+      program.tools = getPrefOrDefault("tools");
+    } else {
+      const toolChoices = supportedTools.map((tool) => ({
+        title: tool.display,
+        value: tool.name,
+      }));
+      const { tools } = await prompts({
+        type: "multiselect",
+        name: "tools",
+        message: "Which tools would you like to use?",
+        choices: toolChoices,
+      });
+      program.tools = tools;
+      preferences.tools = tools;
+    }
+  }
+
   if (!program.openAiKey) {
     const { key } = await prompts(
       {

templates/components/engines/python/agent/__init__.py

+import os
+
+from typing import Any, Optional
+from llama_index.llms import LLM
+from llama_index.agent import AgentRunner
+
+from app.engine.tools import ToolFactory
+from app.engine.index import get_index
+from llama_index.agent import ReActAgent
+from llama_index.tools.query_engine import QueryEngineTool
+
+
+def create_agent_from_llm(
+    llm: Optional[LLM] = None,
+    **kwargs: Any,
+) -> AgentRunner:
+    from llama_index.agent import OpenAIAgent, ReActAgent
+    from llama_index.llms.openai import OpenAI
+    from llama_index.llms.openai_utils import is_function_calling_model
+
+    if isinstance(llm, OpenAI) and is_function_calling_model(llm.model):
+        return OpenAIAgent.from_tools(
+            llm=llm,
+            **kwargs,
+        )
+    else:
+        return ReActAgent.from_tools(
+            llm=llm,
+            **kwargs,
+        )
+
+
+def get_chat_engine():
+    tools = []
+
+    # Add query tool
+    index = get_index()
+    llm = index.service_context.llm
+    query_engine = index.as_query_engine(similarity_top_k=5)
+    query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
+    tools.append(query_engine_tool)
+
+    # Add additional tools
+    tools += ToolFactory.from_env()
+
+    return create_agent_from_llm(
+        llm=llm,
+        tools=tools,
+        verbose=True,
+    )

templates/components/engines/python/agent/tools.py

+import json
+import importlib
+
+from llama_index.tools.tool_spec.base import BaseToolSpec
+from llama_index.tools.function_tool import FunctionTool
+
+
+class ToolFactory:
+
+    @staticmethod
+    def create_tool(tool_name: str, **kwargs) -> list[FunctionTool]:
+        try:
+            module_name = f"llama_hub.tools.{tool_name}.base"
+            module = importlib.import_module(module_name)
+            tool_cls_name = tool_name.title().replace("_", "") + "ToolSpec"
+            tool_class = getattr(module, tool_cls_name)
+            tool_spec: BaseToolSpec = tool_class(**kwargs)
+            return tool_spec.to_tool_list()
+        except (ImportError, AttributeError) as e:
+            raise ValueError(f"Unsupported tool: {tool_name}") from e
+        except TypeError as e:
+            raise ValueError(
+                f"Could not create tool: {tool_name}. With config: {kwargs}"
+            ) from e
+
+    @staticmethod
+    def from_env() -> list[FunctionTool]:
+        tools = []
+        with open("tools_config.json", "r") as f:
+            tool_configs = json.load(f)
+            for name, config in tool_configs.items():
+                tools += ToolFactory.create_tool(name, **config)
+        return tools

templates/components/engines/python/chat/__init__.py

+from app.engine.index import get_index
+
+
+def get_chat_engine():
+    return get_index().as_chat_engine(
+        similarity_top_k=5, chat_mode="condense_plus_context"
+    )

templates/components/vectordbs/python/mongo/index.py

@@ -9,7 +9,7 @@ from llama_index.vector_stores import MongoDBAtlasVectorSearch
 from app.engine.context import create_service_context
 
 
-def get_chat_engine():
+def get_index():
     service_context = create_service_context()
     logger = logging.getLogger("uvicorn")
     logger.info("Connecting to index from MongoDB...")
@@ -20,4 +20,4 @@ def get_chat_engine():
     )
     index = VectorStoreIndex.from_vector_store(store, service_context)
     logger.info("Finished connecting to index from MongoDB.")
-    return index.as_chat_engine(similarity_top_k=5, chat_mode="condense_plus_context")
+    return index

templates/components/vectordbs/python/none/index.py

 import logging
 import os
+
+from app.engine.constants import STORAGE_DIR
+from app.engine.context import create_service_context
 from llama_index import (
     StorageContext,
     load_index_from_storage,
 )
 
-from app.engine.constants import STORAGE_DIR
-from app.engine.context import create_service_context
-
 
-def get_chat_engine():
+def get_index():
     service_context = create_service_context()
     # check if storage already exists
     if not os.path.exists(STORAGE_DIR):
@@ -22,4 +22,4 @@ def get_chat_engine():
     storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
     index = load_index_from_storage(storage_context, service_context=service_context)
     logger.info(f"Finished loading index from {STORAGE_DIR}")
-    return index.as_chat_engine(similarity_top_k=5, chat_mode="condense_plus_context")
+    return index

templates/components/vectordbs/python/pg/index.py

@@ -6,11 +6,11 @@ from app.engine.context import create_service_context
 from app.engine.utils import init_pg_vector_store_from_env
 
 
-def get_chat_engine():
+def get_index():
     service_context = create_service_context()
     logger = logging.getLogger("uvicorn")
     logger.info("Connecting to index from PGVector...")
     store = init_pg_vector_store_from_env()
     index = VectorStoreIndex.from_vector_store(store, service_context)
     logger.info("Finished connecting to index from PGVector.")
-    return index.as_chat_engine(similarity_top_k=5, chat_mode="condense_plus_context")
+    return index

templates/types/simple/fastapi/app/api/routers/chat.py

@@ -5,7 +5,7 @@ from llama_index.chat_engine.types import BaseChatEngine
 from llama_index.llms.base import ChatMessage
 from llama_index.llms.types import MessageRole
 from pydantic import BaseModel
-from app.engine.index import get_chat_engine
+from app.engine import get_chat_engine
 
 chat_router = r = APIRouter()
 

templates/types/simple/fastapi/app/engine/__init__.py

+from llama_index.chat_engine import SimpleChatEngine
+
+from app.context import create_base_context
+
+
+def get_chat_engine():
+    return SimpleChatEngine.from_defaults(service_context=create_base_context())

templates/types/simple/fastapi/app/engine/index.py

-from llama_index.chat_engine import SimpleChatEngine
-
-from app.context import create_base_context
-
-
-def get_chat_engine():
-    return SimpleChatEngine.from_defaults(service_context=create_base_context())

templates/types/simple/fastapi/pyproject.toml

@@ -13,6 +13,8 @@ llama-index = "^0.9.19"
 pypdf = "^3.17.0"
 python-dotenv = "^1.0.0"
 docx2txt = "^0.8"
+llama-hub = "^0.0.77"
+wikipedia = "^1.4.0"
 
 [build-system]
 requires = ["poetry-core"]

templates/types/streaming/fastapi/app/api/routers/chat.py

@@ -3,7 +3,7 @@ from typing import List
 from fastapi.responses import StreamingResponse
 from llama_index.chat_engine.types import BaseChatEngine
 
-from app.engine.index import get_chat_engine
+from app.engine import get_chat_engine
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from llama_index.llms.base import ChatMessage
 from llama_index.llms.types import MessageRole

templates/types/streaming/fastapi/app/engine/__init__.py

+from llama_index.chat_engine import SimpleChatEngine
+
+from app.context import create_base_context
+
+
+def get_chat_engine():
+    return SimpleChatEngine.from_defaults(service_context=create_base_context())

templates/types/streaming/fastapi/app/engine/index.py

-from llama_index.chat_engine import SimpleChatEngine
-
-from app.context import create_base_context
-
-
-def get_chat_engine():
-    return SimpleChatEngine.from_defaults(service_context=create_base_context())