feat: use LlamaCloud for TS/Python (#149)

thucpn · marcusschiesser · web-flow · commit aeb6fef4daa5 · 2024-06-28T22:10:37.000+07:00
---------
Co-authored-by: Marcus Schiesser &lt;mail@marcusschiesser.de&gt;
diff --git a/.changeset/tough-pugs-destroy.md b/.changeset/tough-pugs-destroy.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+use llamacloud for chat
diff --git a/helpers/env-variables.ts b/helpers/env-variables.ts
@@ -133,6 +133,26 @@ const getVectorDBEnvs = (
             "Optional API key for authenticating requests to Qdrant.",
         },
       ];
+    case "llamacloud":
+      return [
+        {
+          name: "LLAMA_CLOUD_INDEX_NAME",
+          description:
+            "The name of the LlamaCloud index to use (part of the LlamaCloud project).",
+          value: "test",
+        },
+        {
+          name: "LLAMA_CLOUD_PROJECT_NAME",
+          description: "The name of the LlamaCloud project.",
+          value: "Default",
+        },
+        {
+          name: "LLAMA_CLOUD_BASE_URL",
+          description:
+            "The base URL for the LlamaCloud API. Only change this for non-production environments",
+          value: "https://api.cloud.llamaindex.ai",
+        },
+      ];
     case "chroma":
       const envs = [
         {
diff --git a/helpers/python.ts b/helpers/python.ts
@@ -118,6 +118,12 @@ const getAdditionalDependencies = (
             version: "^2.9.9",
           });
           break;
+        case "llamacloud":
+          dependencies.push({
+            name: "llama-index-indices-managed-llama-cloud",
+            version: "^0.2.1",
+          });
+          break;
       }
     }
   }
diff --git a/helpers/types.ts b/helpers/types.ts
@@ -26,7 +26,8 @@ export type TemplateVectorDB =
   | "milvus"
   | "astra"
   | "qdrant"
-  | "chroma";
+  | "chroma"
+  | "llamacloud";
 export type TemplatePostInstallAction =
   | "none"
   | "VSCode"
@@ -36,7 +37,7 @@ export type TemplateDataSource = {
   type: TemplateDataSourceType;
   config: TemplateDataSourceConfig;
 };
