Skip to content

Commit

Permalink
feat: add test and fix python dependencies (#304)
Browse files Browse the repository at this point in the history
---------

Co-authored-by: Marcus Schiesser <[email protected]>
  • Loading branch information
leehuwuj and marcusschiesser authored Sep 23, 2024
1 parent 0031e67 commit 40c5c84
Show file tree
Hide file tree
Showing 6 changed files with 202 additions and 14 deletions.
139 changes: 139 additions & 0 deletions e2e/resolve_python_dependencies.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
import { expect, test } from "@playwright/test";
import { exec } from "child_process";
import fs from "fs";
import path from "path";
import util from "util";
import { TemplateFramework, TemplateVectorDB } from "../helpers/types";
import { createTestDir, runCreateLlama } from "./utils";

const execAsync = util.promisify(exec);

const templateFramework: TemplateFramework = process.env.FRAMEWORK
? (process.env.FRAMEWORK as TemplateFramework)
: "fastapi";
const dataSource: string = process.env.DATASOURCE
? process.env.DATASOURCE
: "--example-file";

if (
templateFramework == "fastapi" && // test is only relevant for fastapi
process.version.startsWith("v20.") && // XXX: Only run for Node.js version 20 (CI matrix will trigger other versions)
dataSource === "--example-file" // XXX: this test provides its own data source - only trigger it on one data source (usually the CI matrix will trigger multiple data sources)
) {
// vectorDBs, tools, and data source combinations to test
const vectorDbs: TemplateVectorDB[] = [
"mongo",
"pg",
"pinecone",
"milvus",
"astra",
"qdrant",
"chroma",
"weaviate",
];

const toolOptions = [
"wikipedia.WikipediaToolSpec",
"google.GoogleSearchToolSpec",
];

const dataSources = [
"--example-file",
"--web-source https://www.example.com",
"--db-source mysql+pymysql://user:pass@localhost:3306/mydb",
];

test.describe("Test resolve python dependencies", () => {
for (const vectorDb of vectorDbs) {
for (const tool of toolOptions) {
for (const dataSource of dataSources) {
const dataSourceType = dataSource.split(" ")[0];
const optionDescription = `vectorDb: ${vectorDb}, tools: ${tool}, dataSource: ${dataSourceType}`;

test(`options: ${optionDescription}`, async () => {
const cwd = await createTestDir();

const result = await runCreateLlama(
cwd,
"streaming",
"fastapi",
dataSource,
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(
"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");
}
}

// Check for data source specific dependencies
if (dataSource.includes("--web-source")) {
expect(pyprojectContent).toContain("llama-index-readers-web");
}
if (dataSource.includes("--db-source")) {
expect(pyprojectContent).toContain(
"llama-index-readers-database ",
);
}
});
}
}
}
});
}
20 changes: 17 additions & 3 deletions e2e/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -41,18 +42,31 @@ export async function runCreateLlama(
const name = [
templateType,
templateFramework,
dataSource,
dataSource.split(" ")[0],
templateUI,
appType,
].join("-");

// Handle different data source types
let dataSourceArgs = [];
if (dataSource.includes("--web-source" || "--db-source")) {
const webSource = dataSource.split(" ")[1];
dataSourceArgs.push("--web-source", webSource);
} else if (dataSource.includes("--db-source")) {
const dbSource = dataSource.split(" ")[1];
dataSourceArgs.push("--db-source", dbSource);
} else {
dataSourceArgs.push(dataSource);
}

const commandArgs = [
"create-llama",
name,
"--template",
templateType,
"--framework",
templateFramework,
dataSource,
...dataSourceArgs,
"--vector-db",
vectorDb,
"--open-ai-key",
Expand All @@ -65,7 +79,7 @@ export async function runCreateLlama(
"--post-install-action",
postInstallAction,
"--tools",
"none",
tools ?? "none",
"--no-llama-parse",
"--observability",
"none",
Expand Down
18 changes: 9 additions & 9 deletions helpers/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,28 +36,28 @@ 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;
}
case "pinecone": {
dependencies.push({
name: "llama-index-vector-stores-pinecone",
version: "^0.1.3",
version: "^0.2.1",
});
break;
}
case "milvus": {
dependencies.push({
name: "llama-index-vector-stores-milvus",
version: "^0.1.20",
version: "^0.2.0",
});
dependencies.push({
name: "pymilvus",
Expand All @@ -68,28 +68,28 @@ const getAdditionalDependencies = (
case "astra": {
dependencies.push({
name: "llama-index-vector-stores-astra-db",
version: "^0.1.5",
version: "^0.2.0",
});
break;
}
case "qdrant": {
dependencies.push({
name: "llama-index-vector-stores-qdrant",
version: "^0.2.8",
version: "^0.3.0",
});
break;
}
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;
}
Expand Down Expand Up @@ -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;
}
Expand Down
35 changes: 35 additions & 0 deletions index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,20 @@ const program = new Commander.Command(packageJson.name)
`
Select to use an example PDF as data source.
`,
)
.option(
"--web-source <url>",
`
Specify a website URL to use as a data source.
`,
)
.option(
"--db-source <connection-string>",
`
Specify a database connection string to use as a data source.
`,
)
.option(
Expand Down Expand Up @@ -215,6 +229,27 @@ if (process.argv.includes("--no-files")) {
},
EXAMPLE_FILE,
];
} else if (process.argv.includes("--web-source")) {
program.dataSources = [
{
type: "web",
config: {
baseUrl: program.webSource,
prefix: program.webSource,
depth: 1,
},
},
];
} else if (process.argv.includes("--db-source")) {
program.dataSources = [
{
type: "db",
config: {
uri: program.dbSource,
queries: program.dbQuery || "SELECT * FROM mytable",
},
},
];
}

const packageManager = !!program.useNpm
Expand Down
2 changes: 1 addition & 1 deletion templates/types/multiagent/fastapi/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ readme = "README.md"
generate = "app.engine.generate:generate_datasource"

[tool.poetry.dependencies]
python = "^3.11"
python = ">=3.11,<3.13"
llama-index-agent-openai = ">=0.3.0,<0.4.0"
llama-index = "0.11.11"
fastapi = "^0.112.2"
Expand Down
2 changes: 1 addition & 1 deletion templates/types/streaming/fastapi/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ readme = "README.md"
generate = "app.engine.generate:generate_datasource"

[tool.poetry.dependencies]
python = "^3.11,<4.0"
python = ">=3.11,<3.13"
fastapi = "^0.109.1"
uvicorn = { extras = ["standard"], version = "^0.23.2" }
python-dotenv = "^1.0.0"
Expand Down

0 comments on commit 40c5c84

Please sign in to comment.