feat: add test and fix python dependencies

e2e/resolve_python_dependencies.spec.ts

@@ -0,0 +1,110 @@
+import { expect, test } from "@playwright/test";
+import { exec } from "child_process";
+import fs from "fs";
+import path from "path";
+import util from "util";
+import { TemplateVectorDB } from "../helpers/types";
+import { createTestDir, runCreateLlama } from "./utils";
+
+const execAsync = util.promisify(exec);
+
+const vectorDbs: TemplateVectorDB[] = [
+  "mongo",
+  "pg",
+  "pinecone",
+  "milvus",
+  "astra",
+  "qdrant",
+  "chroma",
+  "weaviate",
+];
+
+const toolOptions = [
+  "wikipedia.WikipediaToolSpec",
+  "google.GoogleSearchToolSpec",
+];
+
+// TODO: Add data sources to the test
+
+test.describe("Test resolve python dependencies", () => {
+  for (const vectorDb of vectorDbs) {
+    for (const tool of toolOptions) {
+      const optionDescription = `vectorDb: ${vectorDb}, tools: ${tool}`;
+
+      test(`options: ${optionDescription}`, async () => {
+        const cwd = await createTestDir();
+
+        const result = await runCreateLlama(
+          cwd,
+          "streaming",
+          "fastapi",
+          "--example-file",
+          vectorDb,
+          3000, // port
+          8000, // externalPort
+          "none", // postInstallAction
+          undefined, // ui
+          "--no-frontend", // appType
+          undefined, // llamaCloudProjectName
+          undefined, // llamaCloudIndexName
+          tool,
+        );
+        const name = result.projectName;
+
+        // Check if the app folder exists
+        const dirExists = fs.existsSync(path.join(cwd, name));
+        expect(dirExists).toBeTruthy();
+
+        // Check if pyproject.toml exists
+        const pyprojectPath = path.join(cwd, name, "pyproject.toml");
+        const pyprojectExists = fs.existsSync(pyprojectPath);
+        expect(pyprojectExists).toBeTruthy();
+
+        // Run poetry lock
+        try {
+          const { stdout, stderr } = await execAsync(
+            // Config poetry to create virtualenv in project directory.
+            // so that we can easily prune the e2e cache to avoid overloading the storage.
+            "poetry config virtualenvs.in-project true && poetry lock --no-update",
+            {
+              cwd: path.join(cwd, name),
+            },
+          );
+          console.log("poetry lock stdout:", stdout);
+          console.error("poetry lock stderr:", stderr);
+        } catch (error) {
+          console.error("Error running poetry lock:", error);
+          throw error;
+        }
+
+        // Check if poetry.lock file was created
+        const poetryLockExists = fs.existsSync(
+          path.join(cwd, name, "poetry.lock"),
+        );
+        expect(poetryLockExists).toBeTruthy();
+
+        // Verify that specific dependencies are in pyproject.toml
+        const pyprojectContent = fs.readFileSync(pyprojectPath, "utf-8");
+        if (vectorDb !== "none") {
+          if (vectorDb === "pg") {
+            expect(pyprojectContent).toContain(
+              "llama-index-vector-stores-postgres",
+            );
+          } else {
+            expect(pyprojectContent).toContain(
+              `llama-index-vector-stores-${vectorDb}`,
+            );
+          }
+        }
+        if (tool !== "none") {
+          if (tool === "wikipedia.WikipediaToolSpec") {
+            expect(pyprojectContent).toContain("wikipedia");
+          }
+          if (tool === "google.GoogleSearchToolSpec") {
+            expect(pyprojectContent).toContain("google");
+          }
+        }
+      });
+    }
+  }
+});

e2e/utils.ts

@@ -32,6 +32,7 @@ export async function runCreateLlama(
   appType?: AppType,
   llamaCloudProjectName?: string,
   llamaCloudIndexName?: string,
+  tools?: string,
 ): Promise<CreateLlamaResult> {
   if (!process.env.OPENAI_API_KEY || !process.env.LLAMA_CLOUD_API_KEY) {
     throw new Error(
@@ -65,7 +66,7 @@ export async function runCreateLlama(
     "--post-install-action",
     postInstallAction,
     "--tools",
-    "none",
+    tools ?? "none",
     "--no-llama-parse",
     "--observability",
     "none",

helpers/python.ts

@@ -36,14 +36,14 @@ const getAdditionalDependencies = (
     case "mongo": {
       dependencies.push({
         name: "llama-index-vector-stores-mongodb",
-        version: "^0.1.3",
+        version: "^0.3.1",
       });
       break;
     }
     case "pg": {
       dependencies.push({
         name: "llama-index-vector-stores-postgres",
-        version: "^0.1.1",
+        version: "^0.2.5",
       });
       break;
     }
@@ -57,7 +57,7 @@ const getAdditionalDependencies = (
     case "milvus": {
       dependencies.push({
         name: "llama-index-vector-stores-milvus",
-        version: "^0.1.20",
+        version: "^0.2.0",
       });
       dependencies.push({
         name: "pymilvus",
@@ -68,7 +68,7 @@ const getAdditionalDependencies = (
     case "astra": {
       dependencies.push({
         name: "llama-index-vector-stores-astra-db",
-        version: "^0.1.5",
+        version: "^0.2.0",
       });
       break;
     }
@@ -82,14 +82,14 @@ const getAdditionalDependencies = (
     case "chroma": {
       dependencies.push({
         name: "llama-index-vector-stores-chroma",
-        version: "^0.1.8",
+        version: "^0.2.0",
       });
       break;
     }
     case "weaviate": {
       dependencies.push({
         name: "llama-index-vector-stores-weaviate",
-        version: "^1.0.2",
+        version: "^1.1.1",
       });
       break;
     }
@@ -130,7 +130,7 @@ const getAdditionalDependencies = (
         case "llamacloud":
           dependencies.push({
             name: "llama-index-indices-managed-llama-cloud",
-            version: "^0.3.0",
+            version: "^0.3.1",
           });
           break;
       }