-export type TemplateDataSourceType = "file" | "web" | "db";
+export type TemplateDataSourceType = "file" | "web" | "db" | "llamacloud";
 export type TemplateObservability = "none" | "opentelemetry";
 // Config for both file and folder
 export type FileSourceConfig = {
diff --git a/helpers/typescript.ts b/helpers/typescript.ts
@@ -1,7 +1,7 @@
 import fs from "fs/promises";
 import os from "os";
 import path from "path";
-import { bold, cyan } from "picocolors";
+import { bold, cyan, yellow } from "picocolors";
 import { assetRelocator, copy } from "../helpers/copy";
 import { callPackageManager } from "../helpers/install";
 import { templatesDir } from "./dir";
@@ -105,7 +105,13 @@ export const installTSTemplate = async ({
   const enginePath = path.join(root, relativeEngineDestPath, "engine");
 
   // copy vector db component
-  console.log("\nUsing vector DB:", vectorDb ?? "none", "\n");
+  if (vectorDb === "llamacloud") {
+    console.log(
+      `\nUsing managed index from LlamaCloud. Ensure the ${yellow("LLAMA_CLOUD_* environment variables are set correctly.")}`,
+    );
+  } else {
+    console.log("\nUsing vector DB:", vectorDb ?? "none");
+  }
   await copy("**", enginePath, {
     parents: true,
     cwd: path.join(compPath, "vectordbs", "typescript", vectorDb ?? "none"),
diff --git a/questions.ts b/questions.ts
@@ -123,7 +123,13 @@ export const getDataSourceChoices = (
   framework: TemplateFramework,
   selectedDataSource: TemplateDataSource[],
 ) => {
+  // If LlamaCloud is already selected, don't show any other options
+  if (selectedDataSource.find((s) => s.type === "llamacloud")) {
+    return [];
+  }
+
   const choices = [];
+
   if (selectedDataSource.length > 0) {
     choices.push({
       title: "No",
@@ -171,6 +177,13 @@ export const getDataSourceChoices = (
       value: "db",
     });
   }
+
+  if (!selectedDataSource.length) {
+    choices.push({
+      title: "Use managed index from LlamaCloud",
+      value: "llamacloud",
+    });
+  }
   return choices;
 };
 
@@ -484,17 +497,19 @@ export const askQuestions = async (
       // continue asking user for data sources if none are initially provided
       while (true) {
         const firstQuestion = program.dataSources.length === 0;
+        const choices = getDataSourceChoices(
+          program.framework,
+          program.dataSources,
+        );
+        if (choices.length === 0) break;
         const { selectedSource } = await prompts(
           {
             type: "select",
             name: "selectedSource",
             message: firstQuestion
               ? "Which data source would you like to use?"
               : "Would you like to add another data source?",
-            choices: getDataSourceChoices(
-              program.framework,
-              program.dataSources,
-            ),
+            choices,
             initial: firstQuestion ? 1 : 0,
           },
           questionHandlers,
@@ -591,51 +606,76 @@ export const askQuestions = async (
               config: await prompts(dbPrompts, questionHandlers),
             });
           }
+          case "llamacloud": {
+            program.dataSources.push({
+              type: "llamacloud",
+              config: {},
+            });
+            program.dataSources.push(EXAMPLE_FILE);
+            break;
+          }
         }
       }
     }
   }
 
-  // Asking for LlamaParse if user selected file or folder data source
-  if (
-    program.dataSources.some((ds) => ds.type === "file") &&
-    program.useLlamaParse === undefined
-  ) {
+  const isUsingLlamaCloud = program.dataSources.some(
+    (ds) => ds.type === "llamacloud",
+  );
+
+  // Asking for LlamaParse if user selected file data source
+  if (isUsingLlamaCloud) {
+    // default to use LlamaParse if using LlamaCloud
+    program.useLlamaParse = preferences.useLlamaParse = true;
+  } else {
+    if (program.dataSources.some((ds) => ds.type === "file")) {
+      if (ciInfo.isCI) {
+        program.useLlamaParse = getPrefOrDefault("useLlamaParse");
+      } else {
+        const { useLlamaParse } = await prompts(
+          {
+            type: "toggle",
+            name: "useLlamaParse",
+            message:
+              "Would you like to use LlamaParse (improved parser for RAG - requires API key)?",
+            initial: false,
+            active: "yes",
+            inactive: "no",
+          },
+          questionHandlers,
+        );
+        program.useLlamaParse = useLlamaParse;
+        preferences.useLlamaParse = useLlamaParse;
+      }
+    }
+  }
+
+  // Ask for LlamaCloud API key when using a LlamaCloud index or LlamaParse
+  if (isUsingLlamaCloud || program.useLlamaParse) {
     if (ciInfo.isCI) {
-      program.useLlamaParse = getPrefOrDefault("useLlamaParse");
       program.llamaCloudKey = getPrefOrDefault("llamaCloudKey");
     } else {
-      const { useLlamaParse } = await prompts(
+      // Ask for LlamaCloud API key
+      const { llamaCloudKey } = await prompts(
         {
-          type: "toggle",
-          name: "useLlamaParse",
+          type: "text",
+          name: "llamaCloudKey",
           message:
-            "Would you like to use LlamaParse (improved parser for RAG - requires API key)?",
-          initial: false,
-          active: "yes",
-          inactive: "no",
+            "Please provide your LlamaCloud API key (leave blank to skip):",
         },
         questionHandlers,
       );
-      program.useLlamaParse = useLlamaParse;
-
-      // Ask for LlamaCloud API key
-      if (useLlamaParse && program.llamaCloudKey === undefined) {
-        const { llamaCloudKey } = await prompts(
-          {
-            type: "text",
-            name: "llamaCloudKey",
-            message:
-              "Please provide your LlamaIndex Cloud API key (leave blank to skip):",
-          },
-          questionHandlers,
-        );
-        program.llamaCloudKey = llamaCloudKey;
-      }
+      program.llamaCloudKey = preferences.llamaCloudKey =
+        llamaCloudKey || process.env.LLAMA_CLOUD_API_KEY;
     }
   }
 
-  if (program.dataSources.length > 0 && !program.vectorDb) {
+  if (isUsingLlamaCloud) {
+    // When using a LlamaCloud index, don't ask for vector database and use code in `llamacloud` folder for vector database
+    const vectorDb = "llamacloud";
+    program.vectorDb = vectorDb;
+    preferences.vectorDb = vectorDb;
+  } else if (program.dataSources.length > 0 && !program.vectorDb) {
     if (ciInfo.isCI) {
       program.vectorDb = getPrefOrDefault("vectorDb");
     } else {
diff --git a/templates/components/engines/typescript/agent/chat.ts b/templates/components/engines/typescript/agent/chat.ts
@@ -2,7 +2,6 @@ import { BaseToolWithCall, OpenAIAgent, QueryEngineTool } from "llamaindex";
 import fs from "node:fs/promises";
 import path from "node:path";
 import { getDataSource } from "./index";
-import { STORAGE_CACHE_DIR } from "./shared";
 import { createTools } from "./tools";
 
 export async function createChatEngine() {
@@ -17,7 +16,7 @@ export async function createChatEngine() {
         queryEngine: index.asQueryEngine(),
         metadata: {
           name: "data_query_engine",
-          description: `A query engine for documents in storage folder: ${STORAGE_CACHE_DIR}`,
+          description: `A query engine for documents from your data source.`,
         },
       }),
     );
diff --git a/templates/components/engines/typescript/chat/chat.ts b/templates/components/engines/typescript/chat/chat.ts
@@ -8,14 +8,14 @@ export async function createChatEngine() {
       `StorageContext is empty - call 'npm run generate' to generate the storage first`,
     );
   }
-  const retriever = index.asRetriever();
-  retriever.similarityTopK = process.env.TOP_K
-    ? parseInt(process.env.TOP_K)
-    : 3;
+  const retriever = index.asRetriever({
+    similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : 3,
+  });
 
   return new ContextChatEngine({
     chatModel: Settings.llm,
     retriever,
-    systemPrompt: process.env.SYSTEM_PROMPT,
+    // disable as a custom system prompt disables the generated context
+    // systemPrompt: process.env.SYSTEM_PROMPT,
   });
 }
diff --git a/templates/components/vectordbs/python/llamacloud/__init__.py b/templates/components/vectordbs/python/llamacloud/__init__.py
diff --git a/templates/components/vectordbs/python/llamacloud/generate.py b/templates/components/vectordbs/python/llamacloud/generate.py
@@ -0,0 +1,45 @@
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import os
+import logging
+from app.settings import init_settings
+from app.engine.loaders import get_documents
+from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger()
+
+
+def generate_datasource():
+    init_settings()
+    logger.info("Generate index for the provided data")
+
+    name = os.getenv("LLAMA_CLOUD_INDEX_NAME")
+    project_name = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
+    api_key = os.getenv("LLAMA_CLOUD_API_KEY")
+    base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
+
+    if name is None or project_name is None or api_key is None:
+        raise ValueError(
+            "Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
+            " to your environment variables or config them in .env file"
+        )
+
+    documents = get_documents()
+
+    LlamaCloudIndex.from_documents(
+        documents=documents,
+        name=name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    logger.info("Finished generating the index")
+
+
+if __name__ == "__main__":
+    generate_datasource()
diff --git a/templates/components/vectordbs/python/llamacloud/index.py b/templates/components/vectordbs/python/llamacloud/index.py
@@ -0,0 +1,28 @@
+import logging
+import os
+from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
+
+
+logger = logging.getLogger("uvicorn")
+
+
+def get_index():
+    name = os.getenv("LLAMA_CLOUD_INDEX_NAME")
+    project_name = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
+    api_key = os.getenv("LLAMA_CLOUD_API_KEY")
+    base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
+
+    if name is None or project_name is None or api_key is None:
+        raise ValueError(
+            "Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
+            " to your environment variables or config them in .env file"
+        )
+
+    index = LlamaCloudIndex(
+        name=name,
+        project_name=project_name,
+        api_key=api_key,
+        base_url=base_url,
+    )
+
+    return index
diff --git a/templates/components/vectordbs/typescript/llamacloud/generate.ts b/templates/components/vectordbs/typescript/llamacloud/generate.ts
@@ -0,0 +1,26 @@
+import * as dotenv from "dotenv";
+import { LlamaCloudIndex } from "llamaindex";
+import { getDocuments } from "./loader";
+import { initSettings } from "./settings";
+import { checkRequiredEnvVars } from "./shared";
+
+dotenv.config();
+
+async function loadAndIndex() {
+  const documents = await getDocuments();
+  await LlamaCloudIndex.fromDocuments({
+    documents,
+    name: process.env.LLAMA_CLOUD_INDEX_NAME!,
+    projectName: process.env.LLAMA_CLOUD_PROJECT_NAME!,
+    apiKey: process.env.LLAMA_CLOUD_API_KEY,
+    baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
+  });
+  console.log(`Successfully created embeddings!`);
+}
+
+(async () => {
+  checkRequiredEnvVars();
+  initSettings();
+  await loadAndIndex();
+  console.log("Finished generating storage.");
+})();
diff --git a/templates/components/vectordbs/typescript/llamacloud/index.ts b/templates/components/vectordbs/typescript/llamacloud/index.ts
@@ -0,0 +1,13 @@
+import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
+import { checkRequiredEnvVars } from "./shared";
+
+export async function getDataSource() {
+  checkRequiredEnvVars();
+  const index = new LlamaCloudIndex({
+    name: process.env.LLAMA_CLOUD_INDEX_NAME!,
+    projectName: process.env.LLAMA_CLOUD_PROJECT_NAME!,
+    apiKey: process.env.LLAMA_CLOUD_API_KEY,
+    baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
+  });
+  return index;
+}
diff --git a/templates/components/vectordbs/typescript/llamacloud/shared.ts b/templates/components/vectordbs/typescript/llamacloud/shared.ts
diff --git a/templates/types/streaming/express/package.json b/templates/types/streaming/express/package.json
diff --git a/templates/types/streaming/fastapi/pyproject.toml b/templates/types/streaming/fastapi/pyproject.toml
diff --git a/templates/types/streaming/nextjs/package.json b/templates/types/streaming/nextjs/package.json

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"create-llama": patch
 +---
++
 +use llamacloud for chat
Original file line number	Diff line number	Diff line change
`@@ -118,6 +118,12 @@ const getAdditionalDependencies = (`
`118`	`118`	`version: "^2.9.9",`
`119`	`119`	`});`
`120`	`120`	`break;`
	`121`	`+ case "llamacloud":`
	`122`	`+ dependencies.push({`
	`123`	`+ name: "llama-index-indices-managed-llama-cloud",`
	`124`	`+ version: "^0.2.1",`
	`125`	`+ });`
	`126`	`+ break;`
`121`	`127`	`}`
`122`	`128`	`}`
`123`	`129`	`}`