mirror of
https://github.com/run-llama/create-llama.git
synced 2026-07-16 03:04:21 -04:00
Compare commits
53 Commits
ms/add-ruff
...
v0.2.8
| Author | SHA1 | Date | |
|---|---|---|---|
| b31fa80326 | |||
| 40c5c8412c | |||
| 0031e674c9 | |||
| 6e9184dd37 | |||
| fa28cb5d0d | |||
| 8c1087f5f1 | |||
| 27333973f1 | |||
| cf3ec97a4c | |||
| 505b8e944a | |||
| 578f7f9e50 | |||
| adc40cf770 | |||
| 7bce7386d5 | |||
| c011455dc4 | |||
| 38a8be8d12 | |||
| 6e70eb4d11 | |||
| 917e862202 | |||
| e363bfeecc | |||
| b6da3c2419 | |||
| 71fbe1b18f | |||
| 8105c5cf06 | |||
| c16deed864 | |||
| 6a409cbbc6 | |||
| a1892bef26 | |||
| 2f7e0220b5 | |||
| 435109fef0 | |||
| b1f3d5222f | |||
| e2c61884ef | |||
| fd4abb3bdd | |||
| bedde2bf20 | |||
| 5cd12fa90d | |||
| 72b71952aa | |||
| 2f8feabcba | |||
| a8a8c247e2 | |||
| 4fa2b76f3d | |||
| 4ead8e14c2 | |||
| 90398400c6 | |||
| 8f670a935c | |||
| f04f60d555 | |||
| 1ffd3c915b | |||
| 57e7638083 | |||
| 22ac2cae61 | |||
| 8077195601 | |||
| 8ce4a8513d | |||
| 1d93775f04 | |||
| 3fb93c7939 | |||
| e248dc56bc | |||
| bd5e39a390 | |||
| de2c7523dd | |||
| 9fd832c8b0 | |||
| b2c76dc7b6 | |||
| 2b7a5d8797 | |||
| d93ec803f5 | |||
| a6023b695b |
@@ -1,5 +0,0 @@
|
||||
---
|
||||
"create-llama": patch
|
||||
---
|
||||
|
||||
Use LlamaCloud pipeline for data ingestion (private file uploads and generate script)
|
||||
@@ -30,3 +30,13 @@ jobs:
|
||||
|
||||
- name: Run Prettier
|
||||
run: pnpm run format
|
||||
|
||||
- name: Run Python format check
|
||||
uses: chartboost/ruff-action@v1
|
||||
with:
|
||||
args: "format --check"
|
||||
|
||||
- name: Run Python lint
|
||||
uses: chartboost/ruff-action@v1
|
||||
with:
|
||||
args: "check"
|
||||
|
||||
+120
@@ -1,5 +1,125 @@
|
||||
# create-llama
|
||||
|
||||
## 0.2.8
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 0031e67: Bump llama-index to 0.11.11 for the multi-agent template
|
||||
|
||||
## 0.2.7
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 505b8e9: bump: use latest ai package version
|
||||
- cf3ec97: Dynamically select model for Groq
|
||||
- 8c1087f: feat: enhance style for markdown
|
||||
|
||||
## 0.2.6
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- adc40cf: fix: vercel ai update crash sending annotations
|
||||
|
||||
## 0.2.5
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 38a8be8: fix: filter in mongo vector store
|
||||
|
||||
## 0.2.4
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 917e862: Fix errors in building the frontend
|
||||
|
||||
## 0.2.3
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- b6da3c2: Ensure the generation script always works
|
||||
|
||||
## 0.2.2
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 8105c5c: Add env config for next questions feature
|
||||
|
||||
## 0.2.1
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 6a409cb: Bump web and database reader packages
|
||||
|
||||
## 0.2.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 435109f: Add multi-agents template based on workflows
|
||||
|
||||
## 0.1.44
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- bedde2b: Change metadata filters to use already existing documents in LlamaCloud Index
|
||||
- 5cd12fa: Use one callback manager per request
|
||||
- 5cd12fa: Bump llama_index version to 0.11.1
|
||||
- fd4abb3: Fix to use filename for uploaded documents in NextJS
|
||||
- 2f8feab: Simplify CLI interface
|
||||
|
||||
## 0.1.43
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 4fa2b76: feat: implement citation for TS
|
||||
|
||||
## 0.1.42
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 8f670a9: Allow relative URL in documents
|
||||
|
||||
## 0.1.41
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 57e7638: Use the retrieval defaults from LlamaCloud
|
||||
|
||||
## 0.1.40
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 8ce4a85: Add UI for extractor template
|
||||
|
||||
## 0.1.39
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 3fb93c7: Use LlamaCloud pipeline for data ingestion in TS (private file uploads and generate script)
|
||||
|
||||
## 0.1.38
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- bd5e39a: Fix error that files in sub folders of 'data' are not displayed
|
||||
|
||||
## 0.1.37
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 9fd832c: Add in-text citation references
|
||||
|
||||
## 0.1.36
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 2b7a5d8: Fix: private file upload not working in Python without LlamaCloud
|
||||
|
||||
## 0.1.35
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 81ef7f0: Use LlamaCloud pipeline for data ingestion (private file uploads and generate script)
|
||||
|
||||
## 0.1.34
|
||||
|
||||
### Patch Changes
|
||||
|
||||
@@ -94,7 +94,7 @@ Need to install the following packages:
|
||||
create-llama@latest
|
||||
Ok to proceed? (y) y
|
||||
✔ What is your project named? … my-app
|
||||
✔ Which template would you like to use? › Agentic RAG (single agent)
|
||||
✔ Which template would you like to use? › Agentic RAG (e.g. chat with docs)
|
||||
✔ Which framework would you like to use? › NextJS
|
||||
✔ Would you like to set up observability? › No
|
||||
✔ Please provide your OpenAI API key (leave blank to skip): …
|
||||
|
||||
@@ -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 { 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";
|
||||
|
||||
// The extractor template currently only works with FastAPI and files (and not on Windows)
|
||||
if (
|
||||
process.platform !== "win32" &&
|
||||
templateFramework !== "nextjs" &&
|
||||
templateFramework !== "express" &&
|
||||
dataSource !== "--no-files"
|
||||
) {
|
||||
test.describe("Test extractor template", async () => {
|
||||
let frontendPort: number;
|
||||
let backendPort: number;
|
||||
let name: string;
|
||||
let appProcess: ChildProcess;
|
||||
let cwd: string;
|
||||
|
||||
// Create extractor app
|
||||
test.beforeAll(async () => {
|
||||
cwd = await createTestDir();
|
||||
frontendPort = Math.floor(Math.random() * 10000) + 10000;
|
||||
backendPort = frontendPort + 1;
|
||||
const result = await runCreateLlama(
|
||||
cwd,
|
||||
"extractor",
|
||||
"fastapi",
|
||||
"--example-file",
|
||||
"none",
|
||||
frontendPort,
|
||||
backendPort,
|
||||
"runApp",
|
||||
);
|
||||
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:${frontendPort}`);
|
||||
await expect(page.getByText("Built by LlamaIndex")).toBeVisible({
|
||||
timeout: 2000 * 60,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
/* 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 type {
|
||||
TemplateFramework,
|
||||
TemplatePostInstallAction,
|
||||
TemplateUI,
|
||||
} from "../helpers";
|
||||
import { createTestDir, runCreateLlama, type AppType } from "./utils";
|
||||
|
||||
const templateFramework: TemplateFramework = "fastapi";
|
||||
const dataSource: string = "--example-file";
|
||||
const templateUI: TemplateUI = "shadcn";
|
||||
const templatePostInstallAction: TemplatePostInstallAction = "runApp";
|
||||
const appType: AppType = "--frontend";
|
||||
const userMessage = "Write a blog post about physical standards for letters";
|
||||
|
||||
test.describe(`Test multiagent template ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
|
||||
test.skip(
|
||||
process.platform !== "linux" ||
|
||||
process.env.FRAMEWORK !== "fastapi" ||
|
||||
process.env.DATASOURCE === "--no-files",
|
||||
"The multiagent template currently only works with FastAPI and files. We also only run on Linux to speed up tests.",
|
||||
);
|
||||
let port: number;
|
||||
let externalPort: number;
|
||||
let cwd: string;
|
||||
let name: string;
|
||||
let appProcess: ChildProcess;
|
||||
// Only test without using vector db for now
|
||||
const vectorDb = "none";
|
||||
|
||||
test.beforeAll(async () => {
|
||||
port = Math.floor(Math.random() * 10000) + 10000;
|
||||
externalPort = port + 1;
|
||||
cwd = await createTestDir();
|
||||
const result = await runCreateLlama(
|
||||
cwd,
|
||||
"multiagent",
|
||||
templateFramework,
|
||||
dataSource,
|
||||
vectorDb,
|
||||
port,
|
||||
externalPort,
|
||||
templatePostInstallAction,
|
||||
templateUI,
|
||||
appType,
|
||||
);
|
||||
name = result.projectName;
|
||||
appProcess = result.appProcess;
|
||||
});
|
||||
|
||||
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:${port}`);
|
||||
await expect(page.getByText("Built by LlamaIndex")).toBeVisible();
|
||||
});
|
||||
|
||||
test("Frontend should be able to submit a message and receive the start of a streamed response", async ({
|
||||
page,
|
||||
}) => {
|
||||
await page.goto(`http://localhost:${port}`);
|
||||
await page.fill("form input", userMessage);
|
||||
|
||||
const responsePromise = page.waitForResponse((res) =>
|
||||
res.url().includes("/api/chat"),
|
||||
);
|
||||
|
||||
await page.click("form button[type=submit]");
|
||||
|
||||
const response = await responsePromise;
|
||||
expect(response.ok()).toBeTruthy();
|
||||
});
|
||||
|
||||
// clean processes
|
||||
test.afterAll(async () => {
|
||||
appProcess?.kill();
|
||||
});
|
||||
});
|
||||
@@ -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 ",
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -6,12 +6,10 @@ import path from "path";
|
||||
import type {
|
||||
TemplateFramework,
|
||||
TemplatePostInstallAction,
|
||||
TemplateType,
|
||||
TemplateUI,
|
||||
} from "../helpers";
|
||||
import { createTestDir, runCreateLlama, type AppType } from "./utils";
|
||||
|
||||
const templateType: TemplateType = "streaming";
|
||||
const templateFramework: TemplateFramework = process.env.FRAMEWORK
|
||||
? (process.env.FRAMEWORK as TemplateFramework)
|
||||
: "fastapi";
|
||||
@@ -27,7 +25,8 @@ const llamaCloudIndexName = "e2e-test";
|
||||
const appType: AppType = templateFramework === "nextjs" ? "" : "--frontend";
|
||||
const userMessage =
|
||||
dataSource !== "--no-files" ? "Physical standard for letters" : "Hello";
|
||||
test.describe(`try create-llama ${templateType} ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
|
||||
|
||||
test.describe(`Test streaming template ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
|
||||
let port: number;
|
||||
let externalPort: number;
|
||||
let cwd: string;
|
||||
@@ -42,15 +41,15 @@ test.describe(`try create-llama ${templateType} ${templateFramework} ${dataSourc
|
||||
cwd = await createTestDir();
|
||||
const result = await runCreateLlama(
|
||||
cwd,
|
||||
templateType,
|
||||
"streaming",
|
||||
templateFramework,
|
||||
dataSource,
|
||||
templateUI,
|
||||
vectorDb,
|
||||
appType,
|
||||
port,
|
||||
externalPort,
|
||||
templatePostInstallAction,
|
||||
templateUI,
|
||||
appType,
|
||||
llamaCloudProjectName,
|
||||
llamaCloudIndexName,
|
||||
);
|
||||
+32
-12
@@ -24,14 +24,15 @@ export async function runCreateLlama(
|
||||
templateType: TemplateType,
|
||||
templateFramework: TemplateFramework,
|
||||
dataSource: string,
|
||||
templateUI: TemplateUI,
|
||||
vectorDb: TemplateVectorDB,
|
||||
appType: AppType,
|
||||
port: number,
|
||||
externalPort: number,
|
||||
postInstallAction: TemplatePostInstallAction,
|
||||
llamaCloudProjectName: string,
|
||||
llamaCloudIndexName: string,
|
||||
templateUI?: TemplateUI,
|
||||
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(
|
||||
@@ -41,25 +42,35 @@ export async function runCreateLlama(
|
||||
const name = [
|
||||
templateType,
|
||||
templateFramework,
|
||||
dataSource,
|
||||
dataSource.split(" ")[0],
|
||||
templateUI,
|
||||
appType,
|
||||
].join("-");
|
||||
const command = [
|
||||
|
||||
// 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,
|
||||
"--ui",
|
||||
templateUI,
|
||||
...dataSourceArgs,
|
||||
"--vector-db",
|
||||
vectorDb,
|
||||
"--open-ai-key",
|
||||
process.env.OPENAI_API_KEY,
|
||||
appType,
|
||||
"--use-pnpm",
|
||||
"--port",
|
||||
port,
|
||||
@@ -68,13 +79,22 @@ export async function runCreateLlama(
|
||||
"--post-install-action",
|
||||
postInstallAction,
|
||||
"--tools",
|
||||
"none",
|
||||
tools ?? "none",
|
||||
"--no-llama-parse",
|
||||
"--observability",
|
||||
"none",
|
||||
"--llama-cloud-key",
|
||||
process.env.LLAMA_CLOUD_API_KEY,
|
||||
].join(" ");
|
||||
];
|
||||
|
||||
if (templateUI) {
|
||||
commandArgs.push("--ui", templateUI);
|
||||
}
|
||||
if (appType) {
|
||||
commandArgs.push(appType);
|
||||
}
|
||||
|
||||
const command = commandArgs.join(" ");
|
||||
console.log(`running command '${command}' in ${cwd}`);
|
||||
const appProcess = exec(command, {
|
||||
cwd,
|
||||
|
||||
+52
-60
@@ -36,74 +36,66 @@ export async function writeLoadersConfig(
|
||||
dataSources: TemplateDataSource[],
|
||||
useLlamaParse?: boolean,
|
||||
) {
|
||||
if (dataSources.length === 0) return; // no datasources, no config needed
|
||||
const loaderConfig = new Document({});
|
||||
// Web loader config
|
||||
const loaderConfig: Record<string, any> = {};
|
||||
|
||||
// Always set file loader config
|
||||
loaderConfig.file = createFileLoaderConfig(useLlamaParse);
|
||||
|
||||
if (dataSources.some((ds) => ds.type === "web")) {
|
||||
const webLoaderConfig = new Document({});
|
||||
|
||||
// Create config for browser driver arguments
|
||||
const driverArgNodeValue = webLoaderConfig.createNode([
|
||||
"--no-sandbox",
|
||||
"--disable-dev-shm-usage",
|
||||
]);
|
||||
driverArgNodeValue.commentBefore =
|
||||
" The arguments to pass to the webdriver. E.g.: add --headless to run in headless mode";
|
||||
webLoaderConfig.set("driver_arguments", driverArgNodeValue);
|
||||
|
||||
// Create config for urls
|
||||
const urlConfigs = dataSources
|
||||
.filter((ds) => ds.type === "web")
|
||||
.map((ds) => {
|
||||
const dsConfig = ds.config as WebSourceConfig;
|
||||
return {
|
||||
base_url: dsConfig.baseUrl,
|
||||
prefix: dsConfig.prefix,
|
||||
depth: dsConfig.depth,
|
||||
};
|
||||
});
|
||||
const urlConfigNode = webLoaderConfig.createNode(urlConfigs);
|
||||
urlConfigNode.commentBefore = ` base_url: The URL to start crawling with
|
||||
prefix: Only crawl URLs matching the specified prefix
|
||||
depth: The maximum depth for BFS traversal
|
||||
You can add more websites by adding more entries (don't forget the - prefix from YAML)`;
|
||||
webLoaderConfig.set("urls", urlConfigNode);
|
||||
|
||||
// Add web config to the loaders config
|
||||
loaderConfig.set("web", webLoaderConfig);
|
||||
loaderConfig.web = createWebLoaderConfig(dataSources);
|
||||
}
|
||||
|
||||
// File loader config
|
||||
if (dataSources.some((ds) => ds.type === "file")) {
|
||||
// Add documentation to web loader config
|
||||
const node = loaderConfig.createNode({
|
||||
use_llama_parse: useLlamaParse,
|
||||
});
|
||||
node.commentBefore = ` use_llama_parse: Use LlamaParse if \`true\`. Needs a \`LLAMA_CLOUD_API_KEY\` from https://cloud.llamaindex.ai set as environment variable`;
|
||||
loaderConfig.set("file", node);
|
||||
}
|
||||
|
||||
// DB loader config
|
||||
const dbLoaders = dataSources.filter((ds) => ds.type === "db");
|
||||
if (dbLoaders.length > 0) {
|
||||
const dbLoaderConfig = new Document({});
|
||||
const configEntries = dbLoaders.map((ds) => {
|
||||
const dsConfig = ds.config as DbSourceConfig;
|
||||
return {
|
||||
uri: dsConfig.uri,
|
||||
queries: [dsConfig.queries],
|
||||
};
|
||||
});
|
||||
|
||||
const node = dbLoaderConfig.createNode(configEntries);
|
||||
node.commentBefore = ` The configuration for the database loader, only supports MySQL and PostgreSQL databases for now.
|
||||
uri: The URI for the database. E.g.: mysql+pymysql://user:password@localhost:3306/db or postgresql+psycopg2://user:password@localhost:5432/db
|
||||
query: The query to fetch data from the database. E.g.: SELECT * FROM table`;
|
||||
loaderConfig.set("db", node);
|
||||
loaderConfig.db = createDbLoaderConfig(dbLoaders);
|
||||
}
|
||||
|
||||
// Create a new Document with the loaderConfig
|
||||
const yamlDoc = new Document(loaderConfig);
|
||||
|
||||
// Write loaders config
|
||||
const loaderConfigPath = path.join(root, "config", "loaders.yaml");
|
||||
await fs.mkdir(path.join(root, "config"), { recursive: true });
|
||||
await fs.writeFile(loaderConfigPath, yaml.stringify(loaderConfig));
|
||||
await fs.writeFile(loaderConfigPath, yaml.stringify(yamlDoc));
|
||||
}
|
||||
|
||||
function createWebLoaderConfig(dataSources: TemplateDataSource[]): any {
|
||||
const webLoaderConfig: Record<string, any> = {};
|
||||
|
||||
// Create config for browser driver arguments
|
||||
webLoaderConfig.driver_arguments = [
|
||||
"--no-sandbox",
|
||||
"--disable-dev-shm-usage",
|
||||
];
|
||||
|
||||
// Create config for urls
|
||||
const urlConfigs = dataSources
|
||||
.filter((ds) => ds.type === "web")
|
||||
.map((ds) => {
|
||||
const dsConfig = ds.config as WebSourceConfig;
|
||||
return {
|
||||
base_url: dsConfig.baseUrl,
|
||||
prefix: dsConfig.prefix,
|
||||
depth: dsConfig.depth,
|
||||
};
|
||||
});
|
||||
webLoaderConfig.urls = urlConfigs;
|
||||
|
||||
return webLoaderConfig;
|
||||
}
|
||||
|
||||
function createFileLoaderConfig(useLlamaParse?: boolean): any {
|
||||
return {
|
||||
use_llama_parse: useLlamaParse,
|
||||
};
|
||||
}
|
||||
|
||||
function createDbLoaderConfig(dbLoaders: TemplateDataSource[]): any {
|
||||
return dbLoaders.map((ds) => {
|
||||
const dsConfig = ds.config as DbSourceConfig;
|
||||
return {
|
||||
uri: dsConfig.uri,
|
||||
queries: [dsConfig.queries],
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
+69
-35
@@ -4,6 +4,7 @@ import { TOOL_SYSTEM_PROMPT_ENV_VAR, Tool } from "./tools";
|
||||
import {
|
||||
InstallTemplateArgs,
|
||||
ModelConfig,
|
||||
TemplateDataSource,
|
||||
TemplateFramework,
|
||||
TemplateObservability,
|
||||
TemplateType,
|
||||
@@ -159,7 +160,7 @@ const getVectorDBEnvs = (
|
||||
{
|
||||
name: "LLAMA_CLOUD_ORGANIZATION_ID",
|
||||
description:
|
||||
"The organization ID for the LlamaCloud project (uses default organization if not specified - Python only)",
|
||||
"The organization ID for the LlamaCloud project (uses default organization if not specified)",
|
||||
},
|
||||
...(framework === "nextjs"
|
||||
? // activate index selector per default (not needed for non-NextJS backends as it's handled by createFrontendEnvFile)
|
||||
@@ -395,7 +396,6 @@ const getEngineEnvs = (): EnvVar[] => {
|
||||
name: "TOP_K",
|
||||
description:
|
||||
"The number of similar embeddings to return when retrieving documents.",
|
||||
value: "3",
|
||||
},
|
||||
{
|
||||
name: "STREAM_TIMEOUT",
|
||||
@@ -423,7 +423,11 @@ const getToolEnvs = (tools?: Tool[]): EnvVar[] => {
|
||||
return toolEnvs;
|
||||
};
|
||||
|
||||
const getSystemPromptEnv = (tools?: Tool[]): EnvVar => {
|
||||
const getSystemPromptEnv = (
|
||||
tools?: Tool[],
|
||||
dataSources?: TemplateDataSource[],
|
||||
framework?: TemplateFramework,
|
||||
): EnvVar[] => {
|
||||
const defaultSystemPrompt =
|
||||
"You are a helpful assistant who helps users with their questions.";
|
||||
|
||||
@@ -442,41 +446,71 @@ const getSystemPromptEnv = (tools?: Tool[]): EnvVar => {
|
||||
? `\"${toolSystemPrompt}\"`
|
||||
: defaultSystemPrompt;
|
||||
|
||||
return {
|
||||
name: "SYSTEM_PROMPT",
|
||||
description: "The system prompt for the AI model.",
|
||||
value: systemPrompt,
|
||||
};
|
||||
const systemPromptEnv = [
|
||||
{
|
||||
name: "SYSTEM_PROMPT",
|
||||
description: "The system prompt for the AI model.",
|
||||
value: systemPrompt,
|
||||
},
|
||||
];
|
||||
|
||||
if (tools?.length == 0 && (dataSources?.length ?? 0 > 0)) {
|
||||
const citationPrompt = `'You have provided information from a knowledge base that has been passed to you in nodes of information.
|
||||
Each node has useful metadata such as node ID, file name, page, etc.
|
||||
Please add the citation to the data node for each sentence or paragraph that you reference in the provided information.
|
||||
The citation format is: . [citation:<node_id>]()
|
||||
Where the <node_id> is the unique identifier of the data node.
|
||||
|
||||
Example:
|
||||
We have two nodes:
|
||||
node_id: xyz
|
||||
file_name: llama.pdf
|
||||
|
||||
node_id: abc
|
||||
file_name: animal.pdf
|
||||
|
||||
User question: Tell me a fun fact about Llama.
|
||||
Your answer:
|
||||
A baby llama is called "Cria" [citation:xyz]().
|
||||
It often live in desert [citation:abc]().
|
||||
It\\'s cute animal.
|
||||
'`;
|
||||
systemPromptEnv.push({
|
||||
name: "SYSTEM_CITATION_PROMPT",
|
||||
description:
|
||||
"An additional system prompt to add citation when responding to user questions.",
|
||||
value: citationPrompt,
|
||||
});
|
||||
}
|
||||
|
||||
return systemPromptEnv;
|
||||
};
|
||||
|
||||
const getTemplateEnvs = (template?: TemplateType): EnvVar[] => {
|
||||
if (template === "multiagent") {
|
||||
return [
|
||||
{
|
||||
name: "MESSAGE_QUEUE_PORT",
|
||||
},
|
||||
{
|
||||
name: "CONTROL_PLANE_PORT",
|
||||
},
|
||||
{
|
||||
name: "HUMAN_CONSUMER_PORT",
|
||||
},
|
||||
{
|
||||
name: "AGENT_QUERY_ENGINE_PORT",
|
||||
value: "8003",
|
||||
},
|
||||
{
|
||||
name: "AGENT_QUERY_ENGINE_DESCRIPTION",
|
||||
value: "Query information from the provided data",
|
||||
},
|
||||
{
|
||||
name: "AGENT_DUMMY_PORT",
|
||||
value: "8004",
|
||||
},
|
||||
];
|
||||
} else {
|
||||
return [];
|
||||
const nextQuestionEnvs: EnvVar[] = [
|
||||
{
|
||||
name: "NEXT_QUESTION_PROMPT",
|
||||
description: `Customize prompt to generate the next question suggestions based on the conversation history.
|
||||
Disable this prompt to disable the next question suggestions feature.`,
|
||||
value: `"You're a helpful assistant! Your task is to suggest the next question that user might ask.
|
||||
Here is the conversation history
|
||||
---------------------
|
||||
{conversation}
|
||||
---------------------
|
||||
Given the conversation history, please give me 3 questions that you might ask next!
|
||||
Your answer should be wrapped in three sticks which follows the following format:
|
||||
\`\`\`
|
||||
<question 1>
|
||||
<question 2>
|
||||
<question 3>
|
||||
\`\`\`"`,
|
||||
},
|
||||
];
|
||||
|
||||
if (template === "multiagent" || template === "streaming") {
|
||||
return nextQuestionEnvs;
|
||||
}
|
||||
return [];
|
||||
};
|
||||
|
||||
const getObservabilityEnvs = (
|
||||
@@ -525,7 +559,7 @@ export const createBackendEnvFile = async (
|
||||
...getToolEnvs(opts.tools),
|
||||
...getTemplateEnvs(opts.template),
|
||||
...getObservabilityEnvs(opts.observability),
|
||||
getSystemPromptEnv(opts.tools),
|
||||
...getSystemPromptEnv(opts.tools, opts.dataSources, opts.framework),
|
||||
];
|
||||
// Render and write env file
|
||||
const content = renderEnvVar(envVars);
|
||||
|
||||
+19
-18
@@ -96,10 +96,11 @@ async function generateContextData(
|
||||
}
|
||||
}
|
||||
|
||||
const copyContextData = async (
|
||||
const prepareContextData = async (
|
||||
root: string,
|
||||
dataSources: TemplateDataSource[],
|
||||
) => {
|
||||
await makeDir(path.join(root, "data"));
|
||||
for (const dataSource of dataSources) {
|
||||
const dataSourceConfig = dataSource?.config as FileSourceConfig;
|
||||
// Copy local data
|
||||
@@ -174,25 +175,25 @@ export const installTemplate = async (
|
||||
await createBackendEnvFile(props.root, props);
|
||||
}
|
||||
|
||||
if (props.dataSources.length > 0) {
|
||||
await prepareContextData(
|
||||
props.root,
|
||||
props.dataSources.filter((ds) => ds.type === "file"),
|
||||
);
|
||||
|
||||
if (
|
||||
props.dataSources.length > 0 &&
|
||||
(props.postInstallAction === "runApp" ||
|
||||
props.postInstallAction === "dependencies")
|
||||
) {
|
||||
console.log("\nGenerating context data...\n");
|
||||
await copyContextData(
|
||||
props.root,
|
||||
props.dataSources.filter((ds) => ds.type === "file"),
|
||||
await generateContextData(
|
||||
props.framework,
|
||||
props.modelConfig,
|
||||
props.packageManager,
|
||||
props.vectorDb,
|
||||
props.llamaCloudKey,
|
||||
props.useLlamaParse,
|
||||
);
|
||||
if (
|
||||
props.postInstallAction === "runApp" ||
|
||||
props.postInstallAction === "dependencies"
|
||||
) {
|
||||
await generateContextData(
|
||||
props.framework,
|
||||
props.modelConfig,
|
||||
props.packageManager,
|
||||
props.vectorDb,
|
||||
props.llamaCloudKey,
|
||||
props.useLlamaParse,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Create outputs directory
|
||||
|
||||
@@ -3,8 +3,55 @@ import prompts from "prompts";
|
||||
import { ModelConfigParams } from ".";
|
||||
import { questionHandlers, toChoice } from "../../questions";
|
||||
|
||||
const MODELS = ["llama3-8b", "llama3-70b", "mixtral-8x7b"];
|
||||
const DEFAULT_MODEL = MODELS[0];
|
||||
import got from "got";
|
||||
import ora from "ora";
|
||||
import { red } from "picocolors";
|
||||
|
||||
const GROQ_API_URL = "https://api.groq.com/openai/v1";
|
||||
|
||||
async function getAvailableModelChoicesGroq(apiKey: string) {
|
||||
if (!apiKey) {
|
||||
throw new Error("Need Groq API key to retrieve model choices");
|
||||
}
|
||||
|
||||
const spinner = ora("Fetching available models from Groq").start();
|
||||
try {
|
||||
const response = await got(`${GROQ_API_URL}/models`, {
|
||||
headers: {
|
||||
Authorization: `Bearer ${apiKey}`,
|
||||
},
|
||||
timeout: 5000,
|
||||
responseType: "json",
|
||||
});
|
||||
const data: any = await response.body;
|
||||
spinner.stop();
|
||||
|
||||
// Filter out the Whisper models
|
||||
return data.data
|
||||
.filter((model: any) => !model.id.toLowerCase().includes("whisper"))
|
||||
.map((el: any) => {
|
||||
return {
|
||||
title: el.id,
|
||||
value: el.id,
|
||||
};
|
||||
});
|
||||
} catch (error: unknown) {
|
||||
spinner.stop();
|
||||
console.log(error);
|
||||
if ((error as any).response?.statusCode === 401) {
|
||||
console.log(
|
||||
red(
|
||||
"Invalid Groq API key provided! Please provide a valid key and try again!",
|
||||
),
|
||||
);
|
||||
} else {
|
||||
console.log(red("Request failed: " + error));
|
||||
}
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
const DEFAULT_MODEL = "llama3-70b-8192";
|
||||
|
||||
// Use huggingface embedding models for now as Groq doesn't support embedding models
|
||||
enum HuggingFaceEmbeddingModelType {
|
||||
@@ -66,12 +113,14 @@ export async function askGroqQuestions({
|
||||
// use default model values in CI or if user should not be asked
|
||||
const useDefaults = ciInfo.isCI || !askModels;
|
||||
if (!useDefaults) {
|
||||
const modelChoices = await getAvailableModelChoicesGroq(config.apiKey!);
|
||||
|
||||
const { model } = await prompts(
|
||||
{
|
||||
type: "select",
|
||||
name: "model",
|
||||
message: "Which LLM model would you like to use?",
|
||||
choices: MODELS.map(toChoice),
|
||||
choices: modelChoices,
|
||||
initial: 0,
|
||||
},
|
||||
questionHandlers,
|
||||
|
||||
+69
-31
@@ -12,6 +12,7 @@ import {
|
||||
InstallTemplateArgs,
|
||||
ModelConfig,
|
||||
TemplateDataSource,
|
||||
TemplateType,
|
||||
TemplateVectorDB,
|
||||
} from "./types";
|
||||
|
||||
@@ -26,6 +27,7 @@ const getAdditionalDependencies = (
|
||||
vectorDb?: TemplateVectorDB,
|
||||
dataSources?: TemplateDataSource[],
|
||||
tools?: Tool[],
|
||||
templateType?: TemplateType,
|
||||
) => {
|
||||
const dependencies: Dependency[] = [];
|
||||
|
||||
@@ -34,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",
|
||||
@@ -66,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;
|
||||
}
|
||||
@@ -107,13 +109,13 @@ const getAdditionalDependencies = (
|
||||
case "web":
|
||||
dependencies.push({
|
||||
name: "llama-index-readers-web",
|
||||
version: "^0.1.6",
|
||||
version: "^0.2.2",
|
||||
});
|
||||
break;
|
||||
case "db":
|
||||
dependencies.push({
|
||||
name: "llama-index-readers-database",
|
||||
version: "^0.1.3",
|
||||
version: "^0.2.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "pymysql",
|
||||
@@ -128,7 +130,7 @@ const getAdditionalDependencies = (
|
||||
case "llamacloud":
|
||||
dependencies.push({
|
||||
name: "llama-index-indices-managed-llama-cloud",
|
||||
version: "^0.2.7",
|
||||
version: "^0.3.1",
|
||||
});
|
||||
break;
|
||||
}
|
||||
@@ -147,77 +149,99 @@ const getAdditionalDependencies = (
|
||||
case "ollama":
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-ollama",
|
||||
version: "0.1.2",
|
||||
version: "0.3.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-ollama",
|
||||
version: "0.1.2",
|
||||
version: "0.3.0",
|
||||
});
|
||||
break;
|
||||
case "openai":
|
||||
dependencies.push({
|
||||
name: "llama-index-agent-openai",
|
||||
version: "0.2.6",
|
||||
});
|
||||
if (templateType !== "multiagent") {
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-openai",
|
||||
version: "^0.2.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-openai",
|
||||
version: "^0.2.3",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-agent-openai",
|
||||
version: "^0.3.0",
|
||||
});
|
||||
}
|
||||
break;
|
||||
case "groq":
|
||||
// Fastembed==0.2.0 does not support python3.13 at the moment
|
||||
// Fixed the python version less than 3.13
|
||||
dependencies.push({
|
||||
name: "python",
|
||||
version: "^3.11,<3.13",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-groq",
|
||||
version: "0.1.4",
|
||||
version: "0.2.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-fastembed",
|
||||
version: "^0.1.4",
|
||||
version: "^0.2.0",
|
||||
});
|
||||
break;
|
||||
case "anthropic":
|
||||
// Fastembed==0.2.0 does not support python3.13 at the moment
|
||||
// Fixed the python version less than 3.13
|
||||
dependencies.push({
|
||||
name: "python",
|
||||
version: "^3.11,<3.13",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-anthropic",
|
||||
version: "0.1.10",
|
||||
version: "0.3.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-fastembed",
|
||||
version: "^0.1.4",
|
||||
version: "^0.2.0",
|
||||
});
|
||||
break;
|
||||
case "gemini":
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-gemini",
|
||||
version: "0.1.10",
|
||||
version: "0.3.4",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-gemini",
|
||||
version: "0.1.6",
|
||||
version: "^0.2.0",
|
||||
});
|
||||
break;
|
||||
case "mistral":
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-mistralai",
|
||||
version: "0.1.17",
|
||||
version: "0.2.1",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-mistralai",
|
||||
version: "0.1.4",
|
||||
version: "0.2.0",
|
||||
});
|
||||
break;
|
||||
case "azure-openai":
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-azure-openai",
|
||||
version: "0.1.10",
|
||||
version: "0.2.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-embeddings-azure-openai",
|
||||
version: "0.1.11",
|
||||
version: "0.2.4",
|
||||
});
|
||||
break;
|
||||
case "t-systems":
|
||||
dependencies.push({
|
||||
name: "llama-index-agent-openai",
|
||||
version: "0.2.2",
|
||||
version: "0.3.0",
|
||||
});
|
||||
dependencies.push({
|
||||
name: "llama-index-llms-openai-like",
|
||||
version: "0.1.3",
|
||||
version: "0.2.0",
|
||||
});
|
||||
break;
|
||||
}
|
||||
@@ -227,7 +251,7 @@ const getAdditionalDependencies = (
|
||||
|
||||
const mergePoetryDependencies = (
|
||||
dependencies: Dependency[],
|
||||
existingDependencies: Record<string, Omit<Dependency, "name">>,
|
||||
existingDependencies: Record<string, Omit<Dependency, "name"> | string>,
|
||||
) => {
|
||||
for (const dependency of dependencies) {
|
||||
let value = existingDependencies[dependency.name] ?? {};
|
||||
@@ -246,7 +270,13 @@ const mergePoetryDependencies = (
|
||||
);
|
||||
}
|
||||
|
||||
existingDependencies[dependency.name] = value;
|
||||
// Serialize separately only if extras are provided
|
||||
if (value.extras && value.extras.length > 0) {
|
||||
existingDependencies[dependency.name] = value;
|
||||
} else {
|
||||
// Otherwise, serialize just the version string
|
||||
existingDependencies[dependency.name] = value.version;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -365,6 +395,13 @@ export const installPythonTemplate = async ({
|
||||
cwd: path.join(compPath, "settings", "python"),
|
||||
});
|
||||
|
||||
// Copy services
|
||||
if (template == "streaming" || template == "multiagent") {
|
||||
await copy("**", path.join(root, "app", "api", "services"), {
|
||||
cwd: path.join(compPath, "services", "python"),
|
||||
});
|
||||
}
|
||||
|
||||
if (template === "streaming") {
|
||||
// For the streaming template only:
|
||||
// Select and copy engine code based on data sources and tools
|
||||
@@ -388,6 +425,7 @@ export const installPythonTemplate = async ({
|
||||
vectorDb,
|
||||
dataSources,
|
||||
tools,
|
||||
template,
|
||||
);
|
||||
|
||||
if (observability && observability !== "none") {
|
||||
|
||||
+59
-48
@@ -23,66 +23,77 @@ const createProcess = (
|
||||
});
|
||||
};
|
||||
|
||||
// eslint-disable-next-line max-params
|
||||
export function runReflexApp(
|
||||
appPath: string,
|
||||
frontendPort?: number,
|
||||
backendPort?: number,
|
||||
) {
|
||||
const commandArgs = ["run", "reflex", "run"];
|
||||
if (frontendPort) {
|
||||
commandArgs.push("--frontend-port", frontendPort.toString());
|
||||
}
|
||||
if (backendPort) {
|
||||
commandArgs.push("--backend-port", backendPort.toString());
|
||||
}
|
||||
return createProcess("poetry", commandArgs, {
|
||||
stdio: "inherit",
|
||||
cwd: appPath,
|
||||
});
|
||||
}
|
||||
|
||||
export function runFastAPIApp(appPath: string, port: number) {
|
||||
const commandArgs = ["run", "uvicorn", "main:app", "--port=" + port];
|
||||
|
||||
return createProcess("poetry", commandArgs, {
|
||||
stdio: "inherit",
|
||||
cwd: appPath,
|
||||
});
|
||||
}
|
||||
|
||||
export function runTSApp(appPath: string, port: number) {
|
||||
return createProcess("npm", ["run", "dev"], {
|
||||
stdio: "inherit",
|
||||
cwd: appPath,
|
||||
env: { ...process.env, PORT: `${port}` },
|
||||
});
|
||||
}
|
||||
|
||||
export async function runApp(
|
||||
appPath: string,
|
||||
template: string,
|
||||
frontend: boolean,
|
||||
framework: TemplateFramework,
|
||||
port?: number,
|
||||
externalPort?: number,
|
||||
): Promise<any> {
|
||||
let backendAppProcess: ChildProcess;
|
||||
let frontendAppProcess: ChildProcess | undefined;
|
||||
const frontendPort = port || 3000;
|
||||
let backendPort = externalPort || 8000;
|
||||
const processes: ChildProcess[] = [];
|
||||
|
||||
// Callback to kill app processes
|
||||
// Callback to kill all sub processes if the main process is killed
|
||||
process.on("exit", () => {
|
||||
console.log("Killing app processes...");
|
||||
backendAppProcess.kill();
|
||||
frontendAppProcess?.kill();
|
||||
processes.forEach((p) => p.kill());
|
||||
});
|
||||
|
||||
let backendCommand = "";
|
||||
let backendArgs: string[];
|
||||
if (framework === "fastapi") {
|
||||
backendCommand = "poetry";
|
||||
backendArgs = [
|
||||
"run",
|
||||
"uvicorn",
|
||||
"main:app",
|
||||
"--host=0.0.0.0",
|
||||
"--port=" + backendPort,
|
||||
];
|
||||
} else if (framework === "nextjs") {
|
||||
backendCommand = "npm";
|
||||
backendArgs = ["run", "dev"];
|
||||
backendPort = frontendPort;
|
||||
} else {
|
||||
backendCommand = "npm";
|
||||
backendArgs = ["run", "dev"];
|
||||
// Default sub app paths
|
||||
const backendPath = path.join(appPath, "backend");
|
||||
const frontendPath = path.join(appPath, "frontend");
|
||||
|
||||
if (template === "extractor") {
|
||||
processes.push(runReflexApp(appPath, port, externalPort));
|
||||
}
|
||||
if (template === "streaming" || template === "multiagent") {
|
||||
if (framework === "fastapi" || framework === "express") {
|
||||
const backendRunner = framework === "fastapi" ? runFastAPIApp : runTSApp;
|
||||
if (frontend) {
|
||||
processes.push(backendRunner(backendPath, externalPort || 8000));
|
||||
processes.push(runTSApp(frontendPath, port || 3000));
|
||||
} else {
|
||||
processes.push(backendRunner(appPath, externalPort || 8000));
|
||||
}
|
||||
} else if (framework === "nextjs") {
|
||||
processes.push(runTSApp(appPath, port || 3000));
|
||||
}
|
||||
}
|
||||
|
||||
if (frontend) {
|
||||
return new Promise((resolve, reject) => {
|
||||
backendAppProcess = createProcess(backendCommand, backendArgs, {
|
||||
stdio: "inherit",
|
||||
cwd: path.join(appPath, "backend"),
|
||||
env: { ...process.env, PORT: `${backendPort}` },
|
||||
});
|
||||
frontendAppProcess = createProcess("npm", ["run", "dev"], {
|
||||
stdio: "inherit",
|
||||
cwd: path.join(appPath, "frontend"),
|
||||
env: { ...process.env, PORT: `${frontendPort}` },
|
||||
});
|
||||
});
|
||||
} else {
|
||||
return new Promise((resolve, reject) => {
|
||||
backendAppProcess = createProcess(backendCommand, backendArgs, {
|
||||
stdio: "inherit",
|
||||
cwd: path.join(appPath),
|
||||
env: { ...process.env, PORT: `${backendPort}` },
|
||||
});
|
||||
});
|
||||
}
|
||||
return Promise.all(processes);
|
||||
}
|
||||
|
||||
+4
-4
@@ -41,7 +41,7 @@ export const supportedTools: Tool[] = [
|
||||
dependencies: [
|
||||
{
|
||||
name: "llama-index-tools-google",
|
||||
version: "0.1.2",
|
||||
version: "^0.2.0",
|
||||
},
|
||||
],
|
||||
supportedFrameworks: ["fastapi"],
|
||||
@@ -83,7 +83,7 @@ For better results, you can specify the region parameter to get results from a s
|
||||
dependencies: [
|
||||
{
|
||||
name: "llama-index-tools-wikipedia",
|
||||
version: "0.1.2",
|
||||
version: "^0.2.0",
|
||||
},
|
||||
],
|
||||
supportedFrameworks: ["fastapi", "express", "nextjs"],
|
||||
@@ -145,7 +145,7 @@ For better results, you can specify the region parameter to get results from a s
|
||||
dependencies: [
|
||||
{
|
||||
name: "llama-index-tools-openapi",
|
||||
version: "0.1.3",
|
||||
version: "0.2.0",
|
||||
},
|
||||
{
|
||||
name: "jsonschema",
|
||||
@@ -153,7 +153,7 @@ For better results, you can specify the region parameter to get results from a s
|
||||
},
|
||||
{
|
||||
name: "llama-index-tools-requests",
|
||||
version: "0.1.3",
|
||||
version: "0.2.0",
|
||||
},
|
||||
],
|
||||
config: {
|
||||
|
||||
@@ -33,7 +33,8 @@ export const installTSTemplate = async ({
|
||||
* Copy the template files to the target directory.
|
||||
*/
|
||||
console.log("\nInitializing project with template:", template, "\n");
|
||||
const templatePath = path.join(templatesDir, "types", template, framework);
|
||||
const type = template === "multiagent" ? "streaming" : template; // use nextjs streaming template for multiagent
|
||||
const templatePath = path.join(templatesDir, "types", type, framework);
|
||||
const copySource = ["**"];
|
||||
|
||||
await copy(copySource, root, {
|
||||
|
||||
@@ -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(
|
||||
@@ -173,7 +187,14 @@ const program = new Commander.Command(packageJson.name)
|
||||
"--ask-models",
|
||||
`
|
||||
|
||||
Select LLM and embedding models.
|
||||
Allow interactive selection of LLM and embedding models of different model providers.
|
||||
`,
|
||||
)
|
||||
.option(
|
||||
"--ask-examples",
|
||||
`
|
||||
|
||||
Allow interactive selection of community templates and LlamaPacks.
|
||||
`,
|
||||
)
|
||||
.allowUnknownOption()
|
||||
@@ -188,10 +209,14 @@ if (process.argv.includes("--tools")) {
|
||||
program.tools = getTools(program.tools.split(","));
|
||||
}
|
||||
}
|
||||
if (process.argv.includes("--no-llama-parse")) {
|
||||
if (
|
||||
process.argv.includes("--no-llama-parse") ||
|
||||
program.template === "extractor"
|
||||
) {
|
||||
program.useLlamaParse = false;
|
||||
}
|
||||
program.askModels = process.argv.includes("--ask-models");
|
||||
program.askExamples = process.argv.includes("--ask-examples");
|
||||
if (process.argv.includes("--no-files")) {
|
||||
program.dataSources = [];
|
||||
} else if (process.argv.includes("--example-file")) {
|
||||
@@ -204,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
|
||||
@@ -341,6 +387,7 @@ Please check ${cyan(
|
||||
console.log(`Running app in ${root}...`);
|
||||
await runApp(
|
||||
root,
|
||||
program.template,
|
||||
program.frontend,
|
||||
program.framework,
|
||||
program.port,
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "create-llama",
|
||||
"version": "0.1.34",
|
||||
"version": "0.2.8",
|
||||
"description": "Create LlamaIndex-powered apps with one command",
|
||||
"keywords": [
|
||||
"rag",
|
||||
|
||||
+49
-36
@@ -28,6 +28,7 @@ export type QuestionArgs = Omit<
|
||||
"appPath" | "packageManager"
|
||||
> & {
|
||||
askModels?: boolean;
|
||||
askExamples?: boolean;
|
||||
};
|
||||
const supportedContextFileTypes = [
|
||||
".pdf",
|
||||
@@ -172,7 +173,7 @@ export const getDataSourceChoices = (
|
||||
);
|
||||
}
|
||||
|
||||
if (framework === "fastapi") {
|
||||
if (framework === "fastapi" && template !== "extractor") {
|
||||
choices.push({
|
||||
title: "Use website content (requires Chrome)",
|
||||
value: "web",
|
||||
@@ -183,7 +184,7 @@ export const getDataSourceChoices = (
|
||||
});
|
||||
}
|
||||
|
||||
if (!selectedDataSource.length) {
|
||||
if (!selectedDataSource.length && template !== "extractor") {
|
||||
choices.push({
|
||||
title: "Use managed index from LlamaCloud",
|
||||
value: "llamacloud",
|
||||
@@ -286,27 +287,25 @@ export const askQuestions = async (
|
||||
},
|
||||
];
|
||||
|
||||
if (program.template !== "multiagent") {
|
||||
const modelConfigured =
|
||||
!program.llamapack && program.modelConfig.isConfigured();
|
||||
// If using LlamaParse, require LlamaCloud API key
|
||||
const llamaCloudKeyConfigured = program.useLlamaParse
|
||||
? program.llamaCloudKey || process.env["LLAMA_CLOUD_API_KEY"]
|
||||
: true;
|
||||
const hasVectorDb = program.vectorDb && program.vectorDb !== "none";
|
||||
// Can run the app if all tools do not require configuration
|
||||
if (
|
||||
!hasVectorDb &&
|
||||
modelConfigured &&
|
||||
llamaCloudKeyConfigured &&
|
||||
!toolsRequireConfig(program.tools)
|
||||
) {
|
||||
actionChoices.push({
|
||||
title:
|
||||
"Generate code, install dependencies, and run the app (~2 min)",
|
||||
value: "runApp",
|
||||
});
|
||||
}
|
||||
const modelConfigured =
|
||||
!program.llamapack && program.modelConfig.isConfigured();
|
||||
// If using LlamaParse, require LlamaCloud API key
|
||||
const llamaCloudKeyConfigured = program.useLlamaParse
|
||||
? program.llamaCloudKey || process.env["LLAMA_CLOUD_API_KEY"]
|
||||
: true;
|
||||
const hasVectorDb = program.vectorDb && program.vectorDb !== "none";
|
||||
// Can run the app if all tools do not require configuration
|
||||
if (
|
||||
!hasVectorDb &&
|
||||
modelConfigured &&
|
||||
llamaCloudKeyConfigured &&
|
||||
!toolsRequireConfig(program.tools)
|
||||
) {
|
||||
actionChoices.push({
|
||||
title:
|
||||
"Generate code, install dependencies, and run the app (~2 min)",
|
||||
value: "runApp",
|
||||
});
|
||||
}
|
||||
|
||||
const { action } = await prompts(
|
||||
@@ -338,20 +337,24 @@ export const askQuestions = async (
|
||||
name: "template",
|
||||
message: "Which template would you like to use?",
|
||||
choices: [
|
||||
{ title: "Agentic RAG (single agent)", value: "streaming" },
|
||||
{ title: "Agentic RAG (e.g. chat with docs)", value: "streaming" },
|
||||
{
|
||||
title: "Multi-agent app (using llama-agents)",
|
||||
title: "Multi-agent app (using workflows)",
|
||||
value: "multiagent",
|
||||
},
|
||||
{ title: "Structured Extractor", value: "extractor" },
|
||||
{
|
||||
title: `Community template from ${styledRepo}`,
|
||||
value: "community",
|
||||
},
|
||||
{
|
||||
title: "Example using a LlamaPack",
|
||||
value: "llamapack",
|
||||
},
|
||||
...(program.askExamples
|
||||
? [
|
||||
{
|
||||
title: `Community template from ${styledRepo}`,
|
||||
value: "community",
|
||||
},
|
||||
{
|
||||
title: "Example using a LlamaPack",
|
||||
value: "llamapack",
|
||||
},
|
||||
]
|
||||
: []),
|
||||
],
|
||||
initial: 0,
|
||||
},
|
||||
@@ -407,9 +410,14 @@ export const askQuestions = async (
|
||||
return; // early return - no further questions needed for llamapack projects
|
||||
}
|
||||
|
||||
if (program.template === "multiagent" || program.template === "extractor") {
|
||||
if (program.template === "multiagent") {
|
||||
// TODO: multi-agents currently only supports FastAPI
|
||||
program.framework = preferences.framework = "fastapi";
|
||||
} else if (program.template === "extractor") {
|
||||
// Extractor 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 = preferences.framework = "fastapi";
|
||||
}
|
||||
if (!program.framework) {
|
||||
if (ciInfo.isCI) {
|
||||
@@ -438,7 +446,7 @@ export const askQuestions = async (
|
||||
|
||||
if (
|
||||
(program.framework === "express" || program.framework === "fastapi") &&
|
||||
program.template === "streaming"
|
||||
(program.template === "streaming" || program.template === "multiagent")
|
||||
) {
|
||||
// if a backend-only framework is selected, ask whether we should create a frontend
|
||||
if (program.frontend === undefined) {
|
||||
@@ -629,6 +637,7 @@ export const askQuestions = async (
|
||||
type: "db",
|
||||
config: await prompts(dbPrompts, questionHandlers),
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "llamacloud": {
|
||||
program.dataSources.push({
|
||||
@@ -652,7 +661,11 @@ export const askQuestions = async (
|
||||
// default to use LlamaParse if using LlamaCloud
|
||||
program.useLlamaParse = preferences.useLlamaParse = true;
|
||||
} else {
|
||||
if (program.useLlamaParse === undefined) {
|
||||
// Extractor template doesn't support LlamaParse and LlamaCloud right now (cannot use asyncio loop in Reflex)
|
||||
if (
|
||||
program.useLlamaParse === undefined &&
|
||||
program.template !== "extractor"
|
||||
) {
|
||||
// if already set useLlamaParse, don't ask again
|
||||
if (program.dataSources.some((ds) => ds.type === "file")) {
|
||||
if (ciInfo.isCI) {
|
||||
|
||||
+13
-8
@@ -1,21 +1,25 @@
|
||||
import os
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.core.agent import AgentRunner
|
||||
from llama_index.core.tools.query_engine import QueryEngineTool
|
||||
|
||||
from app.engine.index import IndexConfig, get_index
|
||||
from app.engine.tools import ToolFactory
|
||||
from app.engine.index import get_index
|
||||
from llama_index.core.agent import AgentRunner
|
||||
from llama_index.core.callbacks import CallbackManager
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.core.tools.query_engine import QueryEngineTool
|
||||
|
||||
|
||||
def get_chat_engine(filters=None, params=None):
|
||||
def get_chat_engine(filters=None, params=None, event_handlers=None):
|
||||
system_prompt = os.getenv("SYSTEM_PROMPT")
|
||||
top_k = os.getenv("TOP_K", "3")
|
||||
top_k = int(os.getenv("TOP_K", 0))
|
||||
tools = []
|
||||
callback_manager = CallbackManager(handlers=event_handlers or [])
|
||||
|
||||
# Add query tool if index exists
|
||||
index = get_index()
|
||||
index_config = IndexConfig(callback_manager=callback_manager, **(params or {}))
|
||||
index = get_index(index_config)
|
||||
if index is not None:
|
||||
query_engine = index.as_query_engine(
|
||||
similarity_top_k=int(top_k), filters=filters
|
||||
filters=filters, **({"similarity_top_k": top_k} if top_k != 0 else {})
|
||||
)
|
||||
query_engine_tool = QueryEngineTool.from_defaults(query_engine=query_engine)
|
||||
tools.append(query_engine_tool)
|
||||
@@ -27,5 +31,6 @@ def get_chat_engine(filters=None, params=None):
|
||||
llm=Settings.llm,
|
||||
tools=tools,
|
||||
system_prompt=system_prompt,
|
||||
callback_manager=callback_manager,
|
||||
verbose=True,
|
||||
)
|
||||
@@ -1,8 +1,6 @@
|
||||
import os
|
||||
import yaml
|
||||
import json
|
||||
import importlib
|
||||
from cachetools import cached, LRUCache
|
||||
from llama_index.core.tools.tool_spec.base import BaseToolSpec
|
||||
from llama_index.core.tools.function_tool import FunctionTool
|
||||
|
||||
@@ -13,7 +11,6 @@ class ToolType:
|
||||
|
||||
|
||||
class ToolFactory:
|
||||
|
||||
TOOL_SOURCE_PACKAGE_MAP = {
|
||||
ToolType.LLAMAHUB: "llama_index.tools",
|
||||
ToolType.LOCAL: "app.engine.tools",
|
||||
|
||||
@@ -3,7 +3,7 @@ import logging
|
||||
import base64
|
||||
import uuid
|
||||
from pydantic import BaseModel
|
||||
from typing import List, Tuple, Dict, Optional
|
||||
from typing import List, Dict, Optional
|
||||
from llama_index.core.tools import FunctionTool
|
||||
from e2b_code_interpreter import CodeInterpreter
|
||||
from e2b_code_interpreter.models import Logs
|
||||
@@ -26,7 +26,6 @@ class E2BToolOutput(BaseModel):
|
||||
|
||||
|
||||
class E2BCodeInterpreter:
|
||||
|
||||
output_dir = "output/tool"
|
||||
|
||||
def __init__(self, api_key: str = None):
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
import os
|
||||
from app.engine.index import get_index
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
def get_chat_engine(filters=None, params=None):
|
||||
system_prompt = os.getenv("SYSTEM_PROMPT")
|
||||
top_k = os.getenv("TOP_K", 3)
|
||||
|
||||
index = get_index(params)
|
||||
if index is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=str(
|
||||
"StorageContext is empty - call 'poetry run generate' to generate the storage first"
|
||||
),
|
||||
)
|
||||
|
||||
return index.as_chat_engine(
|
||||
similarity_top_k=int(top_k),
|
||||
system_prompt=system_prompt,
|
||||
chat_mode="condense_plus_context",
|
||||
filters=filters,
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
|
||||
from app.engine.index import IndexConfig, get_index
|
||||
from app.engine.node_postprocessors import NodeCitationProcessor
|
||||
from fastapi import HTTPException
|
||||
from llama_index.core.callbacks import CallbackManager
|
||||
from llama_index.core.chat_engine import CondensePlusContextChatEngine
|
||||
from llama_index.core.memory import ChatMemoryBuffer
|
||||
from llama_index.core.settings import Settings
|
||||
|
||||
|
||||
def get_chat_engine(filters=None, params=None, event_handlers=None):
|
||||
system_prompt = os.getenv("SYSTEM_PROMPT")
|
||||
citation_prompt = os.getenv("SYSTEM_CITATION_PROMPT", None)
|
||||
top_k = int(os.getenv("TOP_K", 0))
|
||||
llm = Settings.llm
|
||||
memory = ChatMemoryBuffer.from_defaults(
|
||||
token_limit=llm.metadata.context_window - 256
|
||||
)
|
||||
callback_manager = CallbackManager(handlers=event_handlers or [])
|
||||
|
||||
node_postprocessors = []
|
||||
if citation_prompt:
|
||||
node_postprocessors = [NodeCitationProcessor()]
|
||||
system_prompt = f"{system_prompt}\n{citation_prompt}"
|
||||
|
||||
index_config = IndexConfig(callback_manager=callback_manager, **(params or {}))
|
||||
index = get_index(index_config)
|
||||
if index is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=str(
|
||||
"StorageContext is empty - call 'poetry run generate' to generate the storage first"
|
||||
),
|
||||
)
|
||||
|
||||
retriever = index.as_retriever(
|
||||
filters=filters, **({"similarity_top_k": top_k} if top_k != 0 else {})
|
||||
)
|
||||
|
||||
return CondensePlusContextChatEngine(
|
||||
llm=llm,
|
||||
memory=memory,
|
||||
system_prompt=system_prompt,
|
||||
retriever=retriever,
|
||||
node_postprocessors=node_postprocessors,
|
||||
callback_manager=callback_manager,
|
||||
)
|
||||
@@ -0,0 +1,21 @@
|
||||
from typing import List, Optional
|
||||
|
||||
from llama_index.core import QueryBundle
|
||||
from llama_index.core.postprocessor.types import BaseNodePostprocessor
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
|
||||
|
||||
class NodeCitationProcessor(BaseNodePostprocessor):
|
||||
"""
|
||||
Append node_id into metadata for citation purpose.
|
||||
Config SYSTEM_CITATION_PROMPT in your runtime environment variable to enable this feature.
|
||||
"""
|
||||
|
||||
def _postprocess_nodes(
|
||||
self,
|
||||
nodes: List[NodeWithScore],
|
||||
query_bundle: Optional[QueryBundle] = None,
|
||||
) -> List[NodeWithScore]:
|
||||
for node_score in nodes:
|
||||
node_score.node.metadata["node_id"] = node_score.node.node_id
|
||||
return nodes
|
||||
@@ -1,5 +1,6 @@
|
||||
import { ContextChatEngine, Settings } from "llamaindex";
|
||||
import { getDataSource } from "./index";
|
||||
import { nodeCitationProcessor } from "./nodePostprocessors";
|
||||
import { generateFilters } from "./queryFilter";
|
||||
|
||||
export async function createChatEngine(documentIds?: string[], params?: any) {
|
||||
@@ -10,13 +11,22 @@ export async function createChatEngine(documentIds?: string[], params?: any) {
|
||||
);
|
||||
}
|
||||
const retriever = index.asRetriever({
|
||||
similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : 3,
|
||||
similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : undefined,
|
||||
filters: generateFilters(documentIds || []),
|
||||
});
|
||||
|
||||
const systemPrompt = process.env.SYSTEM_PROMPT;
|
||||
const citationPrompt = process.env.SYSTEM_CITATION_PROMPT;
|
||||
const prompt =
|
||||
[systemPrompt, citationPrompt].filter((p) => p).join("\n") || undefined;
|
||||
const nodePostprocessors = citationPrompt
|
||||
? [nodeCitationProcessor]
|
||||
: undefined;
|
||||
|
||||
return new ContextChatEngine({
|
||||
chatModel: Settings.llm,
|
||||
retriever,
|
||||
systemPrompt: process.env.SYSTEM_PROMPT,
|
||||
systemPrompt: prompt,
|
||||
nodePostprocessors,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
import {
|
||||
BaseNodePostprocessor,
|
||||
MessageContent,
|
||||
NodeWithScore,
|
||||
} from "llamaindex";
|
||||
|
||||
class NodeCitationProcessor implements BaseNodePostprocessor {
|
||||
/**
|
||||
* Append node_id into metadata for citation purpose.
|
||||
* Config SYSTEM_CITATION_PROMPT in your runtime environment variable to enable this feature.
|
||||
*/
|
||||
async postprocessNodes(
|
||||
nodes: NodeWithScore[],
|
||||
query?: MessageContent,
|
||||
): Promise<NodeWithScore[]> {
|
||||
for (const nodeScore of nodes) {
|
||||
if (!nodeScore.node || !nodeScore.node.metadata) {
|
||||
continue; // Skip nodes with missing properties
|
||||
}
|
||||
nodeScore.node.metadata["node_id"] = nodeScore.node.id_;
|
||||
}
|
||||
return nodes;
|
||||
}
|
||||
}
|
||||
|
||||
export const nodeCitationProcessor = new NodeCitationProcessor();
|
||||
@@ -1,5 +1,4 @@
|
||||
import fs from "fs";
|
||||
import crypto from "node:crypto";
|
||||
import { getExtractors } from "../../engine/loader";
|
||||
|
||||
const MIME_TYPE_TO_EXT: Record<string, string> = {
|
||||
@@ -11,7 +10,24 @@ const MIME_TYPE_TO_EXT: Record<string, string> = {
|
||||
|
||||
const UPLOADED_FOLDER = "output/uploaded";
|
||||
|
||||
export async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
|
||||
export async function storeAndParseFile(
|
||||
filename: string,
|
||||
fileBuffer: Buffer,
|
||||
mimeType: string,
|
||||
) {
|
||||
const documents = await loadDocuments(fileBuffer, mimeType);
|
||||
await saveDocument(filename, fileBuffer, mimeType);
|
||||
for (const document of documents) {
|
||||
document.metadata = {
|
||||
...document.metadata,
|
||||
file_name: filename,
|
||||
private: "true", // to separate private uploads from public documents
|
||||
};
|
||||
}
|
||||
return documents;
|
||||
}
|
||||
|
||||
async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
|
||||
const extractors = getExtractors();
|
||||
const reader = extractors[MIME_TYPE_TO_EXT[mimeType]];
|
||||
|
||||
@@ -22,11 +38,14 @@ export async function loadDocuments(fileBuffer: Buffer, mimeType: string) {
|
||||
return await reader.loadDataAsContent(fileBuffer);
|
||||
}
|
||||
|
||||
export async function saveDocument(fileBuffer: Buffer, mimeType: string) {
|
||||
async function saveDocument(
|
||||
filename: string,
|
||||
fileBuffer: Buffer,
|
||||
mimeType: string,
|
||||
) {
|
||||
const fileExt = MIME_TYPE_TO_EXT[mimeType];
|
||||
if (!fileExt) throw new Error(`Unsupported document type: ${mimeType}`);
|
||||
|
||||
const filename = `${crypto.randomUUID()}.${fileExt}`;
|
||||
const filepath = `${UPLOADED_FOLDER}/${filename}`;
|
||||
const fileurl = `${process.env.FILESERVER_URL_PREFIX}/${filepath}`;
|
||||
|
||||
|
||||
@@ -5,34 +5,24 @@ import {
|
||||
SimpleNodeParser,
|
||||
VectorStoreIndex,
|
||||
} from "llamaindex";
|
||||
import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
|
||||
|
||||
export async function runPipeline(
|
||||
currentIndex: VectorStoreIndex | LlamaCloudIndex,
|
||||
currentIndex: VectorStoreIndex,
|
||||
documents: Document[],
|
||||
) {
|
||||
if (currentIndex instanceof LlamaCloudIndex) {
|
||||
// LlamaCloudIndex processes the documents automatically
|
||||
// so we don't need ingestion pipeline, just insert the documents directly
|
||||
for (const document of documents) {
|
||||
await currentIndex.insert(document);
|
||||
}
|
||||
} else {
|
||||
// Use ingestion pipeline to process the documents into nodes and add them to the vector store
|
||||
const pipeline = new IngestionPipeline({
|
||||
transformations: [
|
||||
new SimpleNodeParser({
|
||||
chunkSize: Settings.chunkSize,
|
||||
chunkOverlap: Settings.chunkOverlap,
|
||||
}),
|
||||
Settings.embedModel,
|
||||
],
|
||||
});
|
||||
const nodes = await pipeline.run({ documents });
|
||||
await currentIndex.insertNodes(nodes);
|
||||
currentIndex.storageContext.docStore.persist();
|
||||
console.log("Added nodes to the vector store.");
|
||||
}
|
||||
|
||||
// Use ingestion pipeline to process the documents into nodes and add them to the vector store
|
||||
const pipeline = new IngestionPipeline({
|
||||
transformations: [
|
||||
new SimpleNodeParser({
|
||||
chunkSize: Settings.chunkSize,
|
||||
chunkOverlap: Settings.chunkOverlap,
|
||||
}),
|
||||
Settings.embedModel,
|
||||
],
|
||||
});
|
||||
const nodes = await pipeline.run({ documents });
|
||||
await currentIndex.insertNodes(nodes);
|
||||
currentIndex.storageContext.docStore.persist();
|
||||
console.log("Added nodes to the vector store.");
|
||||
return documents.map((document) => document.id_);
|
||||
}
|
||||
|
||||
@@ -1,26 +1,32 @@
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { LLamaCloudFileService, VectorStoreIndex } from "llamaindex";
|
||||
import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
|
||||
import { loadDocuments, saveDocument } from "./helper";
|
||||
import { storeAndParseFile } from "./helper";
|
||||
import { runPipeline } from "./pipeline";
|
||||
|
||||
export async function uploadDocument(
|
||||
index: VectorStoreIndex | LlamaCloudIndex,
|
||||
filename: string,
|
||||
raw: string,
|
||||
): Promise<string[]> {
|
||||
const [header, content] = raw.split(",");
|
||||
const mimeType = header.replace("data:", "").replace(";base64", "");
|
||||
const fileBuffer = Buffer.from(content, "base64");
|
||||
const documents = await loadDocuments(fileBuffer, mimeType);
|
||||
const { filename } = await saveDocument(fileBuffer, mimeType);
|
||||
|
||||
// Update documents with metadata
|
||||
for (const document of documents) {
|
||||
document.metadata = {
|
||||
...document.metadata,
|
||||
file_name: filename,
|
||||
private: "true", // to separate private uploads from public documents
|
||||
};
|
||||
if (index instanceof LlamaCloudIndex) {
|
||||
// trigger LlamaCloudIndex API to upload the file and run the pipeline
|
||||
const projectId = await index.getProjectId();
|
||||
const pipelineId = await index.getPipelineId();
|
||||
return [
|
||||
await LLamaCloudFileService.addFileToPipeline(
|
||||
projectId,
|
||||
pipelineId,
|
||||
new File([fileBuffer], filename, { type: mimeType }),
|
||||
{ private: "true" },
|
||||
),
|
||||
];
|
||||
}
|
||||
|
||||
return await runPipeline(index, documents);
|
||||
// run the pipeline for other vector store indexes
|
||||
const documents = await storeAndParseFile(filename, fileBuffer, mimeType);
|
||||
return runPipeline(index, documents);
|
||||
}
|
||||
|
||||
@@ -1,13 +1,18 @@
|
||||
import { StreamData } from "ai";
|
||||
import {
|
||||
CallbackManager,
|
||||
LLamaCloudFileService,
|
||||
Metadata,
|
||||
MetadataMode,
|
||||
NodeWithScore,
|
||||
ToolCall,
|
||||
ToolOutput,
|
||||
} from "llamaindex";
|
||||
import { LLamaCloudFileService } from "./service";
|
||||
import path from "node:path";
|
||||
import { DATA_DIR } from "../../engine/loader";
|
||||
import { downloadFile } from "./file";
|
||||
|
||||
const LLAMA_CLOUD_DOWNLOAD_FOLDER = "output/llamacloud";
|
||||
|
||||
export function appendSourceData(
|
||||
data: StreamData,
|
||||
@@ -84,7 +89,7 @@ export function createCallbackManager(stream: StreamData) {
|
||||
stream,
|
||||
`Retrieved ${nodes.length} sources to use as context for the query`,
|
||||
);
|
||||
LLamaCloudFileService.downloadFiles(nodes); // don't await to avoid blocking chat streaming
|
||||
downloadFilesFromNodes(nodes); // don't await to avoid blocking chat streaming
|
||||
});
|
||||
|
||||
callbackManager.on("llm-tool-call", (event) => {
|
||||
@@ -116,15 +121,71 @@ function getNodeUrl(metadata: Metadata) {
|
||||
if (fileName && process.env.FILESERVER_URL_PREFIX) {
|
||||
// file_name exists and file server is configured
|
||||
const pipelineId = metadata["pipeline_id"];
|
||||
if (pipelineId && metadata["private"] == null) {
|
||||
// file is from LlamaCloud and was not ingested locally
|
||||
const name = LLamaCloudFileService.toDownloadedName(pipelineId, fileName);
|
||||
return `${process.env.FILESERVER_URL_PREFIX}/output/llamacloud/${name}`;
|
||||
if (pipelineId) {
|
||||
const name = toDownloadedName(pipelineId, fileName);
|
||||
return `${process.env.FILESERVER_URL_PREFIX}/${LLAMA_CLOUD_DOWNLOAD_FOLDER}/${name}`;
|
||||
}
|
||||
const isPrivate = metadata["private"] === "true";
|
||||
const folder = isPrivate ? "output/uploaded" : "data";
|
||||
return `${process.env.FILESERVER_URL_PREFIX}/${folder}/${fileName}`;
|
||||
if (isPrivate) {
|
||||
return `${process.env.FILESERVER_URL_PREFIX}/output/uploaded/${fileName}`;
|
||||
}
|
||||
const filePath = metadata["file_path"];
|
||||
const dataDir = path.resolve(DATA_DIR);
|
||||
|
||||
if (filePath && dataDir) {
|
||||
const relativePath = path.relative(dataDir, filePath);
|
||||
return `${process.env.FILESERVER_URL_PREFIX}/data/${relativePath}`;
|
||||
}
|
||||
}
|
||||
// fallback to URL in metadata (e.g. for websites)
|
||||
return metadata["URL"];
|
||||
}
|
||||
|
||||
async function downloadFilesFromNodes(nodes: NodeWithScore<Metadata>[]) {
|
||||
try {
|
||||
const files = nodesToLlamaCloudFiles(nodes);
|
||||
for (const { pipelineId, fileName, downloadedName } of files) {
|
||||
const downloadUrl = await LLamaCloudFileService.getFileUrl(
|
||||
pipelineId,
|
||||
fileName,
|
||||
);
|
||||
if (downloadUrl) {
|
||||
await downloadFile(
|
||||
downloadUrl,
|
||||
downloadedName,
|
||||
LLAMA_CLOUD_DOWNLOAD_FOLDER,
|
||||
);
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error downloading files from nodes:", error);
|
||||
}
|
||||
}
|
||||
|
||||
function nodesToLlamaCloudFiles(nodes: NodeWithScore<Metadata>[]) {
|
||||
const files: Array<{
|
||||
pipelineId: string;
|
||||
fileName: string;
|
||||
downloadedName: string;
|
||||
}> = [];
|
||||
for (const node of nodes) {
|
||||
const pipelineId = node.node.metadata["pipeline_id"];
|
||||
const fileName = node.node.metadata["file_name"];
|
||||
if (!pipelineId || !fileName) continue;
|
||||
const isDuplicate = files.some(
|
||||
(f) => f.pipelineId === pipelineId && f.fileName === fileName,
|
||||
);
|
||||
if (!isDuplicate) {
|
||||
files.push({
|
||||
pipelineId,
|
||||
fileName,
|
||||
downloadedName: toDownloadedName(pipelineId, fileName),
|
||||
});
|
||||
}
|
||||
}
|
||||
return files;
|
||||
}
|
||||
|
||||
function toDownloadedName(pipelineId: string, fileName: string) {
|
||||
return `${pipelineId}$${fileName}`;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
import fs from "node:fs";
|
||||
import https from "node:https";
|
||||
import path from "node:path";
|
||||
|
||||
export async function downloadFile(
|
||||
urlToDownload: string,
|
||||
filename: string,
|
||||
folder = "output/uploaded",
|
||||
) {
|
||||
try {
|
||||
const downloadedPath = path.join(folder, filename);
|
||||
|
||||
// Check if file already exists
|
||||
if (fs.existsSync(downloadedPath)) return;
|
||||
|
||||
const file = fs.createWriteStream(downloadedPath);
|
||||
https
|
||||
.get(urlToDownload, (response) => {
|
||||
response.pipe(file);
|
||||
file.on("finish", () => {
|
||||
file.close(() => {
|
||||
console.log("File downloaded successfully");
|
||||
});
|
||||
});
|
||||
})
|
||||
.on("error", (err) => {
|
||||
fs.unlink(downloadedPath, () => {
|
||||
console.error("Error downloading file:", err);
|
||||
throw err;
|
||||
});
|
||||
});
|
||||
} catch (error) {
|
||||
throw new Error(`Error downloading file: ${error}`);
|
||||
}
|
||||
}
|
||||
@@ -1,187 +0,0 @@
|
||||
import { Metadata, NodeWithScore } from "llamaindex";
|
||||
import fs from "node:fs";
|
||||
import https from "node:https";
|
||||
import path from "node:path";
|
||||
|
||||
const LLAMA_CLOUD_OUTPUT_DIR = "output/llamacloud";
|
||||
const LLAMA_CLOUD_BASE_URL = "https://cloud.llamaindex.ai/api/v1";
|
||||
const FILE_DELIMITER = "$"; // delimiter between pipelineId and filename
|
||||
|
||||
type LlamaCloudFile = {
|
||||
name: string;
|
||||
file_id: string;
|
||||
project_id: string;
|
||||
};
|
||||
|
||||
type LLamaCloudProject = {
|
||||
id: string;
|
||||
organization_id: string;
|
||||
name: string;
|
||||
is_default: boolean;
|
||||
};
|
||||
|
||||
type LLamaCloudPipeline = {
|
||||
id: string;
|
||||
name: string;
|
||||
project_id: string;
|
||||
};
|
||||
|
||||
export class LLamaCloudFileService {
|
||||
private static readonly headers = {
|
||||
Accept: "application/json",
|
||||
Authorization: `Bearer ${process.env.LLAMA_CLOUD_API_KEY}`,
|
||||
};
|
||||
|
||||
public static async getAllProjectsWithPipelines() {
|
||||
try {
|
||||
const projects = await LLamaCloudFileService.getAllProjects();
|
||||
const pipelines = await LLamaCloudFileService.getAllPipelines();
|
||||
return projects.map((project) => ({
|
||||
...project,
|
||||
pipelines: pipelines.filter((p) => p.project_id === project.id),
|
||||
}));
|
||||
} catch (error) {
|
||||
console.error("Error listing projects and pipelines:", error);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
public static async downloadFiles(nodes: NodeWithScore<Metadata>[]) {
|
||||
const files = LLamaCloudFileService.nodesToDownloadFiles(nodes);
|
||||
if (!files.length) return;
|
||||
console.log("Downloading files from LlamaCloud...");
|
||||
for (const file of files) {
|
||||
await LLamaCloudFileService.downloadFile(file.pipelineId, file.fileName);
|
||||
}
|
||||
}
|
||||
|
||||
public static toDownloadedName(pipelineId: string, fileName: string) {
|
||||
return `${pipelineId}${FILE_DELIMITER}${fileName}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* This function will return an array of unique files to download from LlamaCloud
|
||||
* We only download files that are uploaded directly in LlamaCloud datasources (don't have `private` in metadata)
|
||||
* Files are uploaded directly in LlamaCloud datasources don't have `private` in metadata (public docs)
|
||||
* Files are uploaded from local via `generate` command will have `private=false` (public docs)
|
||||
* Files are uploaded from local via `/chat/upload` endpoint will have `private=true` (private docs)
|
||||
*
|
||||
* @param nodes
|
||||
* @returns list of unique files to download
|
||||
*/
|
||||
private static nodesToDownloadFiles(nodes: NodeWithScore<Metadata>[]) {
|
||||
const downloadFiles: Array<{
|
||||
pipelineId: string;
|
||||
fileName: string;
|
||||
}> = [];
|
||||
for (const node of nodes) {
|
||||
const isLocalFile = node.node.metadata["private"] != null;
|
||||
const pipelineId = node.node.metadata["pipeline_id"];
|
||||
const fileName = node.node.metadata["file_name"];
|
||||
if (isLocalFile || !pipelineId || !fileName) continue;
|
||||
const isDuplicate = downloadFiles.some(
|
||||
(f) => f.pipelineId === pipelineId && f.fileName === fileName,
|
||||
);
|
||||
if (!isDuplicate) {
|
||||
downloadFiles.push({ pipelineId, fileName });
|
||||
}
|
||||
}
|
||||
return downloadFiles;
|
||||
}
|
||||
|
||||
private static async downloadFile(pipelineId: string, fileName: string) {
|
||||
try {
|
||||
const downloadedName = LLamaCloudFileService.toDownloadedName(
|
||||
pipelineId,
|
||||
fileName,
|
||||
);
|
||||
const downloadedPath = path.join(LLAMA_CLOUD_OUTPUT_DIR, downloadedName);
|
||||
|
||||
// Check if file already exists
|
||||
if (fs.existsSync(downloadedPath)) return;
|
||||
|
||||
const urlToDownload = await LLamaCloudFileService.getFileUrlByName(
|
||||
pipelineId,
|
||||
fileName,
|
||||
);
|
||||
if (!urlToDownload) throw new Error("File not found in LlamaCloud");
|
||||
|
||||
const file = fs.createWriteStream(downloadedPath);
|
||||
https
|
||||
.get(urlToDownload, (response) => {
|
||||
response.pipe(file);
|
||||
file.on("finish", () => {
|
||||
file.close(() => {
|
||||
console.log("File downloaded successfully");
|
||||
});
|
||||
});
|
||||
})
|
||||
.on("error", (err) => {
|
||||
fs.unlink(downloadedPath, () => {
|
||||
console.error("Error downloading file:", err);
|
||||
throw err;
|
||||
});
|
||||
});
|
||||
} catch (error) {
|
||||
throw new Error(`Error downloading file from LlamaCloud: ${error}`);
|
||||
}
|
||||
}
|
||||
|
||||
private static async getFileUrlByName(
|
||||
pipelineId: string,
|
||||
name: string,
|
||||
): Promise<string | null> {
|
||||
const files = await LLamaCloudFileService.getAllFiles(pipelineId);
|
||||
const file = files.find((file) => file.name === name);
|
||||
if (!file) return null;
|
||||
return await LLamaCloudFileService.getFileUrlById(
|
||||
file.project_id,
|
||||
file.file_id,
|
||||
);
|
||||
}
|
||||
|
||||
private static async getFileUrlById(
|
||||
projectId: string,
|
||||
fileId: string,
|
||||
): Promise<string> {
|
||||
const url = `${LLAMA_CLOUD_BASE_URL}/files/${fileId}/content?project_id=${projectId}`;
|
||||
const response = await fetch(url, {
|
||||
method: "GET",
|
||||
headers: LLamaCloudFileService.headers,
|
||||
});
|
||||
const data = (await response.json()) as { url: string };
|
||||
return data.url;
|
||||
}
|
||||
|
||||
private static async getAllFiles(
|
||||
pipelineId: string,
|
||||
): Promise<LlamaCloudFile[]> {
|
||||
const url = `${LLAMA_CLOUD_BASE_URL}/pipelines/${pipelineId}/files`;
|
||||
const response = await fetch(url, {
|
||||
method: "GET",
|
||||
headers: LLamaCloudFileService.headers,
|
||||
});
|
||||
const data = await response.json();
|
||||
return data;
|
||||
}
|
||||
|
||||
private static async getAllProjects(): Promise<LLamaCloudProject[]> {
|
||||
const url = `${LLAMA_CLOUD_BASE_URL}/projects`;
|
||||
const response = await fetch(url, {
|
||||
method: "GET",
|
||||
headers: LLamaCloudFileService.headers,
|
||||
});
|
||||
const data = (await response.json()) as LLamaCloudProject[];
|
||||
return data;
|
||||
}
|
||||
|
||||
private static async getAllPipelines(): Promise<LLamaCloudPipeline[]> {
|
||||
const url = `${LLAMA_CLOUD_BASE_URL}/pipelines`;
|
||||
const response = await fetch(url, {
|
||||
method: "GET",
|
||||
headers: LLamaCloudFileService.headers,
|
||||
});
|
||||
const data = (await response.json()) as LLamaCloudPipeline[];
|
||||
return data;
|
||||
}
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
import {
|
||||
StreamData,
|
||||
createCallbacksTransformer,
|
||||
createStreamDataTransformer,
|
||||
trimStartOfStreamHelper,
|
||||
type AIStreamCallbacksAndOptions,
|
||||
} from "ai";
|
||||
import { ChatMessage, EngineResponse } from "llamaindex";
|
||||
import { generateNextQuestions } from "./suggestion";
|
||||
|
||||
export function LlamaIndexStream(
|
||||
response: AsyncIterable<EngineResponse>,
|
||||
data: StreamData,
|
||||
chatHistory: ChatMessage[],
|
||||
opts?: {
|
||||
callbacks?: AIStreamCallbacksAndOptions;
|
||||
},
|
||||
): ReadableStream<Uint8Array> {
|
||||
return createParser(response, data, chatHistory)
|
||||
.pipeThrough(createCallbacksTransformer(opts?.callbacks))
|
||||
.pipeThrough(createStreamDataTransformer());
|
||||
}
|
||||
|
||||
function createParser(
|
||||
res: AsyncIterable<EngineResponse>,
|
||||
data: StreamData,
|
||||
chatHistory: ChatMessage[],
|
||||
) {
|
||||
const it = res[Symbol.asyncIterator]();
|
||||
const trimStartOfStream = trimStartOfStreamHelper();
|
||||
let llmTextResponse = "";
|
||||
|
||||
return new ReadableStream<string>({
|
||||
async pull(controller): Promise<void> {
|
||||
const { value, done } = await it.next();
|
||||
if (done) {
|
||||
controller.close();
|
||||
// LLM stream is done, generate the next questions with a new LLM call
|
||||
chatHistory.push({ role: "assistant", content: llmTextResponse });
|
||||
const questions: string[] = await generateNextQuestions(chatHistory);
|
||||
if (questions.length > 0) {
|
||||
data.appendMessageAnnotation({
|
||||
type: "suggested_questions",
|
||||
data: questions,
|
||||
});
|
||||
}
|
||||
data.close();
|
||||
return;
|
||||
}
|
||||
const text = trimStartOfStream(value.delta ?? "");
|
||||
if (text) {
|
||||
llmTextResponse += text;
|
||||
controller.enqueue(text);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,32 +1,20 @@
|
||||
import { ChatMessage, Settings } from "llamaindex";
|
||||
|
||||
const NEXT_QUESTION_PROMPT_TEMPLATE = `You're a helpful assistant! Your task is to suggest the next question that user might ask.
|
||||
Here is the conversation history
|
||||
---------------------
|
||||
$conversation
|
||||
---------------------
|
||||
Given the conversation history, please give me $number_of_questions questions that you might ask next!
|
||||
Your answer should be wrapped in three sticks which follows the following format:
|
||||
\`\`\`
|
||||
<question 1>
|
||||
<question 2>\`\`\`
|
||||
`;
|
||||
const N_QUESTIONS_TO_GENERATE = 3;
|
||||
|
||||
export async function generateNextQuestions(
|
||||
conversation: ChatMessage[],
|
||||
numberOfQuestions: number = N_QUESTIONS_TO_GENERATE,
|
||||
) {
|
||||
export async function generateNextQuestions(conversation: ChatMessage[]) {
|
||||
const llm = Settings.llm;
|
||||
const NEXT_QUESTION_PROMPT = process.env.NEXT_QUESTION_PROMPT;
|
||||
if (!NEXT_QUESTION_PROMPT) {
|
||||
return [];
|
||||
}
|
||||
|
||||
// Format conversation
|
||||
const conversationText = conversation
|
||||
.map((message) => `${message.role}: ${message.content}`)
|
||||
.join("\n");
|
||||
const message = NEXT_QUESTION_PROMPT_TEMPLATE.replace(
|
||||
"$conversation",
|
||||
const message = NEXT_QUESTION_PROMPT.replace(
|
||||
"{conversation}",
|
||||
conversationText,
|
||||
).replace("$number_of_questions", numberOfQuestions.toString());
|
||||
);
|
||||
|
||||
try {
|
||||
const response = await llm.complete({ prompt: message });
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
from pydantic import BaseModel, validator
|
||||
from llama_index.core.indices.vector_store import VectorStoreIndex
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -2,21 +2,16 @@ import os
|
||||
import logging
|
||||
from typing import Dict
|
||||
from llama_parse import LlamaParse
|
||||
from pydantic import BaseModel, validator
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.config import DATA_DIR
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileLoaderConfig(BaseModel):
|
||||
data_dir: str = "data"
|
||||
use_llama_parse: bool = False
|
||||
|
||||
@validator("data_dir")
|
||||
def data_dir_must_exist(cls, v):
|
||||
if not os.path.isdir(v):
|
||||
raise ValueError(f"Directory '{v}' does not exist")
|
||||
return v
|
||||
|
||||
|
||||
def llama_parse_parser():
|
||||
if os.getenv("LLAMA_CLOUD_API_KEY") is None:
|
||||
@@ -54,7 +49,7 @@ def get_file_documents(config: FileLoaderConfig):
|
||||
|
||||
file_extractor = llama_parse_extractor()
|
||||
reader = SimpleDirectoryReader(
|
||||
config.data_dir,
|
||||
DATA_DIR,
|
||||
recursive=True,
|
||||
filename_as_id=True,
|
||||
raise_on_error=True,
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
import os
|
||||
import json
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
|
||||
+14
-13
@@ -3,18 +3,14 @@ import mimetypes
|
||||
import os
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
import time
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from uuid import uuid4
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
|
||||
from app.engine.index import get_index
|
||||
from app.engine.index import IndexConfig, get_index
|
||||
from llama_index.core import VectorStoreIndex
|
||||
from llama_index.core.ingestion import IngestionPipeline
|
||||
from llama_index.core.readers.file.base import (
|
||||
_try_loading_included_file_formats as get_file_loaders_map,
|
||||
)
|
||||
from llama_index.core.readers.file.base import default_file_metadata_func
|
||||
from llama_index.core.schema import Document
|
||||
from llama_index.indices.managed.llama_cloud.base import LlamaCloudIndex
|
||||
from llama_index.readers.file import FlatReader
|
||||
@@ -50,12 +46,9 @@ class PrivateFileService:
|
||||
return base64.b64decode(data), extension
|
||||
|
||||
@staticmethod
|
||||
def store_and_parse_file(file_data, extension) -> List[Document]:
|
||||
def store_and_parse_file(file_name, file_data, extension) -> List[Document]:
|
||||
# Store file to the private directory
|
||||
os.makedirs(PrivateFileService.PRIVATE_STORE_PATH, exist_ok=True)
|
||||
|
||||
# random file name
|
||||
file_name = f"{uuid4().hex}{extension}"
|
||||
file_path = Path(os.path.join(PrivateFileService.PRIVATE_STORE_PATH, file_name))
|
||||
|
||||
# write file
|
||||
@@ -79,11 +72,17 @@ class PrivateFileService:
|
||||
return documents
|
||||
|
||||
@staticmethod
|
||||
def process_file(file_name: str, base64_content: str, params: Any) -> List[str]:
|
||||
def process_file(
|
||||
file_name: str, base64_content: str, params: Optional[dict] = None
|
||||
) -> List[str]:
|
||||
if params is None:
|
||||
params = {}
|
||||
|
||||
file_data, extension = PrivateFileService.preprocess_base64_file(base64_content)
|
||||
|
||||
# Add the nodes to the index and persist it
|
||||
current_index = get_index(params)
|
||||
index_config = IndexConfig(**params)
|
||||
current_index = get_index(index_config)
|
||||
|
||||
# Insert the documents into the index
|
||||
if isinstance(current_index, LlamaCloudIndex):
|
||||
@@ -106,7 +105,9 @@ class PrivateFileService:
|
||||
]
|
||||
else:
|
||||
# First process documents into nodes
|
||||
documents = PrivateFileService.store_and_parse_file(file_data, extension)
|
||||
documents = PrivateFileService.store_and_parse_file(
|
||||
file_name, file_data, extension
|
||||
)
|
||||
pipeline = IngestionPipeline()
|
||||
nodes = pipeline.run(documents=documents)
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
from app.api.routers.models import Message
|
||||
from llama_index.core.prompts import PromptTemplate
|
||||
from llama_index.core.settings import Settings
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class NextQuestionSuggestion:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the conversation history
|
||||
Disable this feature by removing the NEXT_QUESTION_PROMPT environment variable
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_configured_prompt(cls) -> Optional[str]:
|
||||
prompt = os.getenv("NEXT_QUESTION_PROMPT", None)
|
||||
if not prompt:
|
||||
return None
|
||||
return PromptTemplate(prompt)
|
||||
|
||||
@classmethod
|
||||
async def suggest_next_questions_all_messages(
|
||||
cls,
|
||||
messages: List[Message],
|
||||
) -> Optional[List[str]]:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the conversation history
|
||||
Return None if suggestion is disabled or there is an error
|
||||
"""
|
||||
prompt_template = cls.get_configured_prompt()
|
||||
if not prompt_template:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Reduce the cost by only using the last two messages
|
||||
last_user_message = None
|
||||
last_assistant_message = None
|
||||
for message in reversed(messages):
|
||||
if message.role == "user":
|
||||
last_user_message = f"User: {message.content}"
|
||||
elif message.role == "assistant":
|
||||
last_assistant_message = f"Assistant: {message.content}"
|
||||
if last_user_message and last_assistant_message:
|
||||
break
|
||||
conversation: str = f"{last_user_message}\n{last_assistant_message}"
|
||||
|
||||
# Call the LLM and parse questions from the output
|
||||
prompt = prompt_template.format(conversation=conversation)
|
||||
output = await Settings.llm.acomplete(prompt)
|
||||
questions = cls._extract_questions(output.text)
|
||||
|
||||
return questions
|
||||
except Exception as e:
|
||||
logger.error(f"Error when generating next question: {e}")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _extract_questions(cls, text: str) -> List[str]:
|
||||
content_match = re.search(r"```(.*?)```", text, re.DOTALL)
|
||||
content = content_match.group(1) if content_match else ""
|
||||
return content.strip().split("\n")
|
||||
|
||||
@classmethod
|
||||
async def suggest_next_questions(
|
||||
cls,
|
||||
chat_history: List[Message],
|
||||
response: str,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the chat history and the last response
|
||||
"""
|
||||
messages = chat_history + [Message(role="assistant", content=response)]
|
||||
return await cls.suggest_next_questions_all_messages(messages)
|
||||
@@ -6,11 +6,13 @@ import os
|
||||
DEFAULT_MODEL = "gpt-3.5-turbo"
|
||||
DEFAULT_EMBEDDING_MODEL = "text-embedding-3-large"
|
||||
|
||||
|
||||
class TSIEmbedding(OpenAIEmbedding):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._query_engine = self._text_engine = self.model_name
|
||||
|
||||
|
||||
def llm_config_from_env() -> Dict:
|
||||
from llama_index.core.constants import DEFAULT_TEMPERATURE
|
||||
|
||||
@@ -32,7 +34,7 @@ def llm_config_from_env() -> Dict:
|
||||
|
||||
def embedding_config_from_env() -> Dict:
|
||||
from llama_index.core.constants import DEFAULT_EMBEDDING_DIM
|
||||
|
||||
|
||||
model = os.getenv("EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL)
|
||||
dimension = os.getenv("EMBEDDING_DIM", DEFAULT_EMBEDDING_DIM)
|
||||
api_key = os.getenv("T_SYSTEMS_LLMHUB_API_KEY")
|
||||
@@ -46,6 +48,7 @@ def embedding_config_from_env() -> Dict:
|
||||
}
|
||||
return config
|
||||
|
||||
|
||||
def init_llmhub():
|
||||
from llama_index.llms.openai_like import OpenAILike
|
||||
|
||||
@@ -58,4 +61,4 @@ def init_llmhub():
|
||||
is_chat_model=True,
|
||||
is_function_calling_model=False,
|
||||
context_window=4096,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -126,13 +126,7 @@ def init_fastembed():
|
||||
def init_groq():
|
||||
from llama_index.llms.groq import Groq
|
||||
|
||||
model_map: Dict[str, str] = {
|
||||
"llama3-8b": "llama3-8b-8192",
|
||||
"llama3-70b": "llama3-70b-8192",
|
||||
"mixtral-8x7b": "mixtral-8x7b-32768",
|
||||
}
|
||||
|
||||
Settings.llm = Groq(model=model_map[os.getenv("MODEL")])
|
||||
Settings.llm = Groq(model=os.getenv("MODEL"))
|
||||
# Groq does not provide embeddings, so we use FastEmbed instead
|
||||
init_fastembed()
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
# flake8: noqa: E402
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from app.engine.index import get_index
|
||||
|
||||
@@ -1,41 +1,87 @@
|
||||
import logging
|
||||
import os
|
||||
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
|
||||
from typing import Optional
|
||||
|
||||
from llama_index.core.callbacks import CallbackManager
|
||||
from llama_index.core.ingestion.api_utils import (
|
||||
get_client as llama_cloud_get_client,
|
||||
)
|
||||
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
|
||||
from pydantic import BaseModel, Field, validator
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
def get_client():
|
||||
return llama_cloud_get_client(
|
||||
os.getenv("LLAMA_CLOUD_API_KEY"),
|
||||
os.getenv("LLAMA_CLOUD_BASE_URL"),
|
||||
class LlamaCloudConfig(BaseModel):
|
||||
# Private attributes
|
||||
api_key: str = Field(
|
||||
default=os.getenv("LLAMA_CLOUD_API_KEY"),
|
||||
exclude=True, # Exclude from the model representation
|
||||
)
|
||||
base_url: Optional[str] = Field(
|
||||
default=os.getenv("LLAMA_CLOUD_BASE_URL"),
|
||||
exclude=True,
|
||||
)
|
||||
organization_id: Optional[str] = Field(
|
||||
default=os.getenv("LLAMA_CLOUD_ORGANIZATION_ID"),
|
||||
exclude=True,
|
||||
)
|
||||
# Configuration attributes, can be set by the user
|
||||
pipeline: str = Field(
|
||||
description="The name of the pipeline to use",
|
||||
default=os.getenv("LLAMA_CLOUD_INDEX_NAME"),
|
||||
)
|
||||
project: str = Field(
|
||||
description="The name of the LlamaCloud project",
|
||||
default=os.getenv("LLAMA_CLOUD_PROJECT_NAME"),
|
||||
)
|
||||
|
||||
# Validate and throw error if the env variables are not set before starting the app
|
||||
@validator("pipeline", "project", "api_key", pre=True, always=True)
|
||||
@classmethod
|
||||
def validate_env_vars(cls, value):
|
||||
if value 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"
|
||||
)
|
||||
return value
|
||||
|
||||
def get_index(params=None):
|
||||
configParams = params or {}
|
||||
pipelineConfig = configParams.get("llamaCloudPipeline", {})
|
||||
name = pipelineConfig.get("pipeline", os.getenv("LLAMA_CLOUD_INDEX_NAME"))
|
||||
project_name = pipelineConfig.get("project", os.getenv("LLAMA_CLOUD_PROJECT_NAME"))
|
||||
api_key = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
organization_id = os.getenv("LLAMA_CLOUD_ORGANIZATION_ID")
|
||||
def to_client_kwargs(self) -> dict:
|
||||
return {
|
||||
"api_key": self.api_key,
|
||||
"base_url": self.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,
|
||||
organization_id=organization_id,
|
||||
class IndexConfig(BaseModel):
|
||||
llama_cloud_pipeline_config: LlamaCloudConfig = Field(
|
||||
default=LlamaCloudConfig(),
|
||||
alias="llamaCloudPipeline",
|
||||
)
|
||||
callback_manager: Optional[CallbackManager] = Field(
|
||||
default=None,
|
||||
)
|
||||
|
||||
def to_index_kwargs(self) -> dict:
|
||||
return {
|
||||
"name": self.llama_cloud_pipeline_config.pipeline,
|
||||
"project_name": self.llama_cloud_pipeline_config.project,
|
||||
"api_key": self.llama_cloud_pipeline_config.api_key,
|
||||
"base_url": self.llama_cloud_pipeline_config.base_url,
|
||||
"organization_id": self.llama_cloud_pipeline_config.organization_id,
|
||||
"callback_manager": self.callback_manager,
|
||||
}
|
||||
|
||||
|
||||
def get_index(config: IndexConfig = None):
|
||||
if config is None:
|
||||
config = IndexConfig()
|
||||
index = LlamaCloudIndex(**config.to_index_kwargs())
|
||||
|
||||
return index
|
||||
|
||||
|
||||
def get_client():
|
||||
config = LlamaCloudConfig()
|
||||
return llama_cloud_get_client(**config.to_client_kwargs())
|
||||
|
||||
@@ -5,11 +5,11 @@ def generate_filters(doc_ids):
|
||||
"""
|
||||
Generate public/private document filters based on the doc_ids and the vector store.
|
||||
"""
|
||||
# Using "nin" filter to include the documents don't have the "private" key because they're uploaded in LlamaCloud UI
|
||||
# Using "is_empty" filter to include the documents don't have the "private" key because they're uploaded in LlamaCloud UI
|
||||
public_doc_filter = MetadataFilter(
|
||||
key="private",
|
||||
value=["true"],
|
||||
operator="nin", # type: ignore
|
||||
value=None,
|
||||
operator="is_empty", # type: ignore
|
||||
)
|
||||
selected_doc_filter = MetadataFilter(
|
||||
key="file_id", # Note: LLamaCloud uses "file_id" to reference private document ids as "doc_id" is a restricted field in LlamaCloud
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
# flake8: noqa: E402
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
@@ -26,6 +27,7 @@ def generate_datasource():
|
||||
doc.metadata["private"] = "false"
|
||||
index = VectorStoreIndex.from_documents(
|
||||
documents,
|
||||
show_progress=True,
|
||||
)
|
||||
# store it for later
|
||||
index.storage_context.persist(storage_dir)
|
||||
|
||||
@@ -1,30 +1,43 @@
|
||||
import os
|
||||
import logging
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from typing import Optional
|
||||
|
||||
from cachetools import cached, TTLCache
|
||||
from llama_index.core.storage import StorageContext
|
||||
from cachetools import TTLCache, cached
|
||||
from llama_index.core.callbacks import CallbackManager
|
||||
from llama_index.core.indices import load_index_from_storage
|
||||
from llama_index.core.storage import StorageContext
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class IndexConfig(BaseModel):
|
||||
callback_manager: Optional[CallbackManager] = Field(
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def get_index(config: IndexConfig = None):
|
||||
if config is None:
|
||||
config = IndexConfig()
|
||||
storage_dir = os.getenv("STORAGE_DIR", "storage")
|
||||
# check if storage already exists
|
||||
if not os.path.exists(storage_dir):
|
||||
return None
|
||||
# load the existing index
|
||||
logger.info(f"Loading index from {storage_dir}...")
|
||||
storage_context = get_storage_context(storage_dir)
|
||||
index = load_index_from_storage(
|
||||
storage_context, callback_manager=config.callback_manager
|
||||
)
|
||||
logger.info(f"Finished loading index from {storage_dir}")
|
||||
return index
|
||||
|
||||
|
||||
@cached(
|
||||
TTLCache(maxsize=10, ttl=timedelta(minutes=5).total_seconds()),
|
||||
key=lambda *args, **kwargs: "global_storage_context",
|
||||
)
|
||||
def get_storage_context(persist_dir: str) -> StorageContext:
|
||||
return StorageContext.from_defaults(persist_dir=persist_dir)
|
||||
|
||||
|
||||
def get_index(params=None):
|
||||
storage_dir = os.getenv("STORAGE_DIR", "storage")
|
||||
# check if storage already exists
|
||||
if not os.path.exists(storage_dir):
|
||||
return None
|
||||
# load the existing index
|
||||
logger.info(f"Loading index from {storage_dir}...")
|
||||
storage_context = get_storage_context(storage_dir)
|
||||
index = load_index_from_storage(storage_context)
|
||||
logger.info(f"Finished loading index from {storage_dir}")
|
||||
return index
|
||||
|
||||
@@ -17,9 +17,11 @@ def _create_weaviate_client():
|
||||
client = weaviate.connect_to_weaviate_cloud(cluster_url, auth_credentials)
|
||||
return client
|
||||
|
||||
|
||||
# Global variable to store the Weaviate client
|
||||
client = None
|
||||
|
||||
|
||||
def get_vector_store():
|
||||
global client
|
||||
if client is None:
|
||||
|
||||
@@ -1,24 +1,44 @@
|
||||
import * as dotenv from "dotenv";
|
||||
import { LlamaCloudIndex } from "llamaindex";
|
||||
import * as fs from "fs/promises";
|
||||
import { LLamaCloudFileService } from "llamaindex";
|
||||
import * as path from "path";
|
||||
import { getDataSource } from "./index";
|
||||
import { getDocuments } from "./loader";
|
||||
import { DATA_DIR } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
|
||||
dotenv.config();
|
||||
|
||||
async function loadAndIndex() {
|
||||
const documents = await getDocuments();
|
||||
async function* walk(dir: string): AsyncGenerator<string> {
|
||||
const directory = await fs.opendir(dir);
|
||||
|
||||
await getDataSource();
|
||||
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!`);
|
||||
for await (const dirent of directory) {
|
||||
const entryPath = path.join(dir, dirent.name);
|
||||
|
||||
if (dirent.isDirectory()) {
|
||||
yield* walk(entryPath); // Recursively walk through directories
|
||||
} else if (dirent.isFile()) {
|
||||
yield entryPath; // Yield file paths
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function loadAndIndex() {
|
||||
const index = await getDataSource();
|
||||
const projectId = await index.getProjectId();
|
||||
const pipelineId = await index.getPipelineId();
|
||||
|
||||
// walk through the data directory and upload each file to LlamaCloud
|
||||
for await (const filePath of walk(DATA_DIR)) {
|
||||
const buffer = await fs.readFile(filePath);
|
||||
const filename = path.basename(filePath);
|
||||
const file = new File([buffer], filename);
|
||||
await LLamaCloudFileService.addFileToPipeline(projectId, pipelineId, file, {
|
||||
private: "false",
|
||||
});
|
||||
}
|
||||
|
||||
console.log(`Successfully uploaded documents to LlamaCloud!`);
|
||||
}
|
||||
|
||||
(async () => {
|
||||
|
||||
@@ -18,6 +18,7 @@ export async function getDataSource(params?: LlamaCloudDataSourceParams) {
|
||||
);
|
||||
}
|
||||
const index = new LlamaCloudIndex({
|
||||
organizationId: process.env.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
name: pipelineName,
|
||||
projectName,
|
||||
apiKey,
|
||||
|
||||
@@ -4,15 +4,15 @@ export function generateFilters(documentIds: string[]): MetadataFilters {
|
||||
// public documents don't have the "private" field or it's set to "false"
|
||||
const publicDocumentsFilter: MetadataFilter = {
|
||||
key: "private",
|
||||
value: ["true"],
|
||||
operator: "nin",
|
||||
value: null,
|
||||
operator: "is_empty",
|
||||
};
|
||||
|
||||
// if no documentIds are provided, only retrieve information from public documents
|
||||
if (!documentIds.length) return { filters: [publicDocumentsFilter] };
|
||||
|
||||
const privateDocumentsFilter: MetadataFilter = {
|
||||
key: "doc_id",
|
||||
key: "file_id", // Note: LLamaCloud uses "file_id" to reference private document ids as "doc_id" is a restricted field in LlamaCloud
|
||||
value: documentIds,
|
||||
operator: "in",
|
||||
};
|
||||
|
||||
@@ -1,14 +1,11 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import * as dotenv from "dotenv";
|
||||
import {
|
||||
MongoDBAtlasVectorSearch,
|
||||
VectorStoreIndex,
|
||||
storageContextFromDefaults,
|
||||
} from "llamaindex";
|
||||
import { storageContextFromDefaults, VectorStoreIndex } from "llamaindex";
|
||||
import { MongoDBAtlasVectorSearch } from "llamaindex/storage/vectorStore/MongoDBAtlasVectorStore";
|
||||
import { MongoClient } from "mongodb";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
import { checkRequiredEnvVars, POPULATED_METADATA_FIELDS } from "./shared";
|
||||
|
||||
dotenv.config();
|
||||
|
||||
@@ -30,6 +27,12 @@ async function loadAndIndex() {
|
||||
dbName: databaseName,
|
||||
collectionName: vectorCollectionName, // this is where your embeddings will be stored
|
||||
indexName: indexName, // this is the name of the index you will need to create
|
||||
indexedMetadataFields: POPULATED_METADATA_FIELDS,
|
||||
embeddingDefinition: {
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: 1536,
|
||||
},
|
||||
});
|
||||
|
||||
// now create an index from all the Documents and store them in Atlas
|
||||
|
||||
@@ -1,16 +1,23 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { MongoDBAtlasVectorSearch, VectorStoreIndex } from "llamaindex";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { MongoDBAtlasVectorSearch } from "llamaindex/storage/vectorStore/MongoDBAtlasVectorStore";
|
||||
import { MongoClient } from "mongodb";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
import { checkRequiredEnvVars, POPULATED_METADATA_FIELDS } from "./shared";
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
checkRequiredEnvVars();
|
||||
const client = new MongoClient(process.env.MONGO_URI!);
|
||||
const client = new MongoClient(process.env.MONGODB_URI!);
|
||||
const store = new MongoDBAtlasVectorSearch({
|
||||
mongodbClient: client,
|
||||
dbName: process.env.MONGODB_DATABASE!,
|
||||
collectionName: process.env.MONGODB_VECTORS!,
|
||||
indexName: process.env.MONGODB_VECTOR_INDEX,
|
||||
indexedMetadataFields: POPULATED_METADATA_FIELDS,
|
||||
embeddingDefinition: {
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: 1536,
|
||||
},
|
||||
});
|
||||
|
||||
return await VectorStoreIndex.fromVectorStore(store);
|
||||
|
||||
@@ -5,6 +5,8 @@ const REQUIRED_ENV_VARS = [
|
||||
"MONGODB_VECTOR_INDEX",
|
||||
];
|
||||
|
||||
export const POPULATED_METADATA_FIELDS = ["private", "doc_id"]; // for filtering in MongoDB VectorSearchIndex
|
||||
|
||||
export function checkRequiredEnvVars() {
|
||||
const missingEnvVars = REQUIRED_ENV_VARS.filter((envVar) => {
|
||||
return !process.env[envVar];
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
This is a [LlamaIndex](https://www.llamaindex.ai/) project using [FastAPI](https://fastapi.tiangolo.com/) bootstrapped with [`create-llama`](https://github.com/run-llama/LlamaIndexTS/tree/main/packages/create-llama) featuring [structured extraction](https://docs.llamaindex.ai/en/stable/examples/structured_outputs/structured_outputs/?h=structured+output).
|
||||
This is a [LlamaIndex](https://www.llamaindex.ai/) project using [Reflex](https://reflex.dev/) bootstrapped with [`create-llama`](https://github.com/run-llama/LlamaIndexTS/tree/main/packages/create-llama) featuring [structured extraction](https://docs.llamaindex.ai/en/stable/examples/structured_outputs/structured_outputs/?h=structured+output) in a RAG pipeline.
|
||||
|
||||
## Getting Started
|
||||
|
||||
@@ -8,26 +8,38 @@ First, setup the environment with poetry:
|
||||
|
||||
```shell
|
||||
poetry install
|
||||
poetry shell
|
||||
```
|
||||
|
||||
Then check the parameters that have been pre-configured in the `.env` file in this directory. (E.g. you might need to configure an `OPENAI_API_KEY` if you're using OpenAI as model provider).
|
||||
|
||||
Second, generate the embeddings of the documents in the `./data` directory (if this folder exists - otherwise, skip this step):
|
||||
Second, generate the embeddings of the example document in the `./data` directory:
|
||||
|
||||
```shell
|
||||
poetry run generate
|
||||
```
|
||||
|
||||
Third, run the API in one command:
|
||||
Third, start app with `reflex` command:
|
||||
|
||||
```shell
|
||||
poetry run python main.py
|
||||
poetry run reflex run
|
||||
```
|
||||
|
||||
The example provides the `/api/extractor/query` API endpoint.
|
||||
To deploy the application, refer to the Reflex deployment guide: https://reflex.dev/docs/hosting/deploy-quick-start/
|
||||
|
||||
This query endpoint returns structured data in the format of the [Output](./app/api/routers/output.py) class. Modify this class to change the output format.
|
||||
### UI
|
||||
|
||||
You can now access the UI at http://localhost:3000 to test the structure extractor interactively.
|
||||
|
||||
It allows you to remove and add your own documents, modify the Pydantic model used for structured extraction, and test the RAG pipeline with different queries.
|
||||
|
||||
For example, keep the provided Pydantic model and query: "What is the maximum weight for a parcel?".
|
||||
|
||||
> Note: the Pydantic model used is the last element in the code provided by the user.
|
||||
|
||||
### API
|
||||
|
||||
Alternatively, check the API documentation at http://localhost:8000/docs. This example provides the `/api/extractor/query` API endpoint.
|
||||
Per default, the query endpoint returns structured data in the format of the model [DEFAULT_MODEL](./app/services/model.py) class. Modify this class to change the output format.
|
||||
|
||||
You can test the endpoint with the following curl request:
|
||||
|
||||
@@ -49,15 +61,9 @@ curl --location 'localhost:8000/api/extractor/query' \
|
||||
|
||||
To retrieve a response with low confidence since the question is not related to the provided document in the `./data` directory.
|
||||
|
||||
You can start editing the API endpoint by modifying [`extractor.py`](./app/api/routers/extractor.py). The endpoints auto-update as you save the file.
|
||||
### Development
|
||||
|
||||
Open [http://localhost:8000/docs](http://localhost:8000/docs) with your browser to see the Swagger UI of the API.
|
||||
|
||||
The API allows CORS for all origins to simplify development. You can change this behavior by setting the `ENVIRONMENT` environment variable to `prod`:
|
||||
|
||||
```
|
||||
ENVIRONMENT=prod python main.py
|
||||
```
|
||||
You can start editing the behavior by modifying the [`ExtractorService`](./app/services/extractor.py). The app auto-updates as you save the file.
|
||||
|
||||
## Learn More
|
||||
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.services.model import DEFAULT_MODEL
|
||||
|
||||
|
||||
class RequestData(BaseModel):
|
||||
query: str
|
||||
code: str = DEFAULT_MODEL
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"examples": [
|
||||
{
|
||||
"query": "What's the maximum weight for a parcel?",
|
||||
"code": DEFAULT_MODEL,
|
||||
},
|
||||
],
|
||||
}
|
||||
@@ -1,58 +1,15 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from llama_index.core.settings import Settings
|
||||
from pydantic import BaseModel
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.routers.output import Output
|
||||
from app.engine.index import get_index
|
||||
from app.api.models import RequestData
|
||||
from app.services.extractor import ExtractorService
|
||||
|
||||
extractor_router = r = APIRouter()
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class RequestData(BaseModel):
|
||||
query: str
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"examples": [
|
||||
{"query": "What's the maximum weight for a parcel?"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@r.post("/query")
|
||||
async def query_request(
|
||||
data: RequestData,
|
||||
):
|
||||
# Create a query engine using that returns responses in the format of the Output class
|
||||
query_engine = get_query_engine(Output)
|
||||
|
||||
response = await query_engine.aquery(data.query)
|
||||
|
||||
output_data = response.response.dict()
|
||||
return Output(**output_data)
|
||||
|
||||
|
||||
def get_query_engine(output_cls: BaseModel):
|
||||
top_k = os.getenv("TOP_K", 3)
|
||||
|
||||
index = get_index()
|
||||
if index is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=str(
|
||||
"StorageContext is empty - call 'poetry run generate' to generate the storage first"
|
||||
),
|
||||
)
|
||||
|
||||
sllm = Settings.llm.as_structured_llm(output_cls)
|
||||
|
||||
return index.as_query_engine(
|
||||
similarity_top_k=int(top_k),
|
||||
llm=sllm,
|
||||
response_mode="tree_summarize",
|
||||
)
|
||||
async def query_request(data: RequestData):
|
||||
return await ExtractorService.extract(query=data.query, model_code=data.code)
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.routers.extractor import extractor_router
|
||||
|
||||
api_router = APIRouter()
|
||||
|
||||
api_router.include_router(extractor_router, prefix="/api/extractor")
|
||||
@@ -0,0 +1,22 @@
|
||||
# flake8: noqa: E402
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
import reflex as rx
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.api.routers.extractor import extractor_router
|
||||
from app.settings import init_settings
|
||||
from app.ui.pages import * # Keep this import all pages in the app # noqa: F403
|
||||
|
||||
init_settings()
|
||||
|
||||
|
||||
def add_routers(app: FastAPI):
|
||||
app.include_router(extractor_router, prefix="/api/extractor")
|
||||
|
||||
|
||||
app = rx.App()
|
||||
add_routers(app.api)
|
||||
@@ -0,0 +1 @@
|
||||
DATA_DIR = "data"
|
||||
@@ -0,0 +1 @@
|
||||
from .engine import get_query_engine as get_query_engine
|
||||
@@ -0,0 +1,27 @@
|
||||
import os
|
||||
|
||||
from fastapi import HTTPException
|
||||
from llama_index.core.settings import Settings
|
||||
|
||||
from app.engine.index import get_index
|
||||
|
||||
|
||||
def get_query_engine(output_cls):
|
||||
top_k = int(os.getenv("TOP_K", 0))
|
||||
|
||||
index = get_index()
|
||||
if index is None:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=str(
|
||||
"StorageContext is empty - call 'poetry run generate' to generate the storage first"
|
||||
),
|
||||
)
|
||||
|
||||
sllm = Settings.llm.as_structured_llm(output_cls)
|
||||
|
||||
return index.as_query_engine(
|
||||
llm=sllm,
|
||||
response_mode="tree_summarize",
|
||||
**({"similarity_top_k": top_k} if top_k != 0 else {}),
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
# flake8: noqa: E402
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from llama_index.core.ingestion import IngestionPipeline
|
||||
from llama_index.core.node_parser import SentenceSplitter
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.core.storage import StorageContext
|
||||
from llama_index.core.storage.docstore import SimpleDocumentStore
|
||||
|
||||
from app.engine.loaders import get_documents
|
||||
from app.engine.vectordb import get_vector_store
|
||||
from app.settings import init_settings
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger()
|
||||
|
||||
STORAGE_DIR = os.getenv("STORAGE_DIR", "storage")
|
||||
|
||||
|
||||
def get_doc_store():
|
||||
# If the storage directory is there, load the document store from it.
|
||||
# If not, set up an in-memory document store since we can't load from a directory that doesn't exist.
|
||||
if os.path.exists(STORAGE_DIR):
|
||||
return SimpleDocumentStore.from_persist_dir(STORAGE_DIR)
|
||||
else:
|
||||
return SimpleDocumentStore()
|
||||
|
||||
|
||||
def run_pipeline(docstore, vector_store, documents):
|
||||
pipeline = IngestionPipeline(
|
||||
transformations=[
|
||||
SentenceSplitter(
|
||||
chunk_size=Settings.chunk_size,
|
||||
chunk_overlap=Settings.chunk_overlap,
|
||||
),
|
||||
Settings.embed_model,
|
||||
],
|
||||
docstore=docstore,
|
||||
docstore_strategy="upserts_and_delete",
|
||||
vector_store=vector_store,
|
||||
)
|
||||
|
||||
# Run the ingestion pipeline and store the results
|
||||
nodes = pipeline.run(show_progress=True, documents=documents)
|
||||
|
||||
return nodes
|
||||
|
||||
|
||||
def persist_storage(docstore, vector_store):
|
||||
storage_context = StorageContext.from_defaults(
|
||||
docstore=docstore,
|
||||
vector_store=vector_store,
|
||||
)
|
||||
storage_context.persist(STORAGE_DIR)
|
||||
|
||||
|
||||
def generate_datasource():
|
||||
init_settings()
|
||||
logger.info("Generate index for the provided data")
|
||||
|
||||
# Get the stores and documents or create new ones
|
||||
documents = get_documents()
|
||||
docstore = get_doc_store()
|
||||
vector_store = get_vector_store()
|
||||
|
||||
# Run the ingestion pipeline
|
||||
_ = run_pipeline(docstore, vector_store, documents)
|
||||
|
||||
# Build the index and persist storage
|
||||
persist_storage(docstore, vector_store)
|
||||
|
||||
logger.info("Finished generating the index")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generate_datasource()
|
||||
@@ -0,0 +1,31 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from llama_index.core.callbacks import CallbackManager
|
||||
from llama_index.core.indices import VectorStoreIndex
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.engine.vectordb import get_vector_store
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class IndexConfig(BaseModel):
|
||||
callback_manager: Optional[CallbackManager] = Field(
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def get_index(config: IndexConfig = None):
|
||||
if config is None:
|
||||
config = IndexConfig()
|
||||
logger.info("Connecting vector store...")
|
||||
store = get_vector_store()
|
||||
# Load the index from the vector store
|
||||
# If you are using a vector store that doesn't store text,
|
||||
# you must load the index from both the vector store and the document store
|
||||
index = VectorStoreIndex.from_vector_store(
|
||||
store, callback_manager=config.callback_manager
|
||||
)
|
||||
logger.info("Finished load index from vector store.")
|
||||
return index
|
||||
+2
-1
@@ -1,7 +1,8 @@
|
||||
import logging
|
||||
from llama_index.core.schema import BaseModel, Field
|
||||
from typing import List
|
||||
|
||||
from llama_index.core.schema import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
import logging
|
||||
|
||||
from app.engine import get_query_engine
|
||||
from app.services.model import IMPORTS
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class InvalidModelCode(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ExtractorService:
|
||||
@staticmethod
|
||||
def _parse_code(model_code: str):
|
||||
try:
|
||||
python_code = f"{IMPORTS}\n\n{model_code}"
|
||||
logger.debug(python_code)
|
||||
namespace = {}
|
||||
exec(python_code, namespace)
|
||||
# using the last object that the user defined in `model_code` as pydantic class
|
||||
pydantic_class = namespace[list(namespace.keys())[-1]]
|
||||
class_name = pydantic_class.__name__
|
||||
logger.info(f"Using Pydantic class {class_name} for extraction")
|
||||
return pydantic_class
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
raise InvalidModelCode() from e
|
||||
|
||||
@classmethod
|
||||
async def extract(cls, query: str, model_code: str) -> str:
|
||||
schema_model = cls._parse_code(model_code)
|
||||
# Create a query engine using that returns responses in the format of the schema
|
||||
query_engine = get_query_engine(schema_model)
|
||||
response = await query_engine.aquery(query)
|
||||
output_data = response.response.dict()
|
||||
return schema_model(**output_data).model_dump_json(indent=2)
|
||||
@@ -0,0 +1,22 @@
|
||||
IMPORTS = """
|
||||
from llama_index.core.schema import BaseModel, Field
|
||||
from typing import List, Optional
|
||||
from datetime import date
|
||||
"""
|
||||
|
||||
DEFAULT_MODEL = """class Output(BaseModel):
|
||||
response: str = Field(..., description="The answer to the question.")
|
||||
page_numbers: List[int] = Field(
|
||||
...,
|
||||
description="The page numbers of the sources used to answer this question. Do not include a page number if the context is irrelevant.",
|
||||
)
|
||||
confidence: float = Field(
|
||||
...,
|
||||
ge=0,
|
||||
le=1,
|
||||
description="Confidence value between 0-1 of the correctness of the result.",
|
||||
)
|
||||
confidence_explanation: str = Field(
|
||||
..., description="Explanation for the confidence score"
|
||||
)
|
||||
"""
|
||||
@@ -0,0 +1,9 @@
|
||||
from .extractor import (
|
||||
StructuredQuery as StructuredQuery,
|
||||
extract_data_component as extract_data_component,
|
||||
)
|
||||
from .schema_editor import schema_editor_component as schema_editor_component
|
||||
from .upload import (
|
||||
UploadedFilesState as UploadedFilesState,
|
||||
upload_component as upload_component,
|
||||
)
|
||||
@@ -0,0 +1,87 @@
|
||||
import logging
|
||||
import reflex as rx
|
||||
from app.services.model import DEFAULT_MODEL
|
||||
from app.services.extractor import ExtractorService, InvalidModelCode
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class StructuredQuery(rx.State):
|
||||
query: str
|
||||
response: str
|
||||
loading: bool = False
|
||||
code: str = DEFAULT_MODEL
|
||||
error: str = None
|
||||
|
||||
@rx.background
|
||||
async def handle_query(self):
|
||||
async with self:
|
||||
if not self.query:
|
||||
self.error = "Please enter a query."
|
||||
return
|
||||
self.error = None
|
||||
self.loading = True
|
||||
|
||||
# Extract data
|
||||
# Await long operations outside the context to avoid blocking UI
|
||||
try:
|
||||
response = await ExtractorService.extract(
|
||||
query=self.query, model_code=self.code
|
||||
)
|
||||
except InvalidModelCode:
|
||||
async with self:
|
||||
self.error = "Invalid Python code"
|
||||
response = None
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
logger.error(
|
||||
f"Error occurred: {str(e)}\nStack trace:\n{traceback.format_exc()}"
|
||||
)
|
||||
async with self:
|
||||
self.error = f"Error: {str(e)}"
|
||||
response = None
|
||||
|
||||
async with self:
|
||||
self.response = response
|
||||
self.loading = False
|
||||
|
||||
|
||||
def extract_data_component() -> rx.Component:
|
||||
return rx.vstack(
|
||||
rx.cond(
|
||||
StructuredQuery.error,
|
||||
rx.callout(
|
||||
StructuredQuery.error,
|
||||
icon="triangle_alert",
|
||||
color_scheme="red",
|
||||
role="alert",
|
||||
),
|
||||
),
|
||||
rx.text_area(
|
||||
id="query",
|
||||
placeholder="Enter query",
|
||||
on_change=StructuredQuery.set_query,
|
||||
width="100%",
|
||||
height="10vh",
|
||||
),
|
||||
rx.button(
|
||||
"Query",
|
||||
on_click=StructuredQuery.handle_query,
|
||||
loading=StructuredQuery.loading,
|
||||
),
|
||||
rx.cond(
|
||||
StructuredQuery.response,
|
||||
rx.code_block(
|
||||
StructuredQuery.response,
|
||||
language="json",
|
||||
show_line_numbers=True,
|
||||
wrap_long_lines=True,
|
||||
size="3",
|
||||
resize="vertical",
|
||||
width="100%",
|
||||
height="70vh",
|
||||
),
|
||||
),
|
||||
width="100%",
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Reflex custom component Monaco."""
|
||||
|
||||
# For wrapping react guide, visit https://reflex.dev/docs/wrapping-react/overview/
|
||||
# Taken and modified from https://github.com/Lendemor/reflex-monaco
|
||||
|
||||
import reflex as rx
|
||||
|
||||
|
||||
class MonacoComponent(rx.Component):
|
||||
"""Base Monaco component."""
|
||||
|
||||
library = "@monaco-editor/react@4.6.0"
|
||||
|
||||
# The language to use for the editor.
|
||||
language: rx.Var[str]
|
||||
|
||||
# The theme to use for the editor.
|
||||
theme: rx.Var[str] = rx.color_mode_cond("light", "vs-dark") # type: ignore
|
||||
|
||||
# The width of the editor.
|
||||
line: rx.Var[int] = rx.Var.create_safe(1, _var_is_string=False)
|
||||
|
||||
# The height of the editor.
|
||||
width: rx.Var[str]
|
||||
|
||||
# The height of the editor.
|
||||
height: rx.Var[str]
|
||||
|
||||
|
||||
class MonacoEditor(MonacoComponent):
|
||||
"""The Monaco Editor component."""
|
||||
|
||||
# The React component tag.
|
||||
tag = "MonacoEditor"
|
||||
|
||||
is_default = True
|
||||
|
||||
# The default value to display in the editor.
|
||||
default_value: rx.Var[str]
|
||||
|
||||
# The default language to use for the editor.
|
||||
default_language: rx.Var[str]
|
||||
|
||||
# The path to the default file to load in the editor.
|
||||
default_path: rx.Var[str]
|
||||
|
||||
# The value to display in the editor.
|
||||
value: rx.Var[str]
|
||||
|
||||
# Triggered when the editor value changes.
|
||||
on_change: rx.EventHandler[lambda e: [e]]
|
||||
|
||||
# Triggered when the content is validated. (limited to some languages)
|
||||
on_validate: rx.EventHandler[lambda e: [e]]
|
||||
|
||||
options = {
|
||||
"minimap": {"enabled": False},
|
||||
}
|
||||
|
||||
|
||||
monaco = MonacoEditor.create
|
||||
@@ -0,0 +1,18 @@
|
||||
import reflex as rx
|
||||
|
||||
from app.ui.components.extractor import StructuredQuery
|
||||
from .monaco import monaco
|
||||
|
||||
|
||||
def schema_editor_component() -> rx.Component:
|
||||
return rx.vstack(
|
||||
rx.heading("Pydantic model", size="5"),
|
||||
monaco(
|
||||
default_language="python",
|
||||
default_value=StructuredQuery.code,
|
||||
width="100%",
|
||||
height="50vh",
|
||||
on_change=StructuredQuery.set_code,
|
||||
),
|
||||
width="100%",
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import reflex as rx
|
||||
from app.engine.generate import generate_datasource
|
||||
|
||||
|
||||
class UploadedFile(rx.Base):
|
||||
file_name: str
|
||||
size: int
|
||||
|
||||
|
||||
class UploadedFilesState(rx.State):
|
||||
_uploaded_dir = "data"
|
||||
uploaded_files: List[UploadedFile] = []
|
||||
|
||||
async def handle_upload(self, files: list[rx.UploadFile]):
|
||||
for file in files:
|
||||
upload_data = await file.read()
|
||||
outfile = os.path.join(self._uploaded_dir, file.filename)
|
||||
with open(outfile, "wb") as f:
|
||||
f.write(upload_data)
|
||||
|
||||
new_file = UploadedFile(file_name=file.filename, size=len(upload_data))
|
||||
|
||||
self.uploaded_files.append(new_file)
|
||||
|
||||
# Run indexing
|
||||
try:
|
||||
generate_datasource()
|
||||
except Exception as e:
|
||||
print("Error generating datasource", e)
|
||||
os.remove(outfile)
|
||||
self.uploaded_files.remove(new_file)
|
||||
return rx.toast.error(
|
||||
f"Error generating index for the uploaded files. {str(e)}",
|
||||
position="top-center",
|
||||
)
|
||||
|
||||
return rx.toast.success("Files uploaded successfully", position="top-center")
|
||||
|
||||
def has_files(self) -> bool:
|
||||
return len(self.uploaded_files) > 0
|
||||
|
||||
def load_files(self):
|
||||
self.uploaded_files = []
|
||||
for file in os.listdir(self._uploaded_dir):
|
||||
file_path = os.path.join(self._uploaded_dir, file)
|
||||
if os.path.isfile(file_path):
|
||||
self.uploaded_files.append(
|
||||
UploadedFile(file_name=file, size=os.path.getsize(file_path))
|
||||
)
|
||||
|
||||
def remove_file(self, file_name: str):
|
||||
for file in self.uploaded_files:
|
||||
if file.file_name == file_name:
|
||||
self.uploaded_files.remove(file)
|
||||
os.remove(os.path.join(self._uploaded_dir, file_name))
|
||||
# Run indexing
|
||||
generate_datasource()
|
||||
break
|
||||
|
||||
|
||||
def upload_component() -> rx.Component:
|
||||
return rx.vstack(
|
||||
rx.heading("Upload", size="5"),
|
||||
rx.upload(
|
||||
rx.vstack(
|
||||
rx.text("Drag and drop files here or click to select files"),
|
||||
),
|
||||
on_drop=UploadedFilesState.handle_upload(
|
||||
rx.upload_files(upload_id="upload1")
|
||||
),
|
||||
id="upload1",
|
||||
border="1px dotted rgb(107,99,246)",
|
||||
),
|
||||
rx.foreach(
|
||||
UploadedFilesState.uploaded_files,
|
||||
lambda file: rx.card(
|
||||
rx.stack(
|
||||
rx.text(file.file_name, size="sm"),
|
||||
rx.button(
|
||||
"x",
|
||||
size="sm",
|
||||
on_click=UploadedFilesState.remove_file(file.file_name),
|
||||
),
|
||||
justify="between",
|
||||
width="100%",
|
||||
),
|
||||
width="100%",
|
||||
),
|
||||
),
|
||||
width="100%",
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
from .index import index as index
|
||||
@@ -0,0 +1,49 @@
|
||||
import reflex as rx
|
||||
|
||||
from ..components import (
|
||||
UploadedFilesState,
|
||||
extract_data_component,
|
||||
schema_editor_component,
|
||||
upload_component,
|
||||
)
|
||||
from ..templates import template
|
||||
|
||||
|
||||
@template(
|
||||
route="/",
|
||||
title="Structure extractor",
|
||||
on_load=[
|
||||
UploadedFilesState.load_files,
|
||||
],
|
||||
)
|
||||
def index() -> rx.Component:
|
||||
"""The main index page."""
|
||||
return rx.vstack(
|
||||
rx.vstack(
|
||||
rx.heading("Built by LlamaIndex", size="6"),
|
||||
rx.text(
|
||||
"Upload a file then enter a query. The response will be according to the provided Pydantic model."
|
||||
),
|
||||
background_color="var(--gray-3)",
|
||||
align_items="left",
|
||||
justify_content="left",
|
||||
width="100%",
|
||||
padding="1rem",
|
||||
),
|
||||
rx.stack(
|
||||
rx.vstack(
|
||||
upload_component(),
|
||||
rx.divider(),
|
||||
schema_editor_component(),
|
||||
width="50%",
|
||||
),
|
||||
rx.divider(orientation="vertical"),
|
||||
rx.stack(
|
||||
extract_data_component(),
|
||||
width="50%",
|
||||
),
|
||||
width="100%",
|
||||
padding="1rem",
|
||||
),
|
||||
width="100%",
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
from .template import ThemeState as ThemeState, template as template
|
||||
@@ -0,0 +1,28 @@
|
||||
import reflex as rx
|
||||
|
||||
border_radius = "var(--radius-2)"
|
||||
border = f"1px solid {rx.color('gray', 5)}"
|
||||
text_color = rx.color("gray", 11)
|
||||
gray_color = rx.color("gray", 11)
|
||||
gray_bg_color = rx.color("gray", 3)
|
||||
accent_text_color = rx.color("accent", 10)
|
||||
accent_color = rx.color("accent", 1)
|
||||
accent_bg_color = rx.color("accent", 3)
|
||||
hover_accent_color = {"_hover": {"color": accent_text_color}}
|
||||
hover_accent_bg = {"_hover": {"background_color": accent_color}}
|
||||
content_width_vw = "90vw"
|
||||
sidebar_width = "32em"
|
||||
sidebar_content_width = "16em"
|
||||
color_box_size = ["2.25rem", "2.25rem", "2.5rem"]
|
||||
|
||||
|
||||
template_page_style = {
|
||||
"padding_top": ["1em", "1em", "2em"],
|
||||
"padding_x": ["auto", "auto", "2em"],
|
||||
}
|
||||
|
||||
template_content_style = {
|
||||
"padding": "1em",
|
||||
"margin_bottom": "2em",
|
||||
"min_height": "90vh",
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Common templates used between pages in the app."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Callable
|
||||
|
||||
import reflex as rx
|
||||
|
||||
from . import styles
|
||||
|
||||
# Meta tags for the app.
|
||||
default_meta = [
|
||||
{
|
||||
"name": "viewport",
|
||||
"content": "width=device-width, shrink-to-fit=no, initial-scale=1",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class ThemeState(rx.State):
|
||||
"""The state for the theme of the app."""
|
||||
|
||||
accent_color: str = "crimson"
|
||||
|
||||
gray_color: str = "gray"
|
||||
|
||||
radius: str = "large"
|
||||
|
||||
scaling: str = "100%"
|
||||
|
||||
|
||||
def template(
|
||||
route: str | None = None,
|
||||
title: str | None = None,
|
||||
description: str | None = None,
|
||||
meta: str | None = None,
|
||||
script_tags: list[rx.Component] | None = None,
|
||||
on_load: rx.event.EventHandler | list[rx.event.EventHandler] | None = None,
|
||||
) -> Callable[[Callable[[], rx.Component]], rx.Component]:
|
||||
"""The template for each page of the app.
|
||||
|
||||
Args:
|
||||
route: The route to reach the page.
|
||||
title: The title of the page.
|
||||
description: The description of the page.
|
||||
meta: Additional meta to add to the page.
|
||||
on_load: The event handler(s) called when the page load.
|
||||
script_tags: Scripts to attach to the page.
|
||||
|
||||
Returns:
|
||||
The template with the page content.
|
||||
"""
|
||||
|
||||
def decorator(page_content: Callable[[], rx.Component]) -> rx.Component:
|
||||
"""The template for each page of the app.
|
||||
|
||||
Args:
|
||||
page_content: The content of the page.
|
||||
|
||||
Returns:
|
||||
The template with the page content.
|
||||
"""
|
||||
# Get the meta tags for the page.
|
||||
all_meta = [*default_meta, *(meta or [])]
|
||||
|
||||
def templated_page():
|
||||
return rx.flex(
|
||||
rx.flex(
|
||||
rx.vstack(
|
||||
page_content(),
|
||||
width="100%",
|
||||
**styles.template_content_style,
|
||||
),
|
||||
width="100%",
|
||||
**styles.template_page_style,
|
||||
max_width=[
|
||||
"100%",
|
||||
"100%",
|
||||
"100%",
|
||||
"100%",
|
||||
"100%",
|
||||
],
|
||||
),
|
||||
flex_direction=[
|
||||
"column",
|
||||
"column",
|
||||
"column",
|
||||
"column",
|
||||
"column",
|
||||
"row",
|
||||
],
|
||||
width="100%",
|
||||
margin="auto",
|
||||
position="relative",
|
||||
)
|
||||
|
||||
@rx.page(
|
||||
route=route,
|
||||
title=title,
|
||||
description=description,
|
||||
meta=all_meta,
|
||||
script_tags=script_tags,
|
||||
on_load=on_load,
|
||||
)
|
||||
def theme_wrap():
|
||||
return rx.theme(
|
||||
templated_page(),
|
||||
has_background=True,
|
||||
accent_color=ThemeState.accent_color,
|
||||
gray_color=ThemeState.gray_color,
|
||||
radius=ThemeState.radius,
|
||||
scaling=ThemeState.scaling,
|
||||
)
|
||||
|
||||
return theme_wrap
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,4 @@
|
||||
__pycache__
|
||||
storage
|
||||
.env
|
||||
output
|
||||
@@ -1,45 +0,0 @@
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
import logging
|
||||
import os
|
||||
import uvicorn
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import RedirectResponse
|
||||
from app.api.routers.extractor import extractor_router
|
||||
from app.settings import init_settings
|
||||
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
init_settings()
|
||||
|
||||
environment = os.getenv("ENVIRONMENT", "dev") # Default to 'development' if not set
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
if environment == "dev":
|
||||
logger.warning("Running in development mode - allowing CORS for all origins")
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Redirect to documentation page when accessing base URL
|
||||
@app.get("/")
|
||||
async def redirect_to_docs():
|
||||
return RedirectResponse(url="/docs")
|
||||
|
||||
|
||||
app.include_router(extractor_router, prefix="/api/extractor")
|
||||
|
||||
if __name__ == "__main__":
|
||||
app_host = os.getenv("APP_HOST", "0.0.0.0")
|
||||
app_port = int(os.getenv("APP_PORT", "8000"))
|
||||
reload = True if environment == "dev" else False
|
||||
|
||||
uvicorn.run(app="main:app", host=app_host, port=app_port, reload=reload)
|
||||
@@ -13,8 +13,9 @@ python = "^3.11,<4.0"
|
||||
fastapi = "^0.109.1"
|
||||
uvicorn = { extras = ["standard"], version = "^0.23.2" }
|
||||
python-dotenv = "^1.0.0"
|
||||
llama-index = "^0.10.58"
|
||||
llama-index = "^0.11.1"
|
||||
cachetools = "^5.3.3"
|
||||
reflex = "^0.5.9"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
import reflex as rx
|
||||
|
||||
config = rx.Config(
|
||||
app_name="app",
|
||||
telemetry_enabled=False,
|
||||
)
|
||||
@@ -1,4 +1,18 @@
|
||||
This is a [LlamaIndex](https://www.llamaindex.ai/) project using [FastAPI](https://fastapi.tiangolo.com/) bootstrapped with [`create-llama`](https://github.com/run-llama/LlamaIndexTS/tree/main/packages/create-llama).
|
||||
This is a [LlamaIndex](https://www.llamaindex.ai/) multi-agents project using [Workflows](https://docs.llamaindex.ai/en/stable/understanding/workflows/).
|
||||
|
||||
## Overview
|
||||
|
||||
This example is using three agents to generate a blog post:
|
||||
|
||||
- a researcher that retrieves content via a RAG pipeline,
|
||||
- a writer that specializes in writing blog posts and
|
||||
- a reviewer that is reviewing the blog post.
|
||||
|
||||
There are three different methods how the agents can interact to reach their goal:
|
||||
|
||||
1. [Choreography](./app/examples/choreography.py) - the agents decide themselves to delegate a task to another agent
|
||||
1. [Orchestrator](./app/examples/orchestrator.py) - a central orchestrator decides which agent should execute a task
|
||||
1. [Explicit Workflow](./app/examples/workflow.py) - a pre-defined workflow specific for the task is used to execute the tasks
|
||||
|
||||
## Getting Started
|
||||
|
||||
@@ -8,43 +22,48 @@ First, setup the environment with poetry:
|
||||
|
||||
```shell
|
||||
poetry install
|
||||
poetry shell
|
||||
```
|
||||
|
||||
Then check the parameters that have been pre-configured in the `.env` file in this directory. (E.g. you might need to configure an `OPENAI_API_KEY` if you're using OpenAI as model provider).
|
||||
|
||||
Second, generate the embeddings of the documents in the `./data` directory (if this folder exists - otherwise, skip this step):
|
||||
Second, generate the embeddings of the documents in the `./data` directory:
|
||||
|
||||
```shell
|
||||
poetry run generate
|
||||
```
|
||||
|
||||
Third, run all the services in one command:
|
||||
Third, run the development server:
|
||||
|
||||
```shell
|
||||
poetry run python main.py
|
||||
```
|
||||
|
||||
You can monitor and test the agent services with `llama-agents` monitor TUI:
|
||||
Per default, the example is using the explicit workflow. You can change the example by setting the `EXAMPLE_TYPE` environment variable to `choreography` or `orchestrator`.
|
||||
|
||||
```shell
|
||||
poetry run llama-agents monitor --control-plane-url http://127.0.0.1:8001
|
||||
The example provides one streaming API endpoint `/api/chat`.
|
||||
You can test the endpoint with the following curl request:
|
||||
|
||||
```
|
||||
curl --location 'localhost:8000/api/chat' \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data '{ "messages": [{ "role": "user", "content": "Write a blog post about physical standards for letters" }] }'
|
||||
```
|
||||
|
||||
## Services:
|
||||
You can start editing the API by modifying `app/api/routers/chat.py` or `app/examples/workflow.py`. The API auto-updates as you save the files.
|
||||
|
||||
- Message queue (port 8000): To exchange the message between services
|
||||
- Control plane (port 8001): A gateway to manage the tasks and services.
|
||||
- Human consumer (port 8002): To handle result when the task is completed.
|
||||
- Agent service `query_engine` (port 8003): Agent that can query information from the configured LlamaIndex index.
|
||||
- Agent service `dummy_agent` (port 8004): A dummy agent that does nothing. Good starting point to add more agents.
|
||||
Open [http://localhost:8000/docs](http://localhost:8000/docs) with your browser to see the Swagger UI of the API.
|
||||
|
||||
The ports listed above are set by default, but you can change them in the `.env` file.
|
||||
The API allows CORS for all origins to simplify development. You can change this behavior by setting the `ENVIRONMENT` environment variable to `prod`:
|
||||
|
||||
```
|
||||
ENVIRONMENT=prod poetry run python main.py
|
||||
```
|
||||
|
||||
## Learn More
|
||||
|
||||
To learn more about LlamaIndex, take a look at the following resources:
|
||||
|
||||
- [LlamaIndex Documentation](https://docs.llamaindex.ai) - learn about LlamaIndex.
|
||||
- [Workflows Introduction](https://docs.llamaindex.ai/en/stable/understanding/workflows/) - learn about LlamaIndex workflows.
|
||||
|
||||
You can check out [the LlamaIndex GitHub repository](https://github.com/run-llama/llama_index) - your feedback and contributions are welcome!
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
from llama_agents import AgentService, SimpleMessageQueue
|
||||
from llama_index.core.agent import FunctionCallingAgentWorker
|
||||
from llama_index.core.tools import FunctionTool
|
||||
from llama_index.core.settings import Settings
|
||||
from app.utils import load_from_env
|
||||
|
||||
|
||||
DEFAULT_DUMMY_AGENT_DESCRIPTION = "I'm a dummy agent which does nothing."
|
||||
|
||||
|
||||
def dummy_function():
|
||||
"""
|
||||
This function does nothing.
|
||||
"""
|
||||
return ""
|
||||
|
||||
|
||||
def init_dummy_agent(message_queue: SimpleMessageQueue) -> AgentService:
|
||||
agent = FunctionCallingAgentWorker(
|
||||
tools=[FunctionTool.from_defaults(fn=dummy_function)],
|
||||
llm=Settings.llm,
|
||||
prefix_messages=[],
|
||||
).as_agent()
|
||||
|
||||
return AgentService(
|
||||
service_name="dummy_agent",
|
||||
agent=agent,
|
||||
message_queue=message_queue.client,
|
||||
description=load_from_env("AGENT_DUMMY_DESCRIPTION", throw_error=False)
|
||||
or DEFAULT_DUMMY_AGENT_DESCRIPTION,
|
||||
host=load_from_env("AGENT_DUMMY_HOST", throw_error=False) or "127.0.0.1",
|
||||
port=int(load_from_env("AGENT_DUMMY_PORT")),
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
from typing import Any, List
|
||||
|
||||
from app.agents.planner import StructuredPlannerAgent
|
||||
from app.agents.single import (
|
||||
AgentRunResult,
|
||||
ContextAwareTool,
|
||||
FunctionCallingAgent,
|
||||
)
|
||||
from llama_index.core.tools.types import ToolMetadata, ToolOutput
|
||||
from llama_index.core.tools.utils import create_schema_from_function
|
||||
from llama_index.core.workflow import Context, Workflow
|
||||
|
||||
|
||||
class AgentCallTool(ContextAwareTool):
|
||||
def __init__(self, agent: Workflow) -> None:
|
||||
self.agent = agent
|
||||
name = f"call_{agent.name}"
|
||||
|
||||
async def schema_call(input: str) -> str:
|
||||
pass
|
||||
|
||||
# create the schema without the Context
|
||||
fn_schema = create_schema_from_function(name, schema_call)
|
||||
self._metadata = ToolMetadata(
|
||||
name=name,
|
||||
description=(
|
||||
f"Use this tool to delegate a sub task to the {agent.name} agent."
|
||||
+ (f" The agent is an {agent.role}." if agent.role else "")
|
||||
),
|
||||
fn_schema=fn_schema,
|
||||
)
|
||||
|
||||
# overload the acall function with the ctx argument as it's needed for bubbling the events
|
||||
async def acall(self, ctx: Context, input: str) -> ToolOutput:
|
||||
handler = self.agent.run(input=input)
|
||||
# bubble all events while running the agent to the calling agent
|
||||
async for ev in handler.stream_events():
|
||||
ctx.write_event_to_stream(ev)
|
||||
ret: AgentRunResult = await handler
|
||||
response = ret.response.message.content
|
||||
return ToolOutput(
|
||||
content=str(response),
|
||||
tool_name=self.metadata.name,
|
||||
raw_input={"args": input, "kwargs": {}},
|
||||
raw_output=response,
|
||||
)
|
||||
|
||||
|
||||
class AgentCallingAgent(FunctionCallingAgent):
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
name: str,
|
||||
agents: List[FunctionCallingAgent] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
agents = agents or []
|
||||
tools = [AgentCallTool(agent=agent) for agent in agents]
|
||||
super().__init__(*args, name=name, tools=tools, **kwargs)
|
||||
# call add_workflows so agents will get detected by llama agents automatically
|
||||
self.add_workflows(**{agent.name: agent for agent in agents})
|
||||
|
||||
|
||||
class AgentOrchestrator(StructuredPlannerAgent):
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
name: str = "orchestrator",
|
||||
agents: List[FunctionCallingAgent] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
agents = agents or []
|
||||
tools = [AgentCallTool(agent=agent) for agent in agents]
|
||||
super().__init__(
|
||||
*args,
|
||||
name=name,
|
||||
tools=tools,
|
||||
**kwargs,
|
||||
)
|
||||
# call add_workflows so agents will get detected by llama agents automatically
|
||||
self.add_workflows(**{agent.name: agent for agent in agents})
|
||||
@@ -0,0 +1,324 @@
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from app.agents.single import AgentRunEvent, AgentRunResult, FunctionCallingAgent
|
||||
from llama_index.core.agent.runner.planner import (
|
||||
DEFAULT_INITIAL_PLAN_PROMPT,
|
||||
DEFAULT_PLAN_REFINE_PROMPT,
|
||||
Plan,
|
||||
PlannerAgentState,
|
||||
SubTask,
|
||||
)
|
||||
from llama_index.core.bridge.pydantic import ValidationError
|
||||
from llama_index.core.llms.function_calling import FunctionCallingLLM
|
||||
from llama_index.core.prompts import PromptTemplate
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.core.tools import BaseTool
|
||||
from llama_index.core.workflow import (
|
||||
Context,
|
||||
Event,
|
||||
StartEvent,
|
||||
StopEvent,
|
||||
Workflow,
|
||||
step,
|
||||
)
|
||||
|
||||
|
||||
class ExecutePlanEvent(Event):
|
||||
pass
|
||||
|
||||
|
||||
class SubTaskEvent(Event):
|
||||
sub_task: SubTask
|
||||
|
||||
|
||||
class SubTaskResultEvent(Event):
|
||||
sub_task: SubTask
|
||||
result: AgentRunResult | AsyncGenerator
|
||||
|
||||
|
||||
class PlanEventType(Enum):
|
||||
CREATED = "created"
|
||||
REFINED = "refined"
|
||||
|
||||
|
||||
class PlanEvent(AgentRunEvent):
|
||||
event_type: PlanEventType
|
||||
plan: Plan
|
||||
|
||||
@property
|
||||
def msg(self) -> str:
|
||||
sub_task_names = ", ".join(task.name for task in self.plan.sub_tasks)
|
||||
return f"Plan {self.event_type.value}: Let's do: {sub_task_names}"
|
||||
|
||||
|
||||
class StructuredPlannerAgent(Workflow):
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
name: str,
|
||||
llm: FunctionCallingLLM | None = None,
|
||||
tools: List[BaseTool] | None = None,
|
||||
timeout: float = 360.0,
|
||||
refine_plan: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(*args, timeout=timeout, **kwargs)
|
||||
self.name = name
|
||||
self.refine_plan = refine_plan
|
||||
|
||||
self.tools = tools or []
|
||||
self.planner = Planner(llm=llm, tools=self.tools, verbose=self._verbose)
|
||||
# The executor is keeping the memory of all tool calls and decides to call the right tool for the task
|
||||
self.executor = FunctionCallingAgent(
|
||||
name="executor",
|
||||
llm=llm,
|
||||
tools=self.tools,
|
||||
write_events=False,
|
||||
# it's important to instruct to just return the tool call, otherwise the executor will interpret and change the result
|
||||
system_prompt="You are an expert in completing given tasks by calling the right tool for the task. Just return the result of the tool call. Don't add any information yourself",
|
||||
)
|
||||
self.add_workflows(executor=self.executor)
|
||||
|
||||
@step()
|
||||
async def create_plan(
|
||||
self, ctx: Context, ev: StartEvent
|
||||
) -> ExecutePlanEvent | StopEvent:
|
||||
# set streaming
|
||||
ctx.data["streaming"] = getattr(ev, "streaming", False)
|
||||
ctx.data["task"] = ev.input
|
||||
|
||||
plan_id, plan = await self.planner.create_plan(input=ev.input)
|
||||
ctx.data["act_plan_id"] = plan_id
|
||||
|
||||
# inform about the new plan
|
||||
ctx.write_event_to_stream(
|
||||
PlanEvent(name=self.name, event_type=PlanEventType.CREATED, plan=plan)
|
||||
)
|
||||
if self._verbose:
|
||||
print("=== Executing plan ===\n")
|
||||
return ExecutePlanEvent()
|
||||
|
||||
@step()
|
||||
async def execute_plan(self, ctx: Context, ev: ExecutePlanEvent) -> SubTaskEvent:
|
||||
upcoming_sub_tasks = self.planner.state.get_next_sub_tasks(
|
||||
ctx.data["act_plan_id"]
|
||||
)
|
||||
|
||||
ctx.data["num_sub_tasks"] = len(upcoming_sub_tasks)
|
||||
# send an event per sub task
|
||||
events = [SubTaskEvent(sub_task=sub_task) for sub_task in upcoming_sub_tasks]
|
||||
for event in events:
|
||||
ctx.send_event(event)
|
||||
|
||||
return None
|
||||
|
||||
@step()
|
||||
async def execute_sub_task(
|
||||
self, ctx: Context, ev: SubTaskEvent
|
||||
) -> SubTaskResultEvent:
|
||||
if self._verbose:
|
||||
print(f"=== Executing sub task: {ev.sub_task.name} ===")
|
||||
is_last_tasks = ctx.data["num_sub_tasks"] == self.get_remaining_subtasks(ctx)
|
||||
# TODO: streaming only works without plan refining
|
||||
streaming = is_last_tasks and ctx.data["streaming"] and not self.refine_plan
|
||||
handler = self.executor.run(
|
||||
input=ev.sub_task.input,
|
||||
streaming=streaming,
|
||||
)
|
||||
# bubble all events while running the executor to the planner
|
||||
async for event in handler.stream_events():
|
||||
ctx.write_event_to_stream(event)
|
||||
result: AgentRunResult = await handler
|
||||
if self._verbose:
|
||||
print("=== Done executing sub task ===\n")
|
||||
self.planner.state.add_completed_sub_task(ctx.data["act_plan_id"], ev.sub_task)
|
||||
return SubTaskResultEvent(sub_task=ev.sub_task, result=result)
|
||||
|
||||
@step()
|
||||
async def gather_results(
|
||||
self, ctx: Context, ev: SubTaskResultEvent
|
||||
) -> ExecutePlanEvent | StopEvent:
|
||||
# wait for all sub tasks to finish
|
||||
num_sub_tasks = ctx.data["num_sub_tasks"]
|
||||
results = ctx.collect_events(ev, [SubTaskResultEvent] * num_sub_tasks)
|
||||
if results is None:
|
||||
return None
|
||||
|
||||
upcoming_sub_tasks = self.get_upcoming_sub_tasks(ctx)
|
||||
# if no more tasks to do, stop workflow and send result of last step
|
||||
if upcoming_sub_tasks == 0:
|
||||
return StopEvent(result=results[-1].result)
|
||||
|
||||
if self.refine_plan:
|
||||
# store all results for refining the plan
|
||||
ctx.data["results"] = ctx.data.get("results", {})
|
||||
for result in results:
|
||||
ctx.data["results"][result.sub_task.name] = result.result
|
||||
|
||||
new_plan = await self.planner.refine_plan(
|
||||
ctx.data["task"], ctx.data["act_plan_id"], ctx.data["results"]
|
||||
)
|
||||
# inform about the new plan
|
||||
if new_plan is not None:
|
||||
ctx.write_event_to_stream(
|
||||
PlanEvent(
|
||||
name=self.name, event_type=PlanEventType.REFINED, plan=new_plan
|
||||
)
|
||||
)
|
||||
|
||||
# continue executing plan
|
||||
return ExecutePlanEvent()
|
||||
|
||||
def get_upcoming_sub_tasks(self, ctx: Context):
|
||||
upcoming_sub_tasks = self.planner.state.get_next_sub_tasks(
|
||||
ctx.data["act_plan_id"]
|
||||
)
|
||||
return len(upcoming_sub_tasks)
|
||||
|
||||
def get_remaining_subtasks(self, ctx: Context):
|
||||
remaining_subtasks = self.planner.state.get_remaining_subtasks(
|
||||
ctx.data["act_plan_id"]
|
||||
)
|
||||
return len(remaining_subtasks)
|
||||
|
||||
|
||||
# Concern dealing with creating and refining a plan, extracted from https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/agent/runner/planner.py#L138
|
||||
class Planner:
|
||||
def __init__(
|
||||
self,
|
||||
llm: FunctionCallingLLM | None = None,
|
||||
tools: List[BaseTool] | None = None,
|
||||
initial_plan_prompt: Union[str, PromptTemplate] = DEFAULT_INITIAL_PLAN_PROMPT,
|
||||
plan_refine_prompt: Union[str, PromptTemplate] = DEFAULT_PLAN_REFINE_PROMPT,
|
||||
verbose: bool = True,
|
||||
) -> None:
|
||||
if llm is None:
|
||||
llm = Settings.llm
|
||||
self.llm = llm
|
||||
assert self.llm.metadata.is_function_calling_model
|
||||
|
||||
self.tools = tools or []
|
||||
self.state = PlannerAgentState()
|
||||
self.verbose = verbose
|
||||
|
||||
if isinstance(initial_plan_prompt, str):
|
||||
initial_plan_prompt = PromptTemplate(initial_plan_prompt)
|
||||
self.initial_plan_prompt = initial_plan_prompt
|
||||
|
||||
if isinstance(plan_refine_prompt, str):
|
||||
plan_refine_prompt = PromptTemplate(plan_refine_prompt)
|
||||
self.plan_refine_prompt = plan_refine_prompt
|
||||
|
||||
async def create_plan(self, input: str) -> Tuple[str, Plan]:
|
||||
tools = self.tools
|
||||
tools_str = ""
|
||||
for tool in tools:
|
||||
tools_str += tool.metadata.name + ": " + tool.metadata.description + "\n"
|
||||
|
||||
try:
|
||||
plan = await self.llm.astructured_predict(
|
||||
Plan,
|
||||
self.initial_plan_prompt,
|
||||
tools_str=tools_str,
|
||||
task=input,
|
||||
)
|
||||
except (ValueError, ValidationError):
|
||||
if self.verbose:
|
||||
print("No complex plan predicted. Defaulting to a single task plan.")
|
||||
plan = Plan(
|
||||
sub_tasks=[
|
||||
SubTask(
|
||||
name="default", input=input, expected_output="", dependencies=[]
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
print("=== Initial plan ===")
|
||||
for sub_task in plan.sub_tasks:
|
||||
print(
|
||||
f"{sub_task.name}:\n{sub_task.input} -> {sub_task.expected_output}\ndeps: {sub_task.dependencies}\n\n"
|
||||
)
|
||||
|
||||
plan_id = str(uuid.uuid4())
|
||||
self.state.plan_dict[plan_id] = plan
|
||||
|
||||
return plan_id, plan
|
||||
|
||||
async def refine_plan(
|
||||
self,
|
||||
input: str,
|
||||
plan_id: str,
|
||||
completed_sub_tasks: Dict[str, str],
|
||||
) -> Optional[Plan]:
|
||||
"""Refine a plan."""
|
||||
prompt_args = self.get_refine_plan_prompt_kwargs(
|
||||
plan_id, input, completed_sub_tasks
|
||||
)
|
||||
|
||||
try:
|
||||
new_plan = await self.llm.astructured_predict(
|
||||
Plan, self.plan_refine_prompt, **prompt_args
|
||||
)
|
||||
|
||||
self._update_plan(plan_id, new_plan)
|
||||
|
||||
return new_plan
|
||||
except (ValueError, ValidationError) as e:
|
||||
# likely no new plan predicted
|
||||
if self.verbose:
|
||||
print(f"No new plan predicted: {e}")
|
||||
return None
|
||||
|
||||
def _update_plan(self, plan_id: str, new_plan: Plan) -> None:
|
||||
"""Update the plan."""
|
||||
# update state with new plan
|
||||
self.state.plan_dict[plan_id] = new_plan
|
||||
|
||||
if self.verbose:
|
||||
print("=== Refined plan ===")
|
||||
for sub_task in new_plan.sub_tasks:
|
||||
print(
|
||||
f"{sub_task.name}:\n{sub_task.input} -> {sub_task.expected_output}\ndeps: {sub_task.dependencies}\n\n"
|
||||
)
|
||||
|
||||
def get_refine_plan_prompt_kwargs(
|
||||
self,
|
||||
plan_id: str,
|
||||
task: str,
|
||||
completed_sub_task: Dict[str, str],
|
||||
) -> dict:
|
||||
"""Get the refine plan prompt."""
|
||||
# gather completed sub-tasks and response pairs
|
||||
completed_outputs_str = ""
|
||||
for sub_task_name, task_output in completed_sub_task.items():
|
||||
task_str = f"{sub_task_name}:\n" f"\t{task_output!s}\n"
|
||||
completed_outputs_str += task_str
|
||||
|
||||
# get a string for the remaining sub-tasks
|
||||
remaining_sub_tasks = self.state.get_remaining_subtasks(plan_id)
|
||||
remaining_sub_tasks_str = "" if len(remaining_sub_tasks) != 0 else "None"
|
||||
for sub_task in remaining_sub_tasks:
|
||||
task_str = (
|
||||
f"SubTask(name='{sub_task.name}', "
|
||||
f"input='{sub_task.input}', "
|
||||
f"expected_output='{sub_task.expected_output}', "
|
||||
f"dependencies='{sub_task.dependencies}')\n"
|
||||
)
|
||||
remaining_sub_tasks_str += task_str
|
||||
|
||||
# get the tools string
|
||||
tools = self.tools
|
||||
tools_str = ""
|
||||
for tool in tools:
|
||||
tools_str += tool.metadata.name + ": " + tool.metadata.description + "\n"
|
||||
|
||||
# return the kwargs
|
||||
return {
|
||||
"tools_str": tools_str.strip(),
|
||||
"task": task.strip(),
|
||||
"completed_outputs": completed_outputs_str.strip(),
|
||||
"remaining_sub_tasks": remaining_sub_tasks_str.strip(),
|
||||
}
|
||||
@@ -1,52 +0,0 @@
|
||||
import os
|
||||
from llama_agents import AgentService, SimpleMessageQueue
|
||||
from llama_index.core.agent import FunctionCallingAgentWorker
|
||||
from llama_index.core.tools import QueryEngineTool, ToolMetadata
|
||||
from llama_index.core.settings import Settings
|
||||
from app.engine.index import get_index
|
||||
from app.utils import load_from_env
|
||||
|
||||
|
||||
DEFAULT_QUERY_ENGINE_AGENT_DESCRIPTION = (
|
||||
"Used to answer the questions using the provided context data."
|
||||
)
|
||||
|
||||
|
||||
def get_query_engine_tool() -> QueryEngineTool:
|
||||
"""
|
||||
Provide an agent worker that can be used to query the index.
|
||||
"""
|
||||
index = get_index()
|
||||
if index is None:
|
||||
raise ValueError("Index not found. Please create an index first.")
|
||||
query_engine = index.as_query_engine(similarity_top_k=int(os.getenv("TOP_K", 3)))
|
||||
return QueryEngineTool(
|
||||
query_engine=query_engine,
|
||||
metadata=ToolMetadata(
|
||||
name="context_data",
|
||||
description="""
|
||||
Provide the provided context information.
|
||||
Use a detailed plain text question as input to the tool.
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def init_query_engine_agent(
|
||||
message_queue: SimpleMessageQueue,
|
||||
) -> AgentService:
|
||||
"""
|
||||
Initialize the agent service.
|
||||
"""
|
||||
agent = FunctionCallingAgentWorker(
|
||||
tools=[get_query_engine_tool()], llm=Settings.llm, prefix_messages=[]
|
||||
).as_agent()
|
||||
return AgentService(
|
||||
service_name="context_query_agent",
|
||||
agent=agent,
|
||||
message_queue=message_queue.client,
|
||||
description=load_from_env("AGENT_QUERY_ENGINE_DESCRIPTION", throw_error=False)
|
||||
or DEFAULT_QUERY_ENGINE_AGENT_DESCRIPTION,
|
||||
host=load_from_env("AGENT_QUERY_ENGINE_HOST", throw_error=False) or "127.0.0.1",
|
||||
port=int(load_from_env("AGENT_QUERY_ENGINE_PORT")),
|
||||
)
|
||||
@@ -0,0 +1,245 @@
|
||||
from abc import abstractmethod
|
||||
from typing import Any, AsyncGenerator, List, Optional
|
||||
|
||||
from llama_index.core.llms import ChatMessage, ChatResponse
|
||||
from llama_index.core.llms.function_calling import FunctionCallingLLM
|
||||
from llama_index.core.memory import ChatMemoryBuffer
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.core.tools import ToolOutput, ToolSelection
|
||||
from llama_index.core.tools.types import BaseTool
|
||||
from llama_index.core.tools import FunctionTool
|
||||
|
||||
from llama_index.core.workflow import (
|
||||
Context,
|
||||
Event,
|
||||
StartEvent,
|
||||
StopEvent,
|
||||
Workflow,
|
||||
step,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class InputEvent(Event):
|
||||
input: list[ChatMessage]
|
||||
|
||||
|
||||
class ToolCallEvent(Event):
|
||||
tool_calls: list[ToolSelection]
|
||||
|
||||
|
||||
class AgentRunEvent(Event):
|
||||
name: str
|
||||
_msg: str
|
||||
|
||||
@property
|
||||
def msg(self):
|
||||
return self._msg
|
||||
|
||||
@msg.setter
|
||||
def msg(self, value):
|
||||
self._msg = value
|
||||
|
||||
|
||||
class AgentRunResult(BaseModel):
|
||||
response: ChatResponse
|
||||
sources: list[ToolOutput]
|
||||
|
||||
|
||||
class ContextAwareTool(FunctionTool):
|
||||
@abstractmethod
|
||||
async def acall(self, ctx: Context, input: Any) -> ToolOutput:
|
||||
pass
|
||||
|
||||
|
||||
class FunctionCallingAgent(Workflow):
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
llm: FunctionCallingLLM | None = None,
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
tools: List[BaseTool] | None = None,
|
||||
system_prompt: str | None = None,
|
||||
verbose: bool = False,
|
||||
timeout: float = 360.0,
|
||||
name: str,
|
||||
write_events: bool = True,
|
||||
role: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(*args, verbose=verbose, timeout=timeout, **kwargs)
|
||||
self.tools = tools or []
|
||||
self.name = name
|
||||
self.role = role
|
||||
self.write_events = write_events
|
||||
|
||||
if llm is None:
|
||||
llm = Settings.llm
|
||||
self.llm = llm
|
||||
assert self.llm.metadata.is_function_calling_model
|
||||
|
||||
self.system_prompt = system_prompt
|
||||
|
||||
self.memory = ChatMemoryBuffer.from_defaults(
|
||||
llm=self.llm, chat_history=chat_history
|
||||
)
|
||||
self.sources = []
|
||||
|
||||
@step()
|
||||
async def prepare_chat_history(self, ctx: Context, ev: StartEvent) -> InputEvent:
|
||||
# clear sources
|
||||
self.sources = []
|
||||
|
||||
# set system prompt
|
||||
if self.system_prompt is not None:
|
||||
system_msg = ChatMessage(role="system", content=self.system_prompt)
|
||||
self.memory.put(system_msg)
|
||||
|
||||
# set streaming
|
||||
ctx.data["streaming"] = getattr(ev, "streaming", False)
|
||||
|
||||
# get user input
|
||||
user_input = ev.input
|
||||
user_msg = ChatMessage(role="user", content=user_input)
|
||||
self.memory.put(user_msg)
|
||||
if self.write_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(name=self.name, msg=f"Start to work on: {user_input}")
|
||||
)
|
||||
|
||||
# get chat history
|
||||
chat_history = self.memory.get()
|
||||
return InputEvent(input=chat_history)
|
||||
|
||||
@step()
|
||||
async def handle_llm_input(
|
||||
self, ctx: Context, ev: InputEvent
|
||||
) -> ToolCallEvent | StopEvent:
|
||||
if ctx.data["streaming"]:
|
||||
return await self.handle_llm_input_stream(ctx, ev)
|
||||
|
||||
chat_history = ev.input
|
||||
|
||||
response = await self.llm.achat_with_tools(
|
||||
self.tools, chat_history=chat_history
|
||||
)
|
||||
self.memory.put(response.message)
|
||||
|
||||
tool_calls = self.llm.get_tool_calls_from_response(
|
||||
response, error_on_no_tool_call=False
|
||||
)
|
||||
|
||||
if not tool_calls:
|
||||
if self.write_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(name=self.name, msg="Finished task")
|
||||
)
|
||||
return StopEvent(
|
||||
result=AgentRunResult(response=response, sources=[*self.sources])
|
||||
)
|
||||
else:
|
||||
return ToolCallEvent(tool_calls=tool_calls)
|
||||
|
||||
async def handle_llm_input_stream(
|
||||
self, ctx: Context, ev: InputEvent
|
||||
) -> ToolCallEvent | StopEvent:
|
||||
chat_history = ev.input
|
||||
|
||||
async def response_generator() -> AsyncGenerator:
|
||||
response_stream = await self.llm.astream_chat_with_tools(
|
||||
self.tools, chat_history=chat_history
|
||||
)
|
||||
|
||||
full_response = None
|
||||
yielded_indicator = False
|
||||
async for chunk in response_stream:
|
||||
if "tool_calls" not in chunk.message.additional_kwargs:
|
||||
# Yield a boolean to indicate whether the response is a tool call
|
||||
if not yielded_indicator:
|
||||
yield False
|
||||
yielded_indicator = True
|
||||
|
||||
# if not a tool call, yield the chunks!
|
||||
yield chunk
|
||||
elif not yielded_indicator:
|
||||
# Yield the indicator for a tool call
|
||||
yield True
|
||||
yielded_indicator = True
|
||||
|
||||
full_response = chunk
|
||||
|
||||
# Write the full response to memory
|
||||
self.memory.put(full_response.message)
|
||||
|
||||
# Yield the final response
|
||||
yield full_response
|
||||
|
||||
# Start the generator
|
||||
generator = response_generator()
|
||||
|
||||
# Check for immediate tool call
|
||||
is_tool_call = await generator.__anext__()
|
||||
if is_tool_call:
|
||||
full_response = await generator.__anext__()
|
||||
tool_calls = self.llm.get_tool_calls_from_response(full_response)
|
||||
return ToolCallEvent(tool_calls=tool_calls)
|
||||
|
||||
# If we've reached here, it's not an immediate tool call, so we return the generator
|
||||
if self.write_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(name=self.name, msg="Finished task")
|
||||
)
|
||||
return StopEvent(result=generator)
|
||||
|
||||
@step()
|
||||
async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> InputEvent:
|
||||
tool_calls = ev.tool_calls
|
||||
tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools}
|
||||
|
||||
tool_msgs = []
|
||||
|
||||
# call tools -- safely!
|
||||
for tool_call in tool_calls:
|
||||
tool = tools_by_name.get(tool_call.tool_name)
|
||||
additional_kwargs = {
|
||||
"tool_call_id": tool_call.tool_id,
|
||||
"name": tool.metadata.get_name(),
|
||||
}
|
||||
if not tool:
|
||||
tool_msgs.append(
|
||||
ChatMessage(
|
||||
role="tool",
|
||||
content=f"Tool {tool_call.tool_name} does not exist",
|
||||
additional_kwargs=additional_kwargs,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
if isinstance(tool, ContextAwareTool):
|
||||
# inject context for calling an context aware tool
|
||||
tool_output = await tool.acall(ctx=ctx, **tool_call.tool_kwargs)
|
||||
else:
|
||||
tool_output = await tool.acall(**tool_call.tool_kwargs)
|
||||
self.sources.append(tool_output)
|
||||
tool_msgs.append(
|
||||
ChatMessage(
|
||||
role="tool",
|
||||
content=tool_output.content,
|
||||
additional_kwargs=additional_kwargs,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
tool_msgs.append(
|
||||
ChatMessage(
|
||||
role="tool",
|
||||
content=f"Encountered error in tool call: {e}",
|
||||
additional_kwargs=additional_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
for msg in tool_msgs:
|
||||
self.memory.put(msg)
|
||||
|
||||
chat_history = self.memory.get()
|
||||
return InputEvent(input=chat_history)
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
|
||||
from app.api.routers.models import (
|
||||
ChatData,
|
||||
)
|
||||
from app.api.routers.vercel_response import VercelStreamResponse
|
||||
from app.examples.factory import create_agent
|
||||
from fastapi import APIRouter, HTTPException, Request, status
|
||||
from llama_index.core.workflow import Workflow
|
||||
|
||||
chat_router = r = APIRouter()
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
@r.post("")
|
||||
async def chat(
|
||||
request: Request,
|
||||
data: ChatData,
|
||||
):
|
||||
try:
|
||||
last_message_content = data.get_last_message_content()
|
||||
messages = data.get_history_messages()
|
||||
# TODO: generate filters based on doc_ids
|
||||
# for now just use all documents
|
||||
# doc_ids = data.get_chat_document_ids()
|
||||
# TODO: use params
|
||||
# params = data.data or {}
|
||||
|
||||
agent: Workflow = create_agent(chat_history=messages)
|
||||
handler = agent.run(input=last_message_content, streaming=True)
|
||||
|
||||
return VercelStreamResponse(request, handler, agent.stream_events, data)
|
||||
except Exception as e:
|
||||
logger.exception("Error in agent", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Error in agent: {e}",
|
||||
) from e
|
||||
@@ -0,0 +1,48 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.api.routers.models import ChatConfig
|
||||
|
||||
|
||||
config_router = r = APIRouter()
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
@r.get("")
|
||||
async def chat_config() -> ChatConfig:
|
||||
starter_questions = None
|
||||
conversation_starters = os.getenv("CONVERSATION_STARTERS")
|
||||
if conversation_starters and conversation_starters.strip():
|
||||
starter_questions = conversation_starters.strip().split("\n")
|
||||
return ChatConfig(starter_questions=starter_questions)
|
||||
|
||||
|
||||
try:
|
||||
from app.engine.service import LLamaCloudFileService
|
||||
|
||||
logger.info("LlamaCloud is configured. Adding /config/llamacloud route.")
|
||||
|
||||
@r.get("/llamacloud")
|
||||
async def chat_llama_cloud_config():
|
||||
projects = LLamaCloudFileService.get_all_projects_with_pipelines()
|
||||
pipeline = os.getenv("LLAMA_CLOUD_INDEX_NAME")
|
||||
project = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
|
||||
pipeline_config = None
|
||||
if pipeline and project:
|
||||
pipeline_config = {
|
||||
"pipeline": pipeline,
|
||||
"project": project,
|
||||
}
|
||||
return {
|
||||
"projects": projects,
|
||||
"pipeline": pipeline_config,
|
||||
}
|
||||
|
||||
except ImportError:
|
||||
logger.debug(
|
||||
"LlamaCloud is not configured. Skipping adding /config/llamacloud route."
|
||||
)
|
||||
pass
|
||||
@@ -0,0 +1,227 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Literal, Optional
|
||||
|
||||
from llama_index.core.llms import ChatMessage, MessageRole
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
from pydantic import BaseModel, Field, validator
|
||||
from pydantic.alias_generators import to_camel
|
||||
|
||||
from app.config import DATA_DIR
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class FileContent(BaseModel):
|
||||
type: Literal["text", "ref"]
|
||||
# If the file is pure text then the value is be a string
|
||||
# otherwise, it's a list of document IDs
|
||||
value: str | List[str]
|
||||
|
||||
|
||||
class File(BaseModel):
|
||||
id: str
|
||||
content: FileContent
|
||||
filename: str
|
||||
filesize: int
|
||||
filetype: str
|
||||
|
||||
|
||||
class AnnotationFileData(BaseModel):
|
||||
files: List[File] = Field(
|
||||
default=[],
|
||||
description="List of files",
|
||||
)
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"csvFiles": [
|
||||
{
|
||||
"content": "Name, Age\nAlice, 25\nBob, 30",
|
||||
"filename": "example.csv",
|
||||
"filesize": 123,
|
||||
"id": "123",
|
||||
"type": "text/csv",
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
alias_generator = to_camel
|
||||
|
||||
|
||||
class Annotation(BaseModel):
|
||||
type: str
|
||||
data: AnnotationFileData | List[str]
|
||||
|
||||
def to_content(self) -> str | None:
|
||||
if self.type == "document_file":
|
||||
# We only support generating context content for CSV files for now
|
||||
csv_files = [file for file in self.data.files if file.filetype == "csv"]
|
||||
if len(csv_files) > 0:
|
||||
return "Use data from following CSV raw content\n" + "\n".join(
|
||||
[f"```csv\n{csv_file.content.value}\n```" for csv_file in csv_files]
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"The annotation {self.type} is not supported for generating context content"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
role: MessageRole
|
||||
content: str
|
||||
annotations: List[Annotation] | None = None
|
||||
|
||||
|
||||
class ChatData(BaseModel):
|
||||
messages: List[Message]
|
||||
data: Any = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What standards for letters exist?",
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@validator("messages")
|
||||
def messages_must_not_be_empty(cls, v):
|
||||
if len(v) == 0:
|
||||
raise ValueError("Messages must not be empty")
|
||||
return v
|
||||
|
||||
def get_last_message_content(self) -> str:
|
||||
"""
|
||||
Get the content of the last message along with the data content if available.
|
||||
Fallback to use data content from previous messages
|
||||
"""
|
||||
if len(self.messages) == 0:
|
||||
raise ValueError("There is not any message in the chat")
|
||||
last_message = self.messages[-1]
|
||||
message_content = last_message.content
|
||||
for message in reversed(self.messages):
|
||||
if message.role == MessageRole.USER and message.annotations is not None:
|
||||
annotation_contents = filter(
|
||||
None,
|
||||
[annotation.to_content() for annotation in message.annotations],
|
||||
)
|
||||
if not annotation_contents:
|
||||
continue
|
||||
annotation_text = "\n".join(annotation_contents)
|
||||
message_content = f"{message_content}\n{annotation_text}"
|
||||
break
|
||||
return message_content
|
||||
|
||||
def get_history_messages(self) -> List[ChatMessage]:
|
||||
"""
|
||||
Get the history messages
|
||||
"""
|
||||
return [
|
||||
ChatMessage(role=message.role, content=message.content)
|
||||
for message in self.messages[:-1]
|
||||
]
|
||||
|
||||
def is_last_message_from_user(self) -> bool:
|
||||
return self.messages[-1].role == MessageRole.USER
|
||||
|
||||
def get_chat_document_ids(self) -> List[str]:
|
||||
"""
|
||||
Get the document IDs from the chat messages
|
||||
"""
|
||||
document_ids: List[str] = []
|
||||
for message in self.messages:
|
||||
if message.role == MessageRole.USER and message.annotations is not None:
|
||||
for annotation in message.annotations:
|
||||
if (
|
||||
annotation.type == "document_file"
|
||||
and annotation.data.files is not None
|
||||
):
|
||||
for fi in annotation.data.files:
|
||||
if fi.content.type == "ref":
|
||||
document_ids += fi.content.value
|
||||
return list(set(document_ids))
|
||||
|
||||
|
||||
class SourceNodes(BaseModel):
|
||||
id: str
|
||||
metadata: Dict[str, Any]
|
||||
score: Optional[float]
|
||||
text: str
|
||||
url: Optional[str]
|
||||
|
||||
@classmethod
|
||||
def from_source_node(cls, source_node: NodeWithScore):
|
||||
metadata = source_node.node.metadata
|
||||
url = cls.get_url_from_metadata(metadata)
|
||||
|
||||
return cls(
|
||||
id=source_node.node.node_id,
|
||||
metadata=metadata,
|
||||
score=source_node.score,
|
||||
text=source_node.node.text, # type: ignore
|
||||
url=url,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_url_from_metadata(cls, metadata: Dict[str, Any]) -> str:
|
||||
url_prefix = os.getenv("FILESERVER_URL_PREFIX")
|
||||
if not url_prefix:
|
||||
logger.warning(
|
||||
"Warning: FILESERVER_URL_PREFIX not set in environment variables. Can't use file server"
|
||||
)
|
||||
file_name = metadata.get("file_name")
|
||||
|
||||
if file_name and url_prefix:
|
||||
# file_name exists and file server is configured
|
||||
pipeline_id = metadata.get("pipeline_id")
|
||||
if pipeline_id:
|
||||
# file is from LlamaCloud
|
||||
file_name = f"{pipeline_id}${file_name}"
|
||||
return f"{url_prefix}/output/llamacloud/{file_name}"
|
||||
is_private = metadata.get("private", "false") == "true"
|
||||
if is_private:
|
||||
# file is a private upload
|
||||
return f"{url_prefix}/output/uploaded/{file_name}"
|
||||
# file is from calling the 'generate' script
|
||||
# Get the relative path of file_path to data_dir
|
||||
file_path = metadata.get("file_path")
|
||||
data_dir = os.path.abspath(DATA_DIR)
|
||||
if file_path and data_dir:
|
||||
relative_path = os.path.relpath(file_path, data_dir)
|
||||
return f"{url_prefix}/data/{relative_path}"
|
||||
# fallback to URL in metadata (e.g. for websites)
|
||||
return metadata.get("URL")
|
||||
|
||||
@classmethod
|
||||
def from_source_nodes(cls, source_nodes: List[NodeWithScore]):
|
||||
return [cls.from_source_node(node) for node in source_nodes]
|
||||
|
||||
|
||||
class Result(BaseModel):
|
||||
result: Message
|
||||
nodes: List[SourceNodes]
|
||||
|
||||
|
||||
class ChatConfig(BaseModel):
|
||||
starter_questions: Optional[List[str]] = Field(
|
||||
default=None,
|
||||
description="List of starter questions",
|
||||
serialization_alias="starterQuestions",
|
||||
)
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"starterQuestions": [
|
||||
"What standards for letters exist?",
|
||||
"What are the requirements for a letter to be considered a letter?",
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
import logging
|
||||
from typing import List, Any
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.api.services.file import PrivateFileService
|
||||
|
||||
file_upload_router = r = APIRouter()
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class FileUploadRequest(BaseModel):
|
||||
base64: str
|
||||
filename: str
|
||||
params: Any = None
|
||||
|
||||
|
||||
@r.post("")
|
||||
def upload_file(request: FileUploadRequest) -> List[str]:
|
||||
try:
|
||||
logger.info("Processing file")
|
||||
return PrivateFileService.process_file(
|
||||
request.filename, request.base64, request.params
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing file: {e}", exc_info=True)
|
||||
raise HTTPException(status_code=500, detail="Error processing file")
|
||||
@@ -0,0 +1,121 @@
|
||||
import json
|
||||
import logging
|
||||
from asyncio import Task
|
||||
from typing import AsyncGenerator, List
|
||||
|
||||
from aiostream import stream
|
||||
from app.agents.single import AgentRunEvent, AgentRunResult
|
||||
from app.api.routers.models import ChatData, Message
|
||||
from app.api.services.suggestion import NextQuestionSuggestion
|
||||
from fastapi import Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class VercelStreamResponse(StreamingResponse):
|
||||
"""
|
||||
Class to convert the response from the chat engine to the streaming format expected by Vercel
|
||||
"""
|
||||
|
||||
TEXT_PREFIX = "0:"
|
||||
DATA_PREFIX = "8:"
|
||||
|
||||
@classmethod
|
||||
def convert_text(cls, token: str):
|
||||
# Escape newlines and double quotes to avoid breaking the stream
|
||||
token = json.dumps(token)
|
||||
return f"{cls.TEXT_PREFIX}{token}\n"
|
||||
|
||||
@classmethod
|
||||
def convert_data(cls, data: dict):
|
||||
data_str = json.dumps(data)
|
||||
return f"{cls.DATA_PREFIX}[{data_str}]\n"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request: Request,
|
||||
task: Task[AgentRunResult | AsyncGenerator],
|
||||
events: AsyncGenerator[AgentRunEvent, None],
|
||||
chat_data: ChatData,
|
||||
verbose: bool = True,
|
||||
):
|
||||
content = VercelStreamResponse.content_generator(
|
||||
request, task, events, chat_data, verbose
|
||||
)
|
||||
super().__init__(content=content)
|
||||
|
||||
@classmethod
|
||||
async def content_generator(
|
||||
cls,
|
||||
request: Request,
|
||||
task: Task[AgentRunResult | AsyncGenerator],
|
||||
events: AsyncGenerator[AgentRunEvent, None],
|
||||
chat_data: ChatData,
|
||||
verbose: bool = True,
|
||||
):
|
||||
# Yield the text response
|
||||
async def _chat_response_generator():
|
||||
result = await task
|
||||
final_response = ""
|
||||
|
||||
if isinstance(result, AgentRunResult):
|
||||
for token in result.response.message.content:
|
||||
final_response += token
|
||||
yield cls.convert_text(token)
|
||||
|
||||
if isinstance(result, AsyncGenerator):
|
||||
async for token in result:
|
||||
final_response += token.delta
|
||||
yield cls.convert_text(token.delta)
|
||||
|
||||
# Generate next questions if next question prompt is configured
|
||||
question_data = await cls._generate_next_questions(
|
||||
chat_data.messages, final_response
|
||||
)
|
||||
if question_data:
|
||||
yield cls.convert_data(question_data)
|
||||
|
||||
# TODO: stream sources
|
||||
|
||||
# Yield the events from the event handler
|
||||
async def _event_generator():
|
||||
async for event in events():
|
||||
event_response = cls._event_to_response(event)
|
||||
if verbose:
|
||||
logger.debug(event_response)
|
||||
if event_response is not None:
|
||||
yield cls.convert_data(event_response)
|
||||
|
||||
combine = stream.merge(_chat_response_generator(), _event_generator())
|
||||
|
||||
is_stream_started = False
|
||||
async with combine.stream() as streamer:
|
||||
if not is_stream_started:
|
||||
is_stream_started = True
|
||||
# Stream a blank message to start the stream
|
||||
yield cls.convert_text("")
|
||||
|
||||
async for output in streamer:
|
||||
yield output
|
||||
if await request.is_disconnected():
|
||||
break
|
||||
|
||||
@staticmethod
|
||||
def _event_to_response(event: AgentRunEvent) -> dict:
|
||||
return {
|
||||
"type": "agent",
|
||||
"data": {"agent": event.name, "text": event.msg},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def _generate_next_questions(chat_history: List[Message], response: str):
|
||||
questions = await NextQuestionSuggestion.suggest_next_questions(
|
||||
chat_history, response
|
||||
)
|
||||
if questions:
|
||||
return {
|
||||
"type": "suggested_questions",
|
||||
"data": questions,
|
||||
}
|
||||
return None
|
||||
@@ -0,0 +1 @@
|
||||
DATA_DIR = "data"
|
||||
@@ -1,19 +0,0 @@
|
||||
from llama_index.llms.openai import OpenAI
|
||||
from llama_agents import AgentOrchestrator, ControlPlaneServer
|
||||
from app.core.message_queue import message_queue
|
||||
from app.utils import load_from_env
|
||||
|
||||
|
||||
control_plane_host = (
|
||||
load_from_env("CONTROL_PLANE_HOST", throw_error=False) or "127.0.0.1"
|
||||
)
|
||||
control_plane_port = load_from_env("CONTROL_PLANE_PORT", throw_error=False) or "8001"
|
||||
|
||||
|
||||
# setup control plane
|
||||
control_plane = ControlPlaneServer(
|
||||
message_queue=message_queue,
|
||||
orchestrator=AgentOrchestrator(llm=OpenAI()),
|
||||
host=control_plane_host,
|
||||
port=int(control_plane_port) if control_plane_port else None,
|
||||
)
|
||||
@@ -1,12 +0,0 @@
|
||||
from llama_agents import SimpleMessageQueue
|
||||
from app.utils import load_from_env
|
||||
|
||||
message_queue_host = (
|
||||
load_from_env("MESSAGE_QUEUE_HOST", throw_error=False) or "127.0.0.1"
|
||||
)
|
||||
message_queue_port = load_from_env("MESSAGE_QUEUE_PORT", throw_error=False) or "8000"
|
||||
|
||||
message_queue = SimpleMessageQueue(
|
||||
host=message_queue_host,
|
||||
port=int(message_queue_port) if message_queue_port else None,
|
||||
)
|
||||
@@ -1,88 +0,0 @@
|
||||
import json
|
||||
from logging import getLogger
|
||||
from pathlib import Path
|
||||
from fastapi import FastAPI
|
||||
from typing import Dict, Optional
|
||||
from llama_agents import CallableMessageConsumer, QueueMessage
|
||||
from llama_agents.message_queues.base import BaseMessageQueue
|
||||
from llama_agents.message_consumers.base import BaseMessageQueueConsumer
|
||||
from llama_agents.message_consumers.remote import RemoteMessageConsumer
|
||||
from app.utils import load_from_env
|
||||
from app.core.message_queue import message_queue
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
class TaskResultService:
|
||||
def __init__(
|
||||
self,
|
||||
message_queue: BaseMessageQueue,
|
||||
name: str = "human",
|
||||
host: str = "127.0.0.1",
|
||||
port: Optional[int] = 8002,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.host = host
|
||||
self.port = port
|
||||
|
||||
self._message_queue = message_queue
|
||||
|
||||
# app
|
||||
self._app = FastAPI()
|
||||
self._app.add_api_route(
|
||||
"/", self.home, methods=["GET"], tags=["Human Consumer"]
|
||||
)
|
||||
self._app.add_api_route(
|
||||
"/process_message",
|
||||
self.process_message,
|
||||
methods=["POST"],
|
||||
tags=["Human Consumer"],
|
||||
)
|
||||
|
||||
@property
|
||||
def message_queue(self) -> BaseMessageQueue:
|
||||
return self._message_queue
|
||||
|
||||
def as_consumer(self, remote: bool = False) -> BaseMessageQueueConsumer:
|
||||
if remote:
|
||||
return RemoteMessageConsumer(
|
||||
url=(
|
||||
f"http://{self.host}:{self.port}/process_message"
|
||||
if self.port
|
||||
else f"http://{self.host}/process_message"
|
||||
),
|
||||
message_type=self.name,
|
||||
)
|
||||
|
||||
return CallableMessageConsumer(
|
||||
message_type=self.name,
|
||||
handler=self.process_message,
|
||||
)
|
||||
|
||||
async def process_message(self, message: QueueMessage) -> None:
|
||||
Path("task_results").mkdir(exist_ok=True)
|
||||
with open("task_results/task_results.json", "+a") as f:
|
||||
json.dump(message.model_dump(), f)
|
||||
f.write("\n")
|
||||
|
||||
async def home(self) -> Dict[str, str]:
|
||||
return {"message": "hello, human."}
|
||||
|
||||
async def register_to_message_queue(self) -> None:
|
||||
"""Register to the message queue."""
|
||||
await self.message_queue.register_consumer(self.as_consumer(remote=True))
|
||||
|
||||
|
||||
human_consumer_host = (
|
||||
load_from_env("HUMAN_CONSUMER_HOST", throw_error=False) or "127.0.0.1"
|
||||
)
|
||||
human_consumer_port = load_from_env("HUMAN_CONSUMER_PORT", throw_error=False) or "8002"
|
||||
|
||||
|
||||
human_consumer_server = TaskResultService(
|
||||
message_queue=message_queue,
|
||||
host=human_consumer_host,
|
||||
port=int(human_consumer_port) if human_consumer_port else None,
|
||||
name="human",
|
||||
)
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import List, Optional
|
||||
from app.agents.single import FunctionCallingAgent
|
||||
from app.agents.multi import AgentCallingAgent
|
||||
from app.examples.researcher import create_researcher
|
||||
from llama_index.core.chat_engine.types import ChatMessage
|
||||
|
||||
|
||||
def create_choreography(chat_history: Optional[List[ChatMessage]] = None):
|
||||
researcher = create_researcher(chat_history)
|
||||
reviewer = FunctionCallingAgent(
|
||||
name="reviewer",
|
||||
role="expert in reviewing blog posts",
|
||||
system_prompt="You are an expert in reviewing blog posts. You are given a task to review a blog post. Review the post for logical inconsistencies, ask critical questions, and provide suggestions for improvement. Furthermore, proofread the post for grammar and spelling errors. If the post is good, you can say 'The post is good.'",
|
||||
chat_history=chat_history,
|
||||
)
|
||||
return AgentCallingAgent(
|
||||
name="writer",
|
||||
agents=[researcher, reviewer],
|
||||
role="expert in writing blog posts",
|
||||
system_prompt="""You are an expert in writing blog posts. You are given a task to write a blog post. Before starting to write the post, consult the researcher agent to get the information you need. Don't make up any information yourself.
|
||||
After creating a draft for the post, send it to the reviewer agent to receive some feedback and make sure to incorporate the feedback from the reviewer.
|
||||
You can consult the reviewer and researcher maximal two times. Your output should just contain the blog post.""",
|
||||
# TODO: add chat_history support to AgentCallingAgent
|
||||
# chat_history=chat_history,
|
||||
)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user