Skip to content

Commit

Permalink
feat: Add legal document review use case (#467)
Browse files Browse the repository at this point in the history
  • Loading branch information
leehuwuj authored Dec 24, 2024
1 parent 765d2c4 commit 9077cae
Show file tree
Hide file tree
Showing 65 changed files with 1,444 additions and 110 deletions.
5 changes: 5 additions & 0 deletions .changeset/stale-scissors-turn.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"create-llama": patch
---

Add contract review use case (Python)
60 changes: 0 additions & 60 deletions e2e/shared/extractor_template.spec.ts

This file was deleted.

64 changes: 64 additions & 0 deletions e2e/shared/reflex_template.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
/* eslint-disable turbo/no-undeclared-env-vars */
import { expect, test } from "@playwright/test";
import { ChildProcess } from "child_process";
import fs from "fs";
import path from "path";
import { TemplateAgents, TemplateFramework } from "../../helpers";
import { createTestDir, runCreateLlama } from "../utils";

const templateFramework: TemplateFramework = process.env.FRAMEWORK
? (process.env.FRAMEWORK as TemplateFramework)
: "fastapi";
const dataSource: string = process.env.DATASOURCE
? process.env.DATASOURCE
: "--example-file";
const templateAgents: TemplateAgents[] = ["extractor", "contract_review"];

// The reflex template currently only works with FastAPI and files (and not on Windows)
if (
process.platform !== "win32" &&
templateFramework === "fastapi" &&
dataSource === "--example-file"
) {
for (const agents of templateAgents) {
test.describe(`Test reflex template ${agents} ${templateFramework} ${dataSource}`, async () => {
let appPort: number;
let name: string;
let appProcess: ChildProcess;
let cwd: string;

// Create reflex app
test.beforeAll(async () => {
cwd = await createTestDir();
appPort = Math.floor(Math.random() * 10000) + 10000;
const result = await runCreateLlama({
cwd,
templateType: "reflex",
templateFramework: "fastapi",
dataSource: "--example-file",
vectorDb: "none",
port: appPort,
postInstallAction: "runApp",
agents,
});
name = result.projectName;
appProcess = result.appProcess;
});

test.afterAll(async () => {
appProcess.kill();
});

test("App folder should exist", async () => {
const dirExists = fs.existsSync(path.join(cwd, name));
expect(dirExists).toBeTruthy();
});
test("Frontend should have a title", async ({ page }) => {
await page.goto(`http://localhost:${appPort}`);
await expect(page.getByText("Built by LlamaIndex")).toBeVisible({
timeout: 2000 * 60,
});
});
});
}
}
2 changes: 1 addition & 1 deletion e2e/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ export async function runCreateLlama({
if (observability) {
commandArgs.push("--observability", observability);
}
if (templateType === "multiagent" && agents) {
if ((templateType === "multiagent" || templateType === "reflex") && agents) {
commandArgs.push("--agents", agents);
}

Expand Down
12 changes: 12 additions & 0 deletions helpers/datasources.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ export const EXAMPLE_10K_SEC_FILES: TemplateDataSource[] = [
url: new URL(
"https://s2.q4cdn.com/470004039/files/doc_earnings/2023/q4/filing/_10-K-Q4-2023-As-Filed.pdf",
),
filename: "apple_10k_report.pdf",
},
},
{
Expand All @@ -26,10 +27,21 @@ export const EXAMPLE_10K_SEC_FILES: TemplateDataSource[] = [
url: new URL(
"https://ir.tesla.com/_flysystem/s3/sec/000162828024002390/tsla-20231231-gen.pdf",
),
filename: "tesla_10k_report.pdf",
},
},
];

export const EXAMPLE_GDPR: TemplateDataSource = {
type: "file",
config: {
url: new URL(
"https://eur-lex.europa.eu/legal-content/EN/TXT/PDF/?uri=CELEX:32016R0679",
),
filename: "gdpr.pdf",
},
};

export function getDataSources(
files?: string,
exampleFile?: boolean,
Expand Down
5 changes: 3 additions & 2 deletions helpers/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,8 @@ const prepareContextData = async (
const destPath = path.join(
root,
"data",
path.basename(dataSourceConfig.url.toString()),
dataSourceConfig.filename ??
path.basename(dataSourceConfig.url.toString()),
);
await downloadFile(dataSourceConfig.url.toString(), destPath);
} else {
Expand Down Expand Up @@ -192,7 +193,7 @@ export const installTemplate = async (
if (
props.template === "streaming" ||
props.template === "multiagent" ||
props.template === "extractor"
props.template === "reflex"
) {
await createBackendEnvFile(props.root, props);
}
Expand Down
44 changes: 24 additions & 20 deletions helpers/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -405,8 +405,8 @@ export const installPythonTemplate = async ({
>) => {
console.log("\nInitializing Python project with template:", template, "\n");
let templatePath;
if (template === "extractor") {
templatePath = path.join(templatesDir, "types", "extractor", framework);
if (template === "reflex") {
templatePath = path.join(templatesDir, "types", "reflex");
} else {
templatePath = path.join(templatesDir, "types", "streaming", framework);
}
Expand Down Expand Up @@ -472,24 +472,6 @@ export const installPythonTemplate = async ({
cwd: path.join(compPath, "engines", "python", engine),
});

// Copy agent code
if (template === "multiagent") {
if (agents) {
await copy("**", path.join(root), {
parents: true,
cwd: path.join(compPath, "agents", "python", agents),
rename: assetRelocator,
});
} else {
console.log(
red(
"There is no agent selected for multi-agent template. Please pick an agent to use via --agents flag.",
),
);
process.exit(1);
}
}

// Copy router code
await copyRouterCode(root, tools ?? []);
}
Expand All @@ -503,6 +485,28 @@ export const installPythonTemplate = async ({
});
}

if (template === "multiagent" || template === "reflex") {
if (agents) {
const sourcePath =
template === "multiagent"
? path.join(compPath, "agents", "python", agents)
: path.join(compPath, "reflex", agents);

await copy("**", path.join(root), {
parents: true,
cwd: sourcePath,
rename: assetRelocator,
});
} else {
console.log(
red(
`There is no agent selected for ${template} template. Please pick an agent to use via --agents flag.`,
),
);
process.exit(1);
}
}

console.log("Adding additional dependencies");

const addOnDependencies = getAdditionalDependencies(
Expand Down
8 changes: 4 additions & 4 deletions helpers/run-app.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { SpawnOptions, spawn } from "child_process";
import { TemplateFramework } from "./types";
import { TemplateFramework, TemplateType } from "./types";

const createProcess = (
command: string,
Expand Down Expand Up @@ -58,17 +58,17 @@ export function runTSApp(appPath: string, port: number) {

export async function runApp(
appPath: string,
template: string,
template: TemplateType,
framework: TemplateFramework,
port?: number,
): Promise<void> {
try {
// Start the app
const defaultPort =
framework === "nextjs" || template === "extractor" ? 3000 : 8000;
framework === "nextjs" || template === "reflex" ? 3000 : 8000;

const appRunner =
template === "extractor"
template === "reflex"
? runReflexApp
: framework === "fastapi"
? runFastAPIApp
Expand Down
13 changes: 10 additions & 3 deletions helpers/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,11 @@ export type ModelConfig = {
isConfigured(): boolean;
};
export type TemplateType =
| "extractor"
| "streaming"
| "community"
| "llamapack"
| "multiagent";
| "multiagent"
| "reflex";
export type TemplateFramework = "nextjs" | "express" | "fastapi";
export type TemplateUI = "html" | "shadcn";
export type TemplateVectorDB =
Expand All @@ -49,14 +49,21 @@ export type TemplateDataSource = {
};
export type TemplateDataSourceType = "file" | "web" | "db";
export type TemplateObservability = "none" | "traceloop" | "llamatrace";
export type TemplateAgents = "financial_report" | "blog" | "form_filling";
export type TemplateAgents =
| "financial_report"
| "blog"
| "form_filling"
| "extractor"
| "contract_review";
// Config for both file and folder
export type FileSourceConfig =
| {
path: string;
filename?: string;
}
| {
url: URL;
filename?: string;
};
export type WebSourceConfig = {
baseUrl?: string;
Expand Down
2 changes: 1 addition & 1 deletion helpers/typescript.ts
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ export const installTSTemplate = async ({
} else {
console.log(
red(
"There is no agent selected for multi-agent template. Please pick an agent to use via --agents flag.",
`There is no agent selected for ${template} template. Please pick an agent to use via --agents flag.`,
),
);
process.exit(1);
Expand Down
2 changes: 1 addition & 1 deletion index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ const options = program.opts();

if (
process.argv.includes("--no-llama-parse") ||
options.template === "extractor"
options.template === "reflex"
) {
options.useLlamaParse = false;
}
Expand Down
2 changes: 1 addition & 1 deletion questions/datasources.ts
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ export const getDataSourceChoices = (
);
}

if (framework === "fastapi" && template !== "extractor") {
if (framework === "fastapi" && template !== "reflex") {
choices.push({
title: "Use website content (requires Chrome)",
value: "web",
Expand Down
11 changes: 4 additions & 7 deletions questions/questions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,8 @@ export const askProQuestions = async (program: QuestionArgs) => {
return; // early return - no further questions needed for llamapack projects
}

if (program.template === "extractor") {
// Extractor template only supports FastAPI, empty data sources, and llamacloud
if (program.template === "reflex") {
// Reflex template only supports FastAPI, empty data sources, and llamacloud
// So we just use example file for extractor template, this allows user to choose vector database later
program.dataSources = [EXAMPLE_FILE];
program.framework = "fastapi";
Expand Down Expand Up @@ -354,11 +354,8 @@ export const askProQuestions = async (program: QuestionArgs) => {
// default to use LlamaParse if using LlamaCloud
program.useLlamaParse = true;
} else {
// Extractor template doesn't support LlamaParse and LlamaCloud right now (cannot use asyncio loop in Reflex)
if (
program.useLlamaParse === undefined &&
program.template !== "extractor"
) {
// Reflex template doesn't support LlamaParse and LlamaCloud right now (cannot use asyncio loop in Reflex)
if (program.useLlamaParse === undefined && program.template !== "reflex") {
// if already set useLlamaParse, don't ask again
if (program.dataSources.some((ds) => ds.type === "file")) {
const { useLlamaParse } = await prompts(
Expand Down
Loading

0 comments on commit 9077cae

Please sign in to comment.