mirror of
https://github.com/run-llama/create-llama.git
synced 2026-07-16 11:04:26 -04:00
Compare commits
55 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 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 | |||
| 81ef7f0f93 | |||
| 8faf9170cf | |||
| c49a5e1620 | |||
| 8b2de431f2 | |||
| d746c75e49 | |||
| c87978ab96 |
@@ -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"
|
||||
|
||||
+127
@@ -1,5 +1,132 @@
|
||||
# create-llama
|
||||
|
||||
## 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
|
||||
|
||||
- c49a5e1: Add error handling for generating the next question
|
||||
- c49a5e1: Fix wrong api key variable in Azure OpenAI provider
|
||||
|
||||
## 0.1.33
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- d746c75: Add Weaviate vector store (Typescript)
|
||||
|
||||
## 0.1.32
|
||||
|
||||
### 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();
|
||||
});
|
||||
});
|
||||
@@ -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,
|
||||
);
|
||||
+15
-9
@@ -24,14 +24,14 @@ 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,
|
||||
): Promise<CreateLlamaResult> {
|
||||
if (!process.env.OPENAI_API_KEY || !process.env.LLAMA_CLOUD_API_KEY) {
|
||||
throw new Error(
|
||||
@@ -45,7 +45,7 @@ export async function runCreateLlama(
|
||||
templateUI,
|
||||
appType,
|
||||
].join("-");
|
||||
const command = [
|
||||
const commandArgs = [
|
||||
"create-llama",
|
||||
name,
|
||||
"--template",
|
||||
@@ -53,13 +53,10 @@ export async function runCreateLlama(
|
||||
"--framework",
|
||||
templateFramework,
|
||||
dataSource,
|
||||
"--ui",
|
||||
templateUI,
|
||||
"--vector-db",
|
||||
vectorDb,
|
||||
"--open-ai-key",
|
||||
process.env.OPENAI_API_KEY,
|
||||
appType,
|
||||
"--use-pnpm",
|
||||
"--port",
|
||||
port,
|
||||
@@ -74,7 +71,16 @@ export async function runCreateLlama(
|
||||
"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],
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
+70
-36
@@ -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)
|
||||
@@ -311,7 +312,7 @@ const getModelEnvs = (modelConfig: ModelConfig): EnvVar[] => {
|
||||
...(modelConfig.provider === "azure-openai"
|
||||
? [
|
||||
{
|
||||
name: "AZURE_OPENAI_KEY",
|
||||
name: "AZURE_OPENAI_API_KEY",
|
||||
description: "The Azure OpenAI key to use.",
|
||||
value: modelConfig.apiKey,
|
||||
},
|
||||
@@ -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);
|
||||
|
||||
+28
-24
@@ -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
|
||||
@@ -142,12 +143,15 @@ export const installTemplate = async (
|
||||
|
||||
if (props.framework === "fastapi") {
|
||||
await installPythonTemplate(props);
|
||||
// write loaders configuration (currently Python only)
|
||||
await writeLoadersConfig(
|
||||
props.root,
|
||||
props.dataSources,
|
||||
props.useLlamaParse,
|
||||
);
|
||||
if (props.vectorDb !== "llamacloud") {
|
||||
// write loaders configuration (currently Python only)
|
||||
// not needed for LlamaCloud as it has its own loaders
|
||||
await writeLoadersConfig(
|
||||
props.root,
|
||||
props.dataSources,
|
||||
props.useLlamaParse,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
await installTSTemplate(props);
|
||||
}
|
||||
@@ -171,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
|
||||
|
||||
@@ -9,6 +9,7 @@ const ALL_AZURE_OPENAI_CHAT_MODELS: Record<string, { openAIModel: string }> = {
|
||||
openAIModel: "gpt-3.5-turbo-16k",
|
||||
},
|
||||
"gpt-4o": { openAIModel: "gpt-4o" },
|
||||
"gpt-4o-mini": { openAIModel: "gpt-4o-mini" },
|
||||
"gpt-4": { openAIModel: "gpt-4" },
|
||||
"gpt-4-32k": { openAIModel: "gpt-4-32k" },
|
||||
"gpt-4-turbo": {
|
||||
@@ -26,6 +27,9 @@ const ALL_AZURE_OPENAI_CHAT_MODELS: Record<string, { openAIModel: string }> = {
|
||||
"gpt-4o-2024-05-13": {
|
||||
openAIModel: "gpt-4o-2024-05-13",
|
||||
},
|
||||
"gpt-4o-mini-2024-07-18": {
|
||||
openAIModel: "gpt-4o-mini-2024-07-18",
|
||||
},
|
||||
};
|
||||
|
||||
const ALL_AZURE_OPENAI_EMBEDDING_MODELS: Record<
|
||||
@@ -35,10 +39,6 @@ const ALL_AZURE_OPENAI_EMBEDDING_MODELS: Record<
|
||||
openAIModel: string;
|
||||
}
|
||||
> = {
|
||||
"text-embedding-ada-002": {
|
||||
dimensions: 1536,
|
||||
openAIModel: "text-embedding-ada-002",
|
||||
},
|
||||
"text-embedding-3-small": {
|
||||
dimensions: 1536,
|
||||
openAIModel: "text-embedding-3-small",
|
||||
|
||||
@@ -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,
|
||||
|
||||
+70
-29
@@ -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[] = [];
|
||||
|
||||
@@ -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.0",
|
||||
});
|
||||
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;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -350,18 +380,28 @@ export const installPythonTemplate = async ({
|
||||
cwd: path.join(compPath, "vectordbs", "python", vectorDb ?? "none"),
|
||||
});
|
||||
|
||||
// Copy all loaders to enginePath
|
||||
const loaderPath = path.join(enginePath, "loaders");
|
||||
await copy("**", loaderPath, {
|
||||
parents: true,
|
||||
cwd: path.join(compPath, "loaders", "python"),
|
||||
});
|
||||
if (vectorDb !== "llamacloud") {
|
||||
// Copy all loaders to enginePath
|
||||
// Not needed for LlamaCloud as it has its own loaders
|
||||
const loaderPath = path.join(enginePath, "loaders");
|
||||
await copy("**", loaderPath, {
|
||||
parents: true,
|
||||
cwd: path.join(compPath, "loaders", "python"),
|
||||
});
|
||||
}
|
||||
|
||||
// Copy settings.py to app
|
||||
await copy("**", path.join(root, "app"), {
|
||||
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
|
||||
@@ -385,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, {
|
||||
|
||||
@@ -173,7 +173,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 +195,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")) {
|
||||
@@ -341,6 +352,7 @@ Please check ${cyan(
|
||||
console.log(`Running app in ${root}...`);
|
||||
await runApp(
|
||||
root,
|
||||
program.template,
|
||||
program.frontend,
|
||||
program.framework,
|
||||
program.port,
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "create-llama",
|
||||
"version": "0.1.32",
|
||||
"version": "0.2.7",
|
||||
"description": "Create LlamaIndex-powered apps with one command",
|
||||
"keywords": [
|
||||
"rag",
|
||||
@@ -9,7 +9,7 @@
|
||||
],
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "https://github.com/run-llama/LlamaIndexTS",
|
||||
"url": "https://github.com/run-llama/create-llama",
|
||||
"directory": "packages/create-llama"
|
||||
},
|
||||
"license": "MIT",
|
||||
|
||||
+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,40 +1,28 @@
|
||||
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 });
|
||||
const questions = extractQuestions(response.text);
|
||||
return questions;
|
||||
} catch (error) {
|
||||
console.error("Error: ", error);
|
||||
throw error;
|
||||
console.error("Error when generating the next questions: ", error);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -10,7 +10,15 @@ export function getExtractors() {
|
||||
}
|
||||
|
||||
export async function getDocuments() {
|
||||
return await new SimpleDirectoryReader().loadData({
|
||||
const documents = await new SimpleDirectoryReader().loadData({
|
||||
directoryPath: DATA_DIR,
|
||||
});
|
||||
// Set private=false to mark the document as public (required for filtering)
|
||||
for (const document of documents) {
|
||||
document.metadata = {
|
||||
...document.metadata,
|
||||
private: "false",
|
||||
};
|
||||
}
|
||||
return documents;
|
||||
}
|
||||
|
||||
@@ -23,8 +23,16 @@ export function getExtractors() {
|
||||
export async function getDocuments() {
|
||||
const reader = new SimpleDirectoryReader();
|
||||
const extractors = getExtractors();
|
||||
return await reader.loadData({
|
||||
const documents = await reader.loadData({
|
||||
directoryPath: DATA_DIR,
|
||||
fileExtToReader: extractors,
|
||||
});
|
||||
// Set private=false to mark the document as public (required for filtering)
|
||||
for (const document of documents) {
|
||||
document.metadata = {
|
||||
...document.metadata,
|
||||
private: "false",
|
||||
};
|
||||
}
|
||||
return documents;
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ authors = ["Marcus Schiesser <mail@marcusschiesser.de>"]
|
||||
readme = "README.md"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.11,<3.12"
|
||||
python = "^3.11,<4.0"
|
||||
llama-index = "^0.10.6"
|
||||
llama-index-readers-file = "^0.1.3"
|
||||
python-dotenv = "^1.0.0"
|
||||
|
||||
+36
-25
@@ -1,19 +1,16 @@
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Dict, List
|
||||
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
|
||||
@@ -41,7 +38,7 @@ class PrivateFileService:
|
||||
PRIVATE_STORE_PATH = "output/uploaded"
|
||||
|
||||
@staticmethod
|
||||
def preprocess_base64_file(base64_content: str) -> tuple:
|
||||
def preprocess_base64_file(base64_content: str) -> Tuple[bytes, str | None]:
|
||||
header, data = base64_content.split(",", 1)
|
||||
mime_type = header.split(";")[0].split(":", 1)[1]
|
||||
extension = mimetypes.guess_extension(mime_type)
|
||||
@@ -49,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
|
||||
@@ -78,25 +72,42 @@ class PrivateFileService:
|
||||
return documents
|
||||
|
||||
@staticmethod
|
||||
def process_file(base64_content: str) -> List[str]:
|
||||
file_data, extension = PrivateFileService.preprocess_base64_file(base64_content)
|
||||
documents = PrivateFileService.store_and_parse_file(file_data, extension)
|
||||
def process_file(
|
||||
file_name: str, base64_content: str, params: Optional[dict] = None
|
||||
) -> List[str]:
|
||||
if params is None:
|
||||
params = {}
|
||||
|
||||
# Only process nodes, no store the index
|
||||
pipeline = IngestionPipeline()
|
||||
nodes = pipeline.run(documents=documents)
|
||||
file_data, extension = PrivateFileService.preprocess_base64_file(base64_content)
|
||||
|
||||
# Add the nodes to the index and persist it
|
||||
current_index = get_index()
|
||||
index_config = IndexConfig(**params)
|
||||
current_index = get_index(index_config)
|
||||
|
||||
# Insert the documents into the index
|
||||
if isinstance(current_index, LlamaCloudIndex):
|
||||
# LlamaCloudIndex is a managed index so we don't need to process the nodes
|
||||
# just insert the documents
|
||||
for doc in documents:
|
||||
current_index.insert(doc)
|
||||
from app.engine.service import LLamaCloudFileService
|
||||
|
||||
project_id = current_index._get_project_id()
|
||||
pipeline_id = current_index._get_pipeline_id()
|
||||
# LlamaCloudIndex is a managed index so we can directly use the files
|
||||
upload_file = (file_name, BytesIO(file_data))
|
||||
return [
|
||||
LLamaCloudFileService.add_file_to_pipeline(
|
||||
project_id,
|
||||
pipeline_id,
|
||||
upload_file,
|
||||
custom_metadata={
|
||||
# Set private=true to mark the document as private user docs (required for filtering)
|
||||
"private": "true",
|
||||
},
|
||||
)
|
||||
]
|
||||
else:
|
||||
# Only process nodes, no store the index
|
||||
# First process documents into nodes
|
||||
documents = PrivateFileService.store_and_parse_file(
|
||||
file_name, file_data, extension
|
||||
)
|
||||
pipeline = IngestionPipeline()
|
||||
nodes = pipeline.run(documents=documents)
|
||||
|
||||
@@ -109,5 +120,5 @@ class PrivateFileService:
|
||||
persist_dir=os.environ.get("STORAGE_DIR", "storage")
|
||||
)
|
||||
|
||||
# Return the document ids
|
||||
return [doc.doc_id for doc in documents]
|
||||
# Return the document ids
|
||||
return [doc.doc_id for doc in 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,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -82,7 +82,7 @@ def init_azure_openai():
|
||||
dimensions = os.getenv("EMBEDDING_DIM")
|
||||
|
||||
azure_config = {
|
||||
"api_key": os.environ["AZURE_OPENAI_KEY"],
|
||||
"api_key": os.environ["AZURE_OPENAI_API_KEY"],
|
||||
"azure_endpoint": os.environ["AZURE_OPENAI_ENDPOINT"],
|
||||
"api_version": os.getenv("AZURE_OPENAI_API_VERSION")
|
||||
or os.getenv("OPENAI_API_VERSION"),
|
||||
@@ -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,48 +1,47 @@
|
||||
# flake8: noqa: E402
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from app.engine.index import get_index
|
||||
|
||||
load_dotenv()
|
||||
|
||||
import os
|
||||
import logging
|
||||
from app.settings import init_settings
|
||||
from app.engine.loaders import get_documents
|
||||
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
|
||||
|
||||
from llama_index.core.readers import SimpleDirectoryReader
|
||||
from app.engine.service import LLamaCloudFileService
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger()
|
||||
|
||||
|
||||
def generate_datasource():
|
||||
init_settings()
|
||||
logger.info("Generate index for the provided data")
|
||||
|
||||
name = os.getenv("LLAMA_CLOUD_INDEX_NAME")
|
||||
project_name = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
|
||||
api_key = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
organization_id = os.getenv("LLAMA_CLOUD_ORGANIZATION_ID")
|
||||
index = get_index()
|
||||
project_id = index._get_project_id()
|
||||
pipeline_id = index._get_pipeline_id()
|
||||
|
||||
if name is None or project_name is None or api_key is None:
|
||||
raise ValueError(
|
||||
"Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
|
||||
" to your environment variables or config them in .env file"
|
||||
)
|
||||
|
||||
documents = get_documents()
|
||||
|
||||
# Set private=false to mark the document as public (required for filtering)
|
||||
for doc in documents:
|
||||
doc.metadata["private"] = "false"
|
||||
|
||||
LlamaCloudIndex.from_documents(
|
||||
documents=documents,
|
||||
name=name,
|
||||
project_name=project_name,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
organization_id=organization_id
|
||||
# use SimpleDirectoryReader to retrieve the files to process
|
||||
reader = SimpleDirectoryReader(
|
||||
"data",
|
||||
recursive=True,
|
||||
)
|
||||
files_to_process = reader.input_files
|
||||
|
||||
# add each file to the LlamaCloud pipeline
|
||||
for input_file in files_to_process:
|
||||
with open(input_file, "rb") as f:
|
||||
logger.info(
|
||||
f"Adding file {input_file} to pipeline {index.name} in project {index.project_name}"
|
||||
)
|
||||
LLamaCloudFileService.add_file_to_pipeline(
|
||||
project_id,
|
||||
pipeline_id,
|
||||
f,
|
||||
custom_metadata={
|
||||
# Set private=false to mark the document as public (required for filtering)
|
||||
"private": "false",
|
||||
},
|
||||
)
|
||||
|
||||
logger.info("Finished generating the index")
|
||||
|
||||
|
||||
@@ -1,31 +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_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")
|
||||
|
||||
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 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 to_client_kwargs(self) -> dict:
|
||||
return {
|
||||
"api_key": self.api_key,
|
||||
"base_url": self.base_url,
|
||||
}
|
||||
|
||||
|
||||
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,14 +5,14 @@ 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="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=doc_ids,
|
||||
operator="in", # type: ignore
|
||||
)
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
from io import BytesIO
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
||||
import typing
|
||||
|
||||
from fastapi import BackgroundTasks
|
||||
from llama_cloud import ManagedIngestionStatus, PipelineFileCreateCustomMetadataValue
|
||||
from pydantic import BaseModel
|
||||
import requests
|
||||
from app.api.routers.models import SourceNodes
|
||||
from app.engine.index import get_client
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class LlamaCloudFile(BaseModel):
|
||||
file_name: str
|
||||
pipeline_id: str
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, LlamaCloudFile):
|
||||
return NotImplemented
|
||||
return (
|
||||
self.file_name == other.file_name and self.pipeline_id == other.pipeline_id
|
||||
)
|
||||
|
||||
def __hash__(self):
|
||||
return hash((self.file_name, self.pipeline_id))
|
||||
|
||||
|
||||
class LLamaCloudFileService:
|
||||
LOCAL_STORE_PATH = "output/llamacloud"
|
||||
DOWNLOAD_FILE_NAME_TPL = "{pipeline_id}${filename}"
|
||||
|
||||
@classmethod
|
||||
def get_all_projects_with_pipelines(cls) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
client = get_client()
|
||||
projects = client.projects.list_projects()
|
||||
pipelines = client.pipelines.search_pipelines()
|
||||
return [
|
||||
{
|
||||
**(project.dict()),
|
||||
"pipelines": [
|
||||
{"id": p.id, "name": p.name}
|
||||
for p in pipelines
|
||||
if p.project_id == project.id
|
||||
],
|
||||
}
|
||||
for project in projects
|
||||
]
|
||||
except Exception as error:
|
||||
logger.error(f"Error listing projects and pipelines: {error}")
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def add_file_to_pipeline(
|
||||
cls,
|
||||
project_id: str,
|
||||
pipeline_id: str,
|
||||
upload_file: Union[typing.IO, Tuple[str, BytesIO]],
|
||||
custom_metadata: Optional[Dict[str, PipelineFileCreateCustomMetadataValue]],
|
||||
) -> str:
|
||||
client = get_client()
|
||||
file = client.files.upload_file(project_id=project_id, upload_file=upload_file)
|
||||
files = [
|
||||
{
|
||||
"file_id": file.id,
|
||||
"custom_metadata": {"file_id": file.id, **(custom_metadata or {})},
|
||||
}
|
||||
]
|
||||
files = client.pipelines.add_files_to_pipeline(pipeline_id, request=files)
|
||||
|
||||
# Wait 2s for the file to be processed
|
||||
max_attempts = 20
|
||||
attempt = 0
|
||||
while attempt < max_attempts:
|
||||
result = client.pipelines.get_pipeline_file_status(pipeline_id, file.id)
|
||||
if result.status == ManagedIngestionStatus.ERROR:
|
||||
raise Exception(f"File processing failed: {str(result)}")
|
||||
if result.status == ManagedIngestionStatus.SUCCESS:
|
||||
# File is ingested - return the file id
|
||||
return file.id
|
||||
attempt += 1
|
||||
time.sleep(0.1) # Sleep for 100ms
|
||||
raise Exception(
|
||||
f"File processing did not complete after {max_attempts} attempts."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def download_pipeline_file(
|
||||
cls,
|
||||
file: LlamaCloudFile,
|
||||
force_download: bool = False,
|
||||
):
|
||||
client = get_client()
|
||||
file_name = file.file_name
|
||||
pipeline_id = file.pipeline_id
|
||||
|
||||
# Check is the file already exists
|
||||
downloaded_file_path = cls._get_file_path(file_name, pipeline_id)
|
||||
if os.path.exists(downloaded_file_path) and not force_download:
|
||||
logger.debug(f"File {file_name} already exists in local storage")
|
||||
return
|
||||
try:
|
||||
logger.info(f"Downloading file {file_name} for pipeline {pipeline_id}")
|
||||
files = client.pipelines.list_pipeline_files(pipeline_id)
|
||||
if not files or not isinstance(files, list):
|
||||
raise Exception("No files found in LlamaCloud")
|
||||
for file_entry in files:
|
||||
if file_entry.name == file_name:
|
||||
file_id = file_entry.file_id
|
||||
project_id = file_entry.project_id
|
||||
file_detail = client.files.read_file_content(
|
||||
file_id, project_id=project_id
|
||||
)
|
||||
cls._download_file(file_detail.url, downloaded_file_path)
|
||||
break
|
||||
except Exception as error:
|
||||
logger.info(f"Error fetching file from LlamaCloud: {error}")
|
||||
|
||||
@classmethod
|
||||
def download_files_from_nodes(
|
||||
cls, nodes: List[NodeWithScore], background_tasks: BackgroundTasks
|
||||
):
|
||||
files = cls._get_files_to_download(nodes)
|
||||
for file in files:
|
||||
logger.info(f"Adding download of {file.file_name} to background tasks")
|
||||
background_tasks.add_task(
|
||||
LLamaCloudFileService.download_pipeline_file, file
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_files_to_download(cls, nodes: List[NodeWithScore]) -> Set[LlamaCloudFile]:
|
||||
source_nodes = SourceNodes.from_source_nodes(nodes)
|
||||
llama_cloud_files = [
|
||||
LlamaCloudFile(
|
||||
file_name=node.metadata.get("file_name"),
|
||||
pipeline_id=node.metadata.get("pipeline_id"),
|
||||
)
|
||||
for node in source_nodes
|
||||
if (
|
||||
node.metadata.get("pipeline_id") is not None
|
||||
and node.metadata.get("file_name") is not None
|
||||
)
|
||||
]
|
||||
# Remove duplicates and return
|
||||
return set(llama_cloud_files)
|
||||
|
||||
@classmethod
|
||||
def _get_file_name(cls, name: str, pipeline_id: str) -> str:
|
||||
return cls.DOWNLOAD_FILE_NAME_TPL.format(pipeline_id=pipeline_id, filename=name)
|
||||
|
||||
@classmethod
|
||||
def _get_file_path(cls, name: str, pipeline_id: str) -> str:
|
||||
return os.path.join(cls.LOCAL_STORE_PATH, cls._get_file_name(name, pipeline_id))
|
||||
|
||||
@classmethod
|
||||
def _download_file(cls, url: str, local_file_path: str):
|
||||
logger.info(f"Saving file to {local_file_path}")
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(cls.LOCAL_STORE_PATH, exist_ok=True)
|
||||
# Download the file
|
||||
with requests.get(url, stream=True) as r:
|
||||
r.raise_for_status()
|
||||
with open(local_file_path, "wb") as f:
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
logger.info("File downloaded successfully")
|
||||
@@ -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,30 +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();
|
||||
// Set private=false to mark the document as public (required for filtering)
|
||||
for (const document of documents) {
|
||||
document.metadata = {
|
||||
...document.metadata,
|
||||
private: "false",
|
||||
};
|
||||
async function* walk(dir: string): AsyncGenerator<string> {
|
||||
const directory = await fs.opendir(dir);
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
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!`);
|
||||
}
|
||||
|
||||
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];
|
||||
|
||||
@@ -25,10 +25,6 @@ async function generateDatasource() {
|
||||
persistDir: STORAGE_CACHE_DIR,
|
||||
});
|
||||
const documents = await getDocuments();
|
||||
// Set private=false to mark the document as public (required for filtering)
|
||||
documents.forEach((doc) => {
|
||||
doc.metadata["private"] = "false";
|
||||
});
|
||||
|
||||
await VectorStoreIndex.fromDocuments(documents, {
|
||||
storageContext,
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import * as dotenv from "dotenv";
|
||||
import {
|
||||
VectorStoreIndex,
|
||||
WeaviateVectorStore,
|
||||
storageContextFromDefaults,
|
||||
} from "llamaindex";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { DEFAULT_INDEX_NAME, checkRequiredEnvVars } from "./shared";
|
||||
dotenv.config();
|
||||
|
||||
async function loadAndIndex() {
|
||||
const indexName = process.env.WEAVIATE_INDEX_NAME || DEFAULT_INDEX_NAME;
|
||||
|
||||
// load objects from storage and convert them into LlamaIndex Document objects
|
||||
const documents = await getDocuments();
|
||||
|
||||
const vectorStore = new WeaviateVectorStore({ indexName });
|
||||
|
||||
const storageContext = await storageContextFromDefaults({ vectorStore });
|
||||
await VectorStoreIndex.fromDocuments(documents, {
|
||||
storageContext: storageContext,
|
||||
});
|
||||
console.log(`Successfully upload embeddings to Weaviate index ${indexName}.`);
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
})();
|
||||
@@ -0,0 +1,14 @@
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { WeaviateVectorStore } from "llamaindex/storage/vectorStore/WeaviateVectorStore";
|
||||
import { checkRequiredEnvVars, DEFAULT_INDEX_NAME } from "./shared";
|
||||
|
||||
dotenv.config();
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
checkRequiredEnvVars();
|
||||
const indexName = process.env.WEAVIATE_INDEX_NAME || DEFAULT_INDEX_NAME;
|
||||
const store = new WeaviateVectorStore({ indexName });
|
||||
|
||||
return await VectorStoreIndex.fromVectorStore(store);
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
import { MetadataFilter, MetadataFilters } from "llamaindex";
|
||||
|
||||
export function generateFilters(documentIds: string[]): MetadataFilters {
|
||||
// filter all documents have the private metadata key set to true
|
||||
const publicDocumentsFilter: MetadataFilter = {
|
||||
key: "private",
|
||||
value: "true",
|
||||
operator: "!=",
|
||||
};
|
||||
|
||||
// if no documentIds are provided, only retrieve information from public documents
|
||||
if (!documentIds.length) return { filters: [publicDocumentsFilter] };
|
||||
|
||||
// Weaviate uses 'any' instead of 'in' for the operator
|
||||
const privateDocumentsFilter: MetadataFilter = {
|
||||
key: "doc_id",
|
||||
value: documentIds,
|
||||
operator: "any",
|
||||
};
|
||||
|
||||
// if documentIds are provided, retrieve information from public and private documents
|
||||
return {
|
||||
filters: [publicDocumentsFilter, privateDocumentsFilter],
|
||||
condition: "or",
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
const REQUIRED_ENV_VARS = ["WEAVIATE_CLUSTER_URL", "WEAVIATE_API_KEY"];
|
||||
|
||||
export const DEFAULT_INDEX_NAME = "LlamaIndex";
|
||||
|
||||
export function checkRequiredEnvVars() {
|
||||
const missingEnvVars = REQUIRED_ENV_VARS.filter((envVar) => {
|
||||
return !process.env[envVar];
|
||||
});
|
||||
|
||||
if (missingEnvVars.length > 0) {
|
||||
console.log(
|
||||
`The following environment variables are required but missing: ${missingEnvVars.join(
|
||||
", ",
|
||||
)}`,
|
||||
);
|
||||
throw new Error(
|
||||
`Missing environment variables: ${missingEnvVars.join(", ")}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
@@ -9,12 +9,13 @@ readme = "README.md"
|
||||
generate = "app.engine.generate:generate_datasource"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.11,<3.12"
|
||||
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,83 @@
|
||||
import asyncio
|
||||
from typing import Any, List
|
||||
|
||||
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
|
||||
|
||||
from app.agents.single import (
|
||||
AgentRunResult,
|
||||
ContextAwareTool,
|
||||
FunctionCallingAgent,
|
||||
)
|
||||
from app.agents.planner import StructuredPlannerAgent
|
||||
|
||||
|
||||
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:
|
||||
task = asyncio.create_task(self.agent.run(input=input))
|
||||
# bubble all events while running the agent to the calling agent
|
||||
async for ev in self.agent.stream_events():
|
||||
ctx.write_event_to_stream(ev)
|
||||
ret: AgentRunResult = await task
|
||||
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,328 @@
|
||||
import asyncio
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
from app.agents.single import AgentRunEvent, AgentRunResult, FunctionCallingAgent
|
||||
|
||||
|
||||
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
|
||||
task = asyncio.create_task(
|
||||
self.executor.run(
|
||||
input=ev.sub_task.input,
|
||||
streaming=streaming,
|
||||
)
|
||||
)
|
||||
# bubble all events while running the executor to the planner
|
||||
async for event in self.executor.stream_events():
|
||||
ctx.write_event_to_stream(event)
|
||||
result = await task
|
||||
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,43 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request, status
|
||||
from llama_index.core.workflow import Workflow
|
||||
|
||||
from app.examples.factory import create_agent
|
||||
from app.api.routers.models import (
|
||||
ChatData,
|
||||
)
|
||||
from app.api.routers.vercel_response import VercelStreamResponse
|
||||
|
||||
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)
|
||||
task = asyncio.create_task(
|
||||
agent.run(input=last_message_content, streaming=True)
|
||||
)
|
||||
|
||||
return VercelStreamResponse(request, task, 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?",
|
||||
]
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user