mirror of
https://github.com/run-llama/create-llama.git
synced 2026-07-18 13:05:55 -04:00
Compare commits
30 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4068618b2d | |||
| 54c9e2f95e | |||
| aec1173b71 | |||
| 481663dd63 | |||
| 1ca7dd2e48 | |||
| 3d20990713 | |||
| 8fb69cf807 | |||
| 61af56dac6 | |||
| 4b66039a96 | |||
| ee88f681a6 | |||
| 992c3a95e9 | |||
| 2a4fb702d1 | |||
| 24b9337096 | |||
| fceec69a3a | |||
| 03e5e0a16e | |||
| fe3cd36d3a | |||
| d5d10e9ead | |||
| 5ed925d75f | |||
| ca5df14d41 | |||
| ee69ce7cc1 | |||
| 0e4ecfaf8b | |||
| 3658fec684 | |||
| c3d275abe1 | |||
| 61204a1381 | |||
| 9e723c3a15 | |||
| d5da55b993 | |||
| c1552ebb00 | |||
| 131e63ae4a | |||
| 4e06714cdd | |||
| 18c8d2540c |
@@ -67,13 +67,16 @@ jobs:
|
||||
LLAMA_CLOUD_API_KEY: ${{ secrets.LLAMA_CLOUD_API_KEY }}
|
||||
FRAMEWORK: ${{ matrix.frameworks }}
|
||||
DATASOURCE: ${{ matrix.datasources }}
|
||||
PYTHONIOENCODING: utf-8
|
||||
PYTHONLEGACYWINDOWSSTDIO: utf-8
|
||||
working-directory: .
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
- uses: actions/upload-artifact@v4
|
||||
if: always()
|
||||
with:
|
||||
name: playwright-report-python
|
||||
name: playwright-report-python-${{ matrix.os }}-${{ matrix.frameworks }}-${{ matrix.datasources }}
|
||||
path: ./playwright-report/
|
||||
overwrite: true
|
||||
retention-days: 30
|
||||
|
||||
e2e-typescript:
|
||||
@@ -85,7 +88,7 @@ jobs:
|
||||
node-version: [18, 20]
|
||||
python-version: ["3.11"]
|
||||
os: [macos-latest, windows-latest, ubuntu-22.04]
|
||||
frameworks: ["nextjs", "express"]
|
||||
frameworks: ["nextjs"]
|
||||
datasources: ["--no-files", "--example-file", "--llamacloud"]
|
||||
defaults:
|
||||
run:
|
||||
@@ -136,9 +139,10 @@ jobs:
|
||||
DATASOURCE: ${{ matrix.datasources }}
|
||||
working-directory: .
|
||||
|
||||
- uses: actions/upload-artifact@v3
|
||||
- uses: actions/upload-artifact@v4
|
||||
if: always()
|
||||
with:
|
||||
name: playwright-report-typescript
|
||||
name: playwright-report-typescript-${{ matrix.os }}-${{ matrix.frameworks }}-${{ matrix.datasources }}-node${{ matrix.node-version }}
|
||||
path: ./playwright-report/
|
||||
overwrite: true
|
||||
retention-days: 30
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
name: Release llama-index-server
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "llama-index-server/**"
|
||||
- ".github/workflows/release_llama_index_server.yml"
|
||||
pull_request:
|
||||
types:
|
||||
- closed
|
||||
|
||||
concurrency: ${{ github.workflow }}-${{ github.ref }}
|
||||
|
||||
jobs:
|
||||
release:
|
||||
name: Create Release PR
|
||||
runs-on: ubuntu-latest
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ./llama-index-server
|
||||
if: |
|
||||
github.event_name == 'push' &&
|
||||
!startsWith(github.ref, 'refs/heads/release/llama-index-server-v')
|
||||
|
||||
steps:
|
||||
- name: Checkout Repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install Poetry
|
||||
run: |
|
||||
curl -sSL https://install.python-poetry.org | python3 -
|
||||
|
||||
- name: Install dependencies
|
||||
run: poetry install
|
||||
|
||||
- name: Setup Git
|
||||
run: |
|
||||
git config --global user.email "github-actions[bot]@users.noreply.github.com"
|
||||
git config --global user.name "github-actions[bot]"
|
||||
|
||||
- name: Bump patch version
|
||||
run: |
|
||||
poetry version patch
|
||||
git add pyproject.toml
|
||||
git commit -m "chore(release): bump version to $(poetry version -s)"
|
||||
|
||||
- name: Get current version
|
||||
id: get_version
|
||||
run: |
|
||||
version=$(poetry version -s)
|
||||
echo "current_version=${version}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Create Release PR
|
||||
uses: peter-evans/create-pull-request@v6
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
commit-message: "Release: llama-index-server v${{ steps.get_version.outputs.current_version }}"
|
||||
title: "Release: llama-index-server v${{ steps.get_version.outputs.current_version }}"
|
||||
body: |
|
||||
This PR was automatically created to release a new version of the llama-index-server package.
|
||||
|
||||
Version: ${{ steps.get_version.outputs.current_version }}
|
||||
|
||||
Please review the changes and merge to trigger the release.
|
||||
branch: release/llama-index-server-v${{ steps.get_version.outputs.current_version }}
|
||||
base: main
|
||||
labels: release, llama-index-server
|
||||
|
||||
publish:
|
||||
name: Publish to PyPI
|
||||
runs-on: ubuntu-latest
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ./llama-index-server
|
||||
if: |
|
||||
github.event_name == 'pull_request' &&
|
||||
github.event.pull_request.merged == true &&
|
||||
startsWith(github.event.pull_request.title, 'Release: llama-index-server') &&
|
||||
startsWith(github.event.pull_request.head.ref, 'release/llama-index-server-v')
|
||||
|
||||
steps:
|
||||
- name: Checkout Repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install Poetry
|
||||
run: |
|
||||
curl -sSL https://install.python-poetry.org | python3 -
|
||||
|
||||
- name: Install dependencies
|
||||
run: poetry install
|
||||
|
||||
- name: Get current version
|
||||
id: get_version
|
||||
run: |
|
||||
version=$(poetry version -s)
|
||||
echo "current_version=${version}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Build and publish to PyPI
|
||||
uses: JRubics/poetry-publish@v2.1
|
||||
with:
|
||||
python_version: "3.11"
|
||||
pypi_token: ${{ secrets.PYPI_TOKEN }}
|
||||
package_directory: "llama-index-server"
|
||||
poetry_install_options: "--without dev"
|
||||
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
tag_name: llama-index-server-v${{ steps.get_version.outputs.current_version }}
|
||||
name: "llama-index-server v${{ steps.get_version.outputs.current_version }}"
|
||||
body: |
|
||||
Release of llama-index-server v${{ steps.get_version.outputs.current_version }}
|
||||
draft: false
|
||||
prerelease: false
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
@@ -0,0 +1,111 @@
|
||||
name: Build Package
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
|
||||
env:
|
||||
POETRY_VERSION: "1.8.3"
|
||||
PYTHON_VERSION: "3.9"
|
||||
|
||||
jobs:
|
||||
unit-test:
|
||||
name: Unit Tests
|
||||
runs-on: ${{ matrix.os }}
|
||||
defaults:
|
||||
run:
|
||||
working-directory: llama-index-server
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest]
|
||||
python-version: ["3.9"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Poetry
|
||||
run: pipx install poetry==${{ env.POETRY_VERSION }}
|
||||
|
||||
- name: Set up python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
cache: "poetry"
|
||||
|
||||
- name: Configure Poetry
|
||||
run: |
|
||||
poetry config virtualenvs.create true
|
||||
poetry config virtualenvs.in-project true
|
||||
poetry env use python
|
||||
|
||||
- name: Install dependencies
|
||||
shell: bash
|
||||
run: poetry install --with dev
|
||||
|
||||
- name: Run unit tests
|
||||
shell: bash
|
||||
run: |
|
||||
poetry run pytest tests
|
||||
|
||||
type-check:
|
||||
name: Type Check
|
||||
runs-on: ubuntu-latest
|
||||
defaults:
|
||||
run:
|
||||
working-directory: llama-index-server
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Poetry
|
||||
run: pipx install poetry==${{ env.POETRY_VERSION }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
cache: "poetry"
|
||||
|
||||
- name: Configure Poetry
|
||||
run: |
|
||||
poetry config virtualenvs.create true
|
||||
poetry config virtualenvs.in-project true
|
||||
poetry env use python
|
||||
|
||||
- name: Install dependencies
|
||||
shell: bash
|
||||
run: poetry install --with dev
|
||||
|
||||
- name: Run mypy
|
||||
shell: bash
|
||||
run: poetry run mypy llama_index
|
||||
|
||||
build:
|
||||
needs: [unit-test, type-check]
|
||||
runs-on: ubuntu-latest
|
||||
defaults:
|
||||
run:
|
||||
working-directory: llama-index-server
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: Install Poetry
|
||||
run: pipx install poetry==${{ env.POETRY_VERSION }}
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ env.PYTHON_VERSION }}
|
||||
- name: Clear python cache
|
||||
shell: bash
|
||||
run: poetry cache clear --all pypi
|
||||
- name: Build package
|
||||
shell: bash
|
||||
run: poetry build
|
||||
- name: Test installing built package
|
||||
shell: bash
|
||||
run: python -m pip install .
|
||||
- name: Test import
|
||||
shell: bash
|
||||
working-directory: ${{ vars.RUNNER_TEMP }}
|
||||
run: python -c "from llama_index.server import LlamaIndexServer"
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: llama-index-server
|
||||
path: llama-index-server/dist/
|
||||
@@ -8,6 +8,7 @@ node_modules
|
||||
|
||||
# testing
|
||||
coverage
|
||||
.coverage
|
||||
|
||||
# next.js
|
||||
.next/
|
||||
@@ -48,6 +49,13 @@ e2e/cache
|
||||
|
||||
# Python
|
||||
.mypy_cache/
|
||||
venv/
|
||||
.venv/
|
||||
dist/
|
||||
.__pycache__
|
||||
__pycache__
|
||||
.python-version
|
||||
.ui
|
||||
|
||||
# build artifacts
|
||||
create-llama-*.tgz
|
||||
|
||||
@@ -1,5 +1,34 @@
|
||||
# create-llama
|
||||
|
||||
## 0.5.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 54c9e2f: Simplified generated code using LlamaIndexServer
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 0e4ecfa: fix: add trycatch for generating error
|
||||
- ee69ce7: bump: chat-ui and tailwind v4
|
||||
|
||||
## 0.4.0
|
||||
|
||||
### Minor Changes
|
||||
|
||||
- 61204a1: chore: bump LITS 0.9
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 9e723c3: Standardize the code of the workflow use case (Python)
|
||||
- d5da55b: feat: add components.json to use CLI
|
||||
- c1552eb: chore: move wikipedia tool to create-llama
|
||||
|
||||
## 0.3.28
|
||||
|
||||
### Patch Changes
|
||||
|
||||
- 4e06714: Fix the error: Unable to view file sources due to CORS.
|
||||
|
||||
## 0.3.27
|
||||
|
||||
### Patch Changes
|
||||
|
||||
+2
-2
@@ -90,7 +90,7 @@ export async function createApp({
|
||||
// Install backend
|
||||
await installTemplate({ ...args, backend: true });
|
||||
|
||||
if (frontend && framework === "fastapi") {
|
||||
if (frontend && framework === "fastapi" && template !== "llamaindexserver") {
|
||||
// install frontend
|
||||
const frontendRoot = path.join(root, ".frontend");
|
||||
await makeDir(frontendRoot);
|
||||
@@ -110,7 +110,7 @@ export async function createApp({
|
||||
console.log();
|
||||
}
|
||||
|
||||
if (toolsRequireConfig(tools)) {
|
||||
if (toolsRequireConfig(tools) && template !== "llamaindexserver") {
|
||||
const configFile =
|
||||
framework === "fastapi" ? "config/tools.yaml" : "config/tools.json";
|
||||
console.log(
|
||||
|
||||
+16
-8
@@ -16,15 +16,17 @@ const templateFramework: TemplateFramework = process.env.FRAMEWORK
|
||||
const dataSource: string = "--example-file";
|
||||
const templateUI: TemplateUI = "shadcn";
|
||||
const templatePostInstallAction: TemplatePostInstallAction = "runApp";
|
||||
const appType: AppType = templateFramework === "fastapi" ? "--frontend" : "";
|
||||
const appType: AppType = "--frontend";
|
||||
const userMessage = "Write a blog post about physical standards for letters";
|
||||
const templateUseCases = ["financial_report", "blog", "form_filling"];
|
||||
const templateUseCases = ["financial_report", "agentic_rag", "deep_research"];
|
||||
|
||||
for (const useCase of templateUseCases) {
|
||||
test.describe(`Test multiagent template ${useCase} ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
|
||||
test.describe(`Test use case ${useCase} ${templateFramework} ${dataSource} ${templateUI} ${appType} ${templatePostInstallAction}`, async () => {
|
||||
test.skip(
|
||||
process.platform !== "linux" || process.env.DATASOURCE === "--no-files",
|
||||
"The multiagent template currently only works with files. We also only run on Linux to speed up tests.",
|
||||
process.platform !== "linux" ||
|
||||
process.env.DATASOURCE === "--no-files" ||
|
||||
templateFramework === "express",
|
||||
"The llamaindexserver template currently only works with nextjs, fastapi. We also only run on Linux to speed up tests.",
|
||||
);
|
||||
let port: number;
|
||||
let cwd: string;
|
||||
@@ -38,7 +40,7 @@ for (const useCase of templateUseCases) {
|
||||
cwd = await createTestDir();
|
||||
const result = await runCreateLlama({
|
||||
cwd,
|
||||
templateType: "multiagent",
|
||||
templateType: "llamaindexserver",
|
||||
templateFramework,
|
||||
dataSource,
|
||||
vectorDb,
|
||||
@@ -72,9 +74,9 @@ for (const useCase of templateUseCases) {
|
||||
test.skip(
|
||||
templatePostInstallAction !== "runApp" ||
|
||||
useCase === "financial_report" ||
|
||||
useCase === "form_filling" ||
|
||||
useCase === "deep_research" ||
|
||||
templateFramework === "express",
|
||||
"Skip chat tests for financial report and form filling.",
|
||||
"Skip chat tests for financial report and deep research.",
|
||||
);
|
||||
await page.goto(`http://localhost:${port}`);
|
||||
await page.fill("form textarea", userMessage);
|
||||
@@ -86,6 +88,12 @@ for (const useCase of templateUseCases) {
|
||||
await page.click("form button[type=submit]");
|
||||
|
||||
const response = await responsePromise;
|
||||
console.log(`Response status: ${response.status()}`);
|
||||
const responseBody = await response
|
||||
.text()
|
||||
.catch((e) => `Error reading body: ${e}`);
|
||||
console.log(`Response body: ${responseBody}`);
|
||||
|
||||
expect(response.ok()).toBeTruthy();
|
||||
});
|
||||
|
||||
+6
-1
@@ -113,7 +113,12 @@ export async function runCreateLlama({
|
||||
if (observability) {
|
||||
commandArgs.push("--observability", observability);
|
||||
}
|
||||
if ((templateType === "multiagent" || templateType === "reflex") && useCase) {
|
||||
if (
|
||||
(templateType === "multiagent" ||
|
||||
templateType === "reflex" ||
|
||||
templateType === "llamaindexserver") &&
|
||||
useCase
|
||||
) {
|
||||
commandArgs.push("--use-case", useCase);
|
||||
}
|
||||
|
||||
|
||||
+59
-36
@@ -44,6 +44,7 @@ const renderEnvVar = (envVars: EnvVar[]): string => {
|
||||
const getVectorDBEnvs = (
|
||||
vectorDb?: TemplateVectorDB,
|
||||
framework?: TemplateFramework,
|
||||
template?: TemplateType,
|
||||
): EnvVar[] => {
|
||||
if (!vectorDb || !framework) {
|
||||
return [];
|
||||
@@ -168,7 +169,7 @@ const getVectorDBEnvs = (
|
||||
description:
|
||||
"The organization ID for the LlamaCloud project (uses default organization if not specified)",
|
||||
},
|
||||
...(framework === "nextjs"
|
||||
...(framework === "nextjs" && template !== "llamaindexserver"
|
||||
? // activate index selector per default (not needed for non-NextJS backends as it's handled by createFrontendEnvFile)
|
||||
[
|
||||
{
|
||||
@@ -223,13 +224,15 @@ Otherwise, use CHROMA_HOST and CHROMA_PORT config above`,
|
||||
},
|
||||
];
|
||||
default:
|
||||
return [
|
||||
{
|
||||
name: "STORAGE_CACHE_DIR",
|
||||
description: "The directory to store the local storage cache.",
|
||||
value: ".cache",
|
||||
},
|
||||
];
|
||||
return template !== "llamaindexserver"
|
||||
? [
|
||||
{
|
||||
name: "STORAGE_CACHE_DIR",
|
||||
description: "The directory to store the local storage cache.",
|
||||
value: ".cache",
|
||||
},
|
||||
]
|
||||
: [];
|
||||
}
|
||||
};
|
||||
|
||||
@@ -382,38 +385,42 @@ const getModelEnvs = (modelConfig: ModelConfig): EnvVar[] => {
|
||||
|
||||
const getFrameworkEnvs = (
|
||||
framework: TemplateFramework,
|
||||
template: TemplateType,
|
||||
port?: number,
|
||||
): EnvVar[] => {
|
||||
const sPort = port?.toString() || "8000";
|
||||
const result: EnvVar[] = [
|
||||
{
|
||||
name: "FILESERVER_URL_PREFIX",
|
||||
description:
|
||||
"FILESERVER_URL_PREFIX is the URL prefix of the server storing the images generated by the interpreter.",
|
||||
value:
|
||||
framework === "nextjs"
|
||||
? // FIXME: if we are using nextjs, port should be 3000
|
||||
"http://localhost:3000/api/files"
|
||||
: `http://localhost:${sPort}/api/files`,
|
||||
},
|
||||
];
|
||||
const result: EnvVar[] =
|
||||
template !== "llamaindexserver"
|
||||
? [
|
||||
{
|
||||
name: "FILESERVER_URL_PREFIX",
|
||||
description:
|
||||
"FILESERVER_URL_PREFIX is the URL prefix of the server storing the images generated by the interpreter.",
|
||||
value:
|
||||
framework === "nextjs"
|
||||
? // FIXME: if we are using nextjs, port should be 3000
|
||||
"http://localhost:3000/api/files"
|
||||
: `http://localhost:${sPort}/api/files`,
|
||||
},
|
||||
]
|
||||
: [];
|
||||
if (framework === "fastapi") {
|
||||
result.push(
|
||||
...[
|
||||
{
|
||||
name: "APP_HOST",
|
||||
description: "The address to start the backend app.",
|
||||
description: "The address to start the FastAPI app.",
|
||||
value: "0.0.0.0",
|
||||
},
|
||||
{
|
||||
name: "APP_PORT",
|
||||
description: "The port to start the backend app.",
|
||||
description: "The port to start the FastAPI app.",
|
||||
value: sPort,
|
||||
},
|
||||
],
|
||||
);
|
||||
}
|
||||
if (framework === "nextjs") {
|
||||
if (framework === "nextjs" && template !== "llamaindexserver") {
|
||||
result.push({
|
||||
name: "NEXT_PUBLIC_CHAT_API",
|
||||
description:
|
||||
@@ -569,25 +576,41 @@ export const createBackendEnvFile = async (
|
||||
| "port"
|
||||
| "tools"
|
||||
| "observability"
|
||||
| "useLlamaParse"
|
||||
>,
|
||||
) => {
|
||||
// Init env values
|
||||
const envFileName = ".env";
|
||||
const envVars: EnvVar[] = [
|
||||
{
|
||||
name: "LLAMA_CLOUD_API_KEY",
|
||||
description: `The Llama Cloud API key.`,
|
||||
value: opts.llamaCloudKey,
|
||||
},
|
||||
// Add environment variables of each component
|
||||
...getModelEnvs(opts.modelConfig),
|
||||
...getEngineEnvs(),
|
||||
...getVectorDBEnvs(opts.vectorDb, opts.framework),
|
||||
...getFrameworkEnvs(opts.framework, opts.port),
|
||||
...(opts.useLlamaParse
|
||||
? [
|
||||
{
|
||||
name: "LLAMA_CLOUD_API_KEY",
|
||||
description: `The Llama Cloud API key.`,
|
||||
value: opts.llamaCloudKey,
|
||||
},
|
||||
]
|
||||
: []),
|
||||
...getVectorDBEnvs(opts.vectorDb, opts.framework, opts.template),
|
||||
...getToolEnvs(opts.tools),
|
||||
...getTemplateEnvs(opts.template),
|
||||
...getObservabilityEnvs(opts.observability),
|
||||
...getSystemPromptEnv(opts.tools, opts.dataSources, opts.template),
|
||||
...getFrameworkEnvs(opts.framework, opts.template, opts.port),
|
||||
// Add environment variables of each component
|
||||
...(opts.template === "llamaindexserver"
|
||||
? [
|
||||
{
|
||||
name: "OPENAI_API_KEY",
|
||||
description: "The OpenAI API key to use.",
|
||||
value: opts.modelConfig.apiKey,
|
||||
},
|
||||
]
|
||||
: [
|
||||
// don't use this stuff for llama-indexserver
|
||||
...getModelEnvs(opts.modelConfig),
|
||||
...getEngineEnvs(),
|
||||
...getTemplateEnvs(opts.template),
|
||||
...getObservabilityEnvs(opts.observability),
|
||||
...getSystemPromptEnv(opts.tools, opts.dataSources, opts.template),
|
||||
]),
|
||||
];
|
||||
// Render and write env file
|
||||
const content = renderEnvVar(envVars);
|
||||
|
||||
+19
-17
@@ -1,7 +1,7 @@
|
||||
import { callPackageManager } from "./install";
|
||||
|
||||
import path from "path";
|
||||
import { cyan } from "picocolors";
|
||||
import picocolors, { cyan } from "picocolors";
|
||||
|
||||
import fsExtra from "fs-extra";
|
||||
import { writeLoadersConfig } from "./datasources";
|
||||
@@ -41,7 +41,11 @@ const checkForGenerateScript = (
|
||||
missingSettings.push("your LLAMA_CLOUD_API_KEY");
|
||||
}
|
||||
|
||||
if (vectorDb !== "none" && vectorDb !== "llamacloud") {
|
||||
if (
|
||||
vectorDb !== undefined &&
|
||||
vectorDb !== "none" &&
|
||||
vectorDb !== "llamacloud"
|
||||
) {
|
||||
missingSettings.push("your Vector DB environment variables");
|
||||
}
|
||||
|
||||
@@ -92,7 +96,7 @@ async function generateContextData(
|
||||
}
|
||||
|
||||
const settingsMessage = `After setting ${missingSettings.join(" and ")}, run ${runGenerate} to generate the context data.`;
|
||||
console.log(`\n${settingsMessage}\n\n`);
|
||||
console.log(picocolors.yellow(`\n${settingsMessage}\n\n`));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -166,6 +170,17 @@ export const installTemplate = async (
|
||||
|
||||
if (props.framework === "fastapi") {
|
||||
await installPythonTemplate(props);
|
||||
} else {
|
||||
await installTSTemplate(props);
|
||||
}
|
||||
|
||||
// write configurations
|
||||
if (props.template !== "llamaindexserver") {
|
||||
await writeToolsConfig(
|
||||
props.root,
|
||||
props.tools,
|
||||
props.framework === "fastapi" ? ConfigFileType.YAML : ConfigFileType.JSON,
|
||||
);
|
||||
if (props.vectorDb !== "llamacloud") {
|
||||
// write loaders configuration (currently Python only)
|
||||
// not needed for LlamaCloud as it has its own loaders
|
||||
@@ -175,26 +190,13 @@ export const installTemplate = async (
|
||||
props.useLlamaParse,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
await installTSTemplate(props);
|
||||
}
|
||||
|
||||
// write tools configuration
|
||||
await writeToolsConfig(
|
||||
props.root,
|
||||
props.tools,
|
||||
props.framework === "fastapi" ? ConfigFileType.YAML : ConfigFileType.JSON,
|
||||
);
|
||||
|
||||
if (props.backend) {
|
||||
// This is a backend, so we need to copy the test data and create the env file.
|
||||
|
||||
// Copy the environment file to the target directory.
|
||||
if (
|
||||
props.template === "streaming" ||
|
||||
props.template === "multiagent" ||
|
||||
props.template === "reflex"
|
||||
) {
|
||||
if (props.template !== "community" && props.template !== "llamapack") {
|
||||
await createBackendEnvFile(props.root, props);
|
||||
}
|
||||
|
||||
|
||||
+140
-52
@@ -12,6 +12,7 @@ import {
|
||||
InstallTemplateArgs,
|
||||
ModelConfig,
|
||||
TemplateDataSource,
|
||||
TemplateObservability,
|
||||
TemplateType,
|
||||
TemplateVectorDB,
|
||||
} from "./types";
|
||||
@@ -29,6 +30,7 @@ const getAdditionalDependencies = (
|
||||
dataSources?: TemplateDataSource[],
|
||||
tools?: Tool[],
|
||||
templateType?: TemplateType,
|
||||
observability?: TemplateObservability,
|
||||
) => {
|
||||
const dependencies: Dependency[] = [];
|
||||
|
||||
@@ -103,7 +105,7 @@ const getAdditionalDependencies = (
|
||||
case "llamacloud":
|
||||
dependencies.push({
|
||||
name: "llama-index-indices-managed-llama-cloud",
|
||||
version: "^0.6.3",
|
||||
version: "0.6.3",
|
||||
});
|
||||
break;
|
||||
}
|
||||
@@ -268,6 +270,21 @@ const getAdditionalDependencies = (
|
||||
break;
|
||||
}
|
||||
|
||||
if (observability && observability !== "none") {
|
||||
if (observability === "traceloop") {
|
||||
dependencies.push({
|
||||
name: "traceloop-sdk",
|
||||
version: "^0.15.11",
|
||||
});
|
||||
}
|
||||
if (observability === "llamatrace") {
|
||||
dependencies.push({
|
||||
name: "llama-index-callbacks-arize-phoenix",
|
||||
version: "^0.3.0",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return dependencies;
|
||||
};
|
||||
|
||||
@@ -379,47 +396,24 @@ export const installPythonDependencies = (
|
||||
}
|
||||
};
|
||||
|
||||
export const installPythonTemplate = async ({
|
||||
appName,
|
||||
const installLegacyPythonTemplate = async ({
|
||||
root,
|
||||
template,
|
||||
framework,
|
||||
vectorDb,
|
||||
postInstallAction,
|
||||
modelConfig,
|
||||
dataSources,
|
||||
tools,
|
||||
useLlamaParse,
|
||||
useCase,
|
||||
observability,
|
||||
}: Pick<
|
||||
InstallTemplateArgs,
|
||||
| "appName"
|
||||
| "root"
|
||||
| "template"
|
||||
| "framework"
|
||||
| "vectorDb"
|
||||
| "postInstallAction"
|
||||
| "modelConfig"
|
||||
| "dataSources"
|
||||
| "tools"
|
||||
| "useLlamaParse"
|
||||
| "useCase"
|
||||
| "observability"
|
||||
>) => {
|
||||
console.log("\nInitializing Python project with template:", template, "\n");
|
||||
let templatePath;
|
||||
if (template === "reflex") {
|
||||
templatePath = path.join(templatesDir, "types", "reflex");
|
||||
} else {
|
||||
templatePath = path.join(templatesDir, "types", "streaming", framework);
|
||||
}
|
||||
await copy("**", root, {
|
||||
parents: true,
|
||||
cwd: templatePath,
|
||||
rename: assetRelocator,
|
||||
});
|
||||
|
||||
const compPath = path.join(templatesDir, "components");
|
||||
const enginePath = path.join(root, "app", "engine");
|
||||
|
||||
@@ -509,34 +503,7 @@ export const installPythonTemplate = async ({
|
||||
}
|
||||
}
|
||||
|
||||
console.log("Adding additional dependencies");
|
||||
|
||||
const addOnDependencies = getAdditionalDependencies(
|
||||
modelConfig,
|
||||
vectorDb,
|
||||
dataSources,
|
||||
tools,
|
||||
template,
|
||||
);
|
||||
|
||||
if (observability && observability !== "none") {
|
||||
if (observability === "traceloop") {
|
||||
addOnDependencies.push({
|
||||
name: "traceloop-sdk",
|
||||
version: "^0.15.11",
|
||||
});
|
||||
}
|
||||
|
||||
if (observability === "llamatrace") {
|
||||
addOnDependencies.push({
|
||||
name: "llama-index-callbacks-arize-phoenix",
|
||||
version: "^0.3.0",
|
||||
constraints: {
|
||||
python: ">=3.11,<3.13",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
const templateObservabilityPath = path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
@@ -548,6 +515,127 @@ export const installPythonTemplate = async ({
|
||||
cwd: templateObservabilityPath,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
const installLlamaIndexServerTemplate = async ({
|
||||
root,
|
||||
useCase,
|
||||
useLlamaParse,
|
||||
}: Pick<InstallTemplateArgs, "root" | "useCase" | "useLlamaParse">) => {
|
||||
if (!useCase) {
|
||||
console.log(
|
||||
red(
|
||||
`There is no use case selected. Please pick a use case to use via --use-case flag.`,
|
||||
),
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
await copy("workflow.py", path.join(root, "app"), {
|
||||
parents: true,
|
||||
cwd: path.join(templatesDir, "components", "workflows", "python", useCase),
|
||||
});
|
||||
|
||||
if (useLlamaParse) {
|
||||
await copy("index.py", path.join(root, "app"), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"vectordbs",
|
||||
"llamaindexserver",
|
||||
"llamacloud",
|
||||
"python",
|
||||
),
|
||||
});
|
||||
// TODO: Consider moving generate.py to app folder.
|
||||
await copy("generate.py", path.join(root), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"vectordbs",
|
||||
"llamaindexserver",
|
||||
"llamacloud",
|
||||
"python",
|
||||
),
|
||||
});
|
||||
}
|
||||
// Copy README.md
|
||||
await copy("README-template.md", path.join(root), {
|
||||
parents: true,
|
||||
cwd: path.join(templatesDir, "components", "workflows", "python", useCase),
|
||||
rename: assetRelocator,
|
||||
});
|
||||
};
|
||||
|
||||
export const installPythonTemplate = async ({
|
||||
appName,
|
||||
root,
|
||||
template,
|
||||
framework,
|
||||
vectorDb,
|
||||
postInstallAction,
|
||||
modelConfig,
|
||||
dataSources,
|
||||
tools,
|
||||
useLlamaParse,
|
||||
useCase,
|
||||
observability,
|
||||
}: Pick<
|
||||
InstallTemplateArgs,
|
||||
| "appName"
|
||||
| "root"
|
||||
| "template"
|
||||
| "framework"
|
||||
| "vectorDb"
|
||||
| "postInstallAction"
|
||||
| "modelConfig"
|
||||
| "dataSources"
|
||||
| "tools"
|
||||
| "useLlamaParse"
|
||||
| "useCase"
|
||||
| "observability"
|
||||
>) => {
|
||||
console.log("\nInitializing Python project with template:", template, "\n");
|
||||
let templatePath;
|
||||
if (template === "reflex") {
|
||||
templatePath = path.join(templatesDir, "types", "reflex");
|
||||
} else {
|
||||
templatePath = path.join(templatesDir, "types", template, framework);
|
||||
}
|
||||
await copy("**", root, {
|
||||
parents: true,
|
||||
cwd: templatePath,
|
||||
rename: assetRelocator,
|
||||
});
|
||||
|
||||
if (template === "llamaindexserver") {
|
||||
await installLlamaIndexServerTemplate({
|
||||
root,
|
||||
useCase,
|
||||
useLlamaParse,
|
||||
});
|
||||
} else {
|
||||
await installLegacyPythonTemplate({
|
||||
root,
|
||||
template,
|
||||
vectorDb,
|
||||
dataSources,
|
||||
tools,
|
||||
useCase,
|
||||
observability,
|
||||
});
|
||||
}
|
||||
|
||||
console.log("Adding additional dependencies");
|
||||
const addOnDependencies = getAdditionalDependencies(
|
||||
modelConfig,
|
||||
vectorDb,
|
||||
dataSources,
|
||||
tools,
|
||||
template,
|
||||
);
|
||||
|
||||
await addDependencies(root, addOnDependencies);
|
||||
|
||||
|
||||
+10
-3
@@ -124,7 +124,7 @@ For better results, you can specify the region parameter to get results from a s
|
||||
dependencies: [
|
||||
{
|
||||
name: "e2b_code_interpreter",
|
||||
version: "1.0.3",
|
||||
version: "1.1.1",
|
||||
},
|
||||
],
|
||||
supportedFrameworks: ["fastapi", "express", "nextjs"],
|
||||
@@ -155,7 +155,7 @@ For better results, you can specify the region parameter to get results from a s
|
||||
dependencies: [
|
||||
{
|
||||
name: "e2b_code_interpreter",
|
||||
version: "1.0.3",
|
||||
version: "1.1.1",
|
||||
},
|
||||
],
|
||||
supportedFrameworks: ["fastapi", "express", "nextjs"],
|
||||
@@ -325,9 +325,16 @@ export const writeToolsConfig = async (
|
||||
yaml.stringify(configContent),
|
||||
);
|
||||
} else {
|
||||
// For Typescript, we treat llamahub tools as local tools
|
||||
const tsConfigContent = {
|
||||
local: {
|
||||
...configContent.local,
|
||||
...configContent.llamahub,
|
||||
},
|
||||
};
|
||||
await fs.writeFile(
|
||||
path.join(configPath, "tools.json"),
|
||||
JSON.stringify(configContent, null, 2),
|
||||
JSON.stringify(tsConfigContent, null, 2),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
+4
-2
@@ -24,7 +24,8 @@ export type TemplateType =
|
||||
| "community"
|
||||
| "llamapack"
|
||||
| "multiagent"
|
||||
| "reflex";
|
||||
| "reflex"
|
||||
| "llamaindexserver";
|
||||
export type TemplateFramework = "nextjs" | "express" | "fastapi";
|
||||
export type TemplateUI = "html" | "shadcn";
|
||||
export type TemplateVectorDB =
|
||||
@@ -55,7 +56,8 @@ export type TemplateUseCase =
|
||||
| "deep_research"
|
||||
| "form_filling"
|
||||
| "extractor"
|
||||
| "contract_review";
|
||||
| "contract_review"
|
||||
| "agentic_rag";
|
||||
// Config for both file and folder
|
||||
export type FileSourceConfig =
|
||||
| {
|
||||
|
||||
+244
-52
@@ -6,43 +6,100 @@ import { assetRelocator, copy } from "../helpers/copy";
|
||||
import { callPackageManager } from "../helpers/install";
|
||||
import { templatesDir } from "./dir";
|
||||
import { PackageManager } from "./get-pkg-manager";
|
||||
import { InstallTemplateArgs } from "./types";
|
||||
import { InstallTemplateArgs, ModelProvider, TemplateVectorDB } from "./types";
|
||||
|
||||
/**
|
||||
* Install a LlamaIndex internal template to a given `root` directory.
|
||||
*/
|
||||
export const installTSTemplate = async ({
|
||||
appName,
|
||||
const installLlamaIndexServerTemplate = async ({
|
||||
root,
|
||||
useCase,
|
||||
vectorDb,
|
||||
}: Pick<InstallTemplateArgs, "root" | "useCase" | "vectorDb">) => {
|
||||
if (!useCase) {
|
||||
console.log(
|
||||
red(
|
||||
`There is no use case selected. Please pick a use case to use via --use-case flag.`,
|
||||
),
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
if (!vectorDb) {
|
||||
console.log(
|
||||
red(
|
||||
`There is no vector db selected. Please pick a vector db to use via --vector-db flag.`,
|
||||
),
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
await copy("workflow.ts", path.join(root, "src", "app"), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"workflows",
|
||||
"typescript",
|
||||
useCase,
|
||||
),
|
||||
});
|
||||
|
||||
if (vectorDb === "llamacloud") {
|
||||
await copy("generate.ts", path.join(root, "src"), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"vectordbs",
|
||||
"llamaindexserver",
|
||||
"llamacloud",
|
||||
"typescript",
|
||||
),
|
||||
});
|
||||
|
||||
await copy("index.ts", path.join(root, "src", "app"), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"vectordbs",
|
||||
"llamaindexserver",
|
||||
"llamacloud",
|
||||
"typescript",
|
||||
),
|
||||
rename: () => "data.ts",
|
||||
});
|
||||
}
|
||||
// Copy README.md
|
||||
await copy("README-template.md", path.join(root), {
|
||||
parents: true,
|
||||
cwd: path.join(
|
||||
templatesDir,
|
||||
"components",
|
||||
"workflows",
|
||||
"typescript",
|
||||
useCase,
|
||||
),
|
||||
rename: assetRelocator,
|
||||
});
|
||||
};
|
||||
|
||||
const installLegacyTSTemplate = async ({
|
||||
root,
|
||||
packageManager,
|
||||
isOnline,
|
||||
template,
|
||||
backend,
|
||||
framework,
|
||||
ui,
|
||||
vectorDb,
|
||||
postInstallAction,
|
||||
backend,
|
||||
observability,
|
||||
tools,
|
||||
dataSources,
|
||||
useLlamaParse,
|
||||
useCase,
|
||||
}: InstallTemplateArgs & { backend: boolean }) => {
|
||||
console.log(bold(`Using ${packageManager}.`));
|
||||
|
||||
/**
|
||||
* Copy the template files to the target directory.
|
||||
*/
|
||||
console.log("\nInitializing project with template:", template, "\n");
|
||||
const templatePath = path.join(templatesDir, "types", "streaming", framework);
|
||||
const copySource = ["**"];
|
||||
|
||||
await copy(copySource, root, {
|
||||
parents: true,
|
||||
cwd: templatePath,
|
||||
rename: assetRelocator,
|
||||
});
|
||||
|
||||
modelConfig,
|
||||
relativeEngineDestPath,
|
||||
}: InstallTemplateArgs & {
|
||||
backend: boolean;
|
||||
relativeEngineDestPath: string;
|
||||
}) => {
|
||||
/**
|
||||
* If next.js is used, update its configuration if necessary
|
||||
*/
|
||||
@@ -97,10 +154,6 @@ export const installTSTemplate = async ({
|
||||
}
|
||||
|
||||
const compPath = path.join(templatesDir, "components");
|
||||
const relativeEngineDestPath =
|
||||
framework === "nextjs"
|
||||
? path.join("app", "api", "chat")
|
||||
: path.join("src", "controllers");
|
||||
const enginePath = path.join(root, relativeEngineDestPath, "engine");
|
||||
|
||||
// copy llamaindex code for TS templates
|
||||
@@ -181,6 +234,12 @@ export const installTSTemplate = async ({
|
||||
cwd: path.join(compPath, "loaders", "typescript", loaderFolder),
|
||||
});
|
||||
|
||||
// copy provider settings
|
||||
await copy("**", enginePath, {
|
||||
parents: true,
|
||||
cwd: path.join(compPath, "providers", "typescript", modelConfig.provider),
|
||||
});
|
||||
|
||||
// Select and copy engine code based on data sources and tools
|
||||
let engine;
|
||||
tools = tools ?? [];
|
||||
@@ -229,6 +288,75 @@ export const installTSTemplate = async ({
|
||||
await fs.rm(path.join(root, "app", "api"), { recursive: true });
|
||||
await fs.rm(path.join(root, "config"), { recursive: true, force: true });
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Install a LlamaIndex internal template to a given `root` directory.
|
||||
*/
|
||||
export const installTSTemplate = async ({
|
||||
appName,
|
||||
root,
|
||||
packageManager,
|
||||
isOnline,
|
||||
template,
|
||||
framework,
|
||||
ui,
|
||||
vectorDb,
|
||||
postInstallAction,
|
||||
backend,
|
||||
observability,
|
||||
tools,
|
||||
dataSources,
|
||||
useLlamaParse,
|
||||
useCase,
|
||||
modelConfig,
|
||||
}: InstallTemplateArgs & { backend: boolean }) => {
|
||||
console.log(bold(`Using ${packageManager}.`));
|
||||
|
||||
/**
|
||||
* 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 copySource = ["**"];
|
||||
|
||||
await copy(copySource, root, {
|
||||
parents: true,
|
||||
cwd: templatePath,
|
||||
rename: assetRelocator,
|
||||
});
|
||||
|
||||
const relativeEngineDestPath =
|
||||
framework === "nextjs"
|
||||
? path.join("app", "api", "chat")
|
||||
: path.join("src", "controllers");
|
||||
|
||||
if (template === "llamaindexserver") {
|
||||
await installLlamaIndexServerTemplate({
|
||||
root,
|
||||
useCase,
|
||||
vectorDb,
|
||||
});
|
||||
} else {
|
||||
await installLegacyTSTemplate({
|
||||
appName,
|
||||
root,
|
||||
packageManager,
|
||||
isOnline,
|
||||
template,
|
||||
backend,
|
||||
framework,
|
||||
ui,
|
||||
vectorDb,
|
||||
observability,
|
||||
tools,
|
||||
dataSources,
|
||||
useLlamaParse,
|
||||
useCase,
|
||||
modelConfig,
|
||||
relativeEngineDestPath,
|
||||
});
|
||||
}
|
||||
|
||||
const packageJson = await updatePackageJson({
|
||||
root,
|
||||
@@ -239,6 +367,9 @@ export const installTSTemplate = async ({
|
||||
ui,
|
||||
observability,
|
||||
vectorDb,
|
||||
backend,
|
||||
modelConfig,
|
||||
template,
|
||||
});
|
||||
|
||||
if (
|
||||
@@ -249,6 +380,68 @@ export const installTSTemplate = async ({
|
||||
}
|
||||
};
|
||||
|
||||
const providerDependencies: {
|
||||
[key in ModelProvider]?: Record<string, string>;
|
||||
} = {
|
||||
openai: {
|
||||
"@llamaindex/openai": "^0.2.0",
|
||||
},
|
||||
gemini: {
|
||||
"@llamaindex/google": "^0.2.0",
|
||||
},
|
||||
ollama: {
|
||||
"@llamaindex/ollama": "^0.1.0",
|
||||
},
|
||||
mistral: {
|
||||
"@llamaindex/mistral": "^0.2.0",
|
||||
},
|
||||
"azure-openai": {
|
||||
"@llamaindex/openai": "^0.2.0",
|
||||
},
|
||||
groq: {
|
||||
"@llamaindex/groq": "^0.0.61",
|
||||
"@llamaindex/huggingface": "^0.1.0", // groq uses huggingface as default embedding model
|
||||
},
|
||||
anthropic: {
|
||||
"@llamaindex/anthropic": "^0.3.0",
|
||||
"@llamaindex/huggingface": "^0.1.0", // anthropic uses huggingface as default embedding model
|
||||
},
|
||||
};
|
||||
|
||||
const vectorDbDependencies: Record<TemplateVectorDB, Record<string, string>> = {
|
||||
astra: {
|
||||
"@llamaindex/astra": "^0.0.5",
|
||||
},
|
||||
chroma: {
|
||||
"@llamaindex/chroma": "^0.0.5",
|
||||
},
|
||||
llamacloud: {},
|
||||
milvus: {
|
||||
"@zilliz/milvus2-sdk-node": "^2.4.6",
|
||||
"@llamaindex/milvus": "^0.1.0",
|
||||
},
|
||||
mongo: {
|
||||
mongodb: "6.7.0",
|
||||
"@llamaindex/mongodb": "^0.0.5",
|
||||
},
|
||||
none: {},
|
||||
pg: {
|
||||
pg: "^8.12.0",
|
||||
pgvector: "^0.2.0",
|
||||
"@llamaindex/postgres": "^0.0.33",
|
||||
},
|
||||
pinecone: {
|
||||
"@llamaindex/pinecone": "^0.0.5",
|
||||
},
|
||||
qdrant: {
|
||||
"@qdrant/js-client-rest": "^1.11.0",
|
||||
"@llamaindex/qdrant": "^0.1.0",
|
||||
},
|
||||
weaviate: {
|
||||
"@llamaindex/weaviate": "^0.0.5",
|
||||
},
|
||||
};
|
||||
|
||||
async function updatePackageJson({
|
||||
root,
|
||||
appName,
|
||||
@@ -258,6 +451,9 @@ async function updatePackageJson({
|
||||
ui,
|
||||
observability,
|
||||
vectorDb,
|
||||
backend,
|
||||
modelConfig,
|
||||
template,
|
||||
}: Pick<
|
||||
InstallTemplateArgs,
|
||||
| "root"
|
||||
@@ -267,8 +463,11 @@ async function updatePackageJson({
|
||||
| "ui"
|
||||
| "observability"
|
||||
| "vectorDb"
|
||||
| "modelConfig"
|
||||
| "template"
|
||||
> & {
|
||||
relativeEngineDestPath: string;
|
||||
backend: boolean;
|
||||
}): Promise<any> {
|
||||
const packageJsonFile = path.join(root, "package.json");
|
||||
const packageJson: any = JSON.parse(
|
||||
@@ -277,7 +476,7 @@ async function updatePackageJson({
|
||||
packageJson.name = appName;
|
||||
packageJson.version = "0.1.0";
|
||||
|
||||
if (relativeEngineDestPath) {
|
||||
if (relativeEngineDestPath && template !== "llamaindexserver") {
|
||||
// TODO: move script to {root}/scripts for all frameworks
|
||||
// add generate script if using context engine
|
||||
packageJson.scripts = {
|
||||
@@ -308,32 +507,25 @@ async function updatePackageJson({
|
||||
};
|
||||
}
|
||||
|
||||
if (vectorDb === "pg") {
|
||||
if (backend) {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
pg: "^8.12.0",
|
||||
pgvector: "^0.2.0",
|
||||
"@llamaindex/readers": "^2.0.0",
|
||||
};
|
||||
}
|
||||
|
||||
if (vectorDb === "qdrant") {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
"@qdrant/js-client-rest": "^1.11.0",
|
||||
};
|
||||
}
|
||||
if (vectorDb === "mongo") {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
mongodb: "^6.7.0",
|
||||
};
|
||||
}
|
||||
if (vectorDb && vectorDb in vectorDbDependencies) {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
...vectorDbDependencies[vectorDb],
|
||||
};
|
||||
}
|
||||
|
||||
if (vectorDb === "milvus") {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
"@zilliz/milvus2-sdk-node": "^2.4.6",
|
||||
};
|
||||
if (modelConfig.provider && modelConfig.provider in providerDependencies) {
|
||||
packageJson.dependencies = {
|
||||
...packageJson.dependencies,
|
||||
...providerDependencies[modelConfig.provider],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if (observability === "traceloop") {
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
# LlamaIndex Server
|
||||
|
||||
LlamaIndexServer is a FastAPI-based application that allows you to quickly launch your [LlamaIndex Workflows](https://docs.llamaindex.ai/en/stable/module_guides/workflow/#workflows) and [Agent Workflows](https://docs.llamaindex.ai/en/stable/understanding/agent/multi_agent/) as an API server with an optional chat UI. It provides a complete environment for running LlamaIndex workflows with both API endpoints and a user interface for interaction.
|
||||
|
||||
## Features
|
||||
|
||||
- Serving a workflow as a chatbot
|
||||
- Built on FastAPI for high performance and easy API development
|
||||
- Optional built-in chat UI
|
||||
- Prebuilt development code
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install llama-index-server
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
# main.py
|
||||
from llama_index.core.agent.workflow import AgentWorkflow
|
||||
from llama_index.core.workflow import Workflow
|
||||
from llama_index.core.tools import FunctionTool
|
||||
from llama_index.server import LlamaIndexServer
|
||||
|
||||
|
||||
# Define a factory function that returns a Workflow or AgentWorkflow
|
||||
def create_workflow() -> Workflow:
|
||||
def fetch_weather(city: str) -> str:
|
||||
return f"The weather in {city} is sunny"
|
||||
|
||||
return AgentWorkflow.from_tools(
|
||||
tools=[
|
||||
FunctionTool.from_defaults(
|
||||
fn=fetch_weather,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
# Create an API server for the workflow
|
||||
app = LlamaIndexServer(
|
||||
workflow_factory=create_workflow, # Supports Workflow or AgentWorkflow
|
||||
env="dev", # Enable development mode
|
||||
include_ui=True, # Include chat UI
|
||||
starter_questions=["What can you do?", "How do I use this?"],
|
||||
verbose=True
|
||||
)
|
||||
```
|
||||
|
||||
## Running the Server
|
||||
|
||||
- In the same directory as `main.py`, run the following command to start the server:
|
||||
|
||||
```bash
|
||||
fastapi dev
|
||||
```
|
||||
|
||||
- Making a request to the server:
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/api/chat" -H "Content-Type: application/json" -d '{"message": "What is the weather in Tokyo?"}'
|
||||
```
|
||||
|
||||
- See the API documentation at `http://localhost:8000/docs`
|
||||
- Access the chat UI at `http://localhost:8000/` (Make sure you set the `env="dev"` or `include_ui=True` in the server configuration)
|
||||
|
||||
## Configuration Options
|
||||
|
||||
The LlamaIndexServer accepts the following configuration parameters:
|
||||
|
||||
- `workflow_factory`: A callable that creates a workflow instance for each request
|
||||
- `logger`: Optional logger instance (defaults to uvicorn logger)
|
||||
- `use_default_routers`: Whether to include default routers (chat, static file serving)
|
||||
- `env`: Environment setting ('dev' enables CORS and UI by default)
|
||||
- `include_ui`: Whether to include the chat UI
|
||||
- `starter_questions`: List of starter questions for the chat UI
|
||||
- `verbose`: Enable verbose logging
|
||||
- `api_prefix`: API route prefix (default: "/api")
|
||||
- `server_url`: The deployment URL of the server (default is None)
|
||||
- `ui_path`: Path for downloaded UI static files (default: ".ui")
|
||||
|
||||
## Default Routers and Features
|
||||
|
||||
### Chat Router
|
||||
|
||||
The server includes a default chat router at `/api/chat` for handling chat interactions.
|
||||
|
||||
### Static File Serving
|
||||
|
||||
- The server automatically mounts the `data` and `output` folders at `{server_url}{api_prefix}/files/data` (default: `/api/files/data`) and `{server_url}{api_prefix}/files/output` (default: `/api/files/output`) respectively.
|
||||
- Your workflows can use both folders to store and access files. As a convention, the `data` folder is used for documents that are ingested and the `output` folder is used for documents that are generated by the workflow.
|
||||
- The example workflows from `create-llama` (see below) are following this pattern.
|
||||
|
||||
### Chat UI
|
||||
|
||||
When enabled, the server provides a chat interface at the root path (`/`) with:
|
||||
|
||||
- Configurable starter questions
|
||||
- Real-time chat interface
|
||||
- API endpoint integration
|
||||
|
||||
## Development Mode
|
||||
|
||||
In development mode (`env="dev"`), the server:
|
||||
|
||||
- Enables CORS for all origins
|
||||
- Automatically includes the chat UI
|
||||
- Provides more verbose logging
|
||||
|
||||
## API Endpoints
|
||||
|
||||
The server provides the following default endpoints:
|
||||
|
||||
- `/api/chat`: Chat interaction endpoint
|
||||
- `/api/files/data/*`: Access to data directory files
|
||||
- `/api/files/output/*`: Access to output directory files
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. Always provide a workflow factory that creates fresh workflow instances
|
||||
2. Use environment variables for sensitive configuration
|
||||
3. Enable verbose logging during development
|
||||
4. Configure CORS appropriately for your deployment environment
|
||||
5. Use starter questions to guide users in the chat UI
|
||||
|
||||
## Getting Started with a New Project
|
||||
|
||||
Want to start a new project with LlamaIndexServer? Check out our [create-llama](https://github.com/run-llama/create-llama) tool to quickly generate a new project with LlamaIndexServer.
|
||||
@@ -0,0 +1,3 @@
|
||||
from .server import LlamaIndexServer
|
||||
|
||||
__all__ = ["LlamaIndexServer"]
|
||||
@@ -0,0 +1,13 @@
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
from llama_index.server.api.callbacks.llamacloud import LlamaCloudFileDownload
|
||||
from llama_index.server.api.callbacks.source_nodes import SourceNodesFromToolCall
|
||||
from llama_index.server.api.callbacks.suggest_next_questions import (
|
||||
SuggestNextQuestions,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"EventCallback",
|
||||
"SourceNodesFromToolCall",
|
||||
"SuggestNextQuestions",
|
||||
"LlamaCloudFileDownload",
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class EventCallback(ABC):
|
||||
"""
|
||||
Base class for event callbacks during event streaming.
|
||||
"""
|
||||
|
||||
async def run(self, event: Any) -> Any:
|
||||
"""
|
||||
Called for each event in the stream.
|
||||
Default behavior: pass through the event unchanged.
|
||||
"""
|
||||
return event
|
||||
|
||||
async def on_complete(self, final_response: str) -> Any:
|
||||
"""
|
||||
Called when the stream is complete.
|
||||
Default behavior: return None.
|
||||
"""
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def from_default(self, *args: Any, **kwargs: Any) -> "EventCallback":
|
||||
"""
|
||||
Create a new instance of the processor from default values.
|
||||
"""
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from typing import Any, List
|
||||
|
||||
from fastapi import BackgroundTasks
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
from llama_index.server.services.llamacloud.file import LlamaCloudFileService
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class LlamaCloudFileDownload(EventCallback):
|
||||
"""
|
||||
Processor for handling LlamaCloud file downloads from source nodes.
|
||||
"""
|
||||
|
||||
def __init__(self, background_tasks: BackgroundTasks) -> None:
|
||||
self.background_tasks = background_tasks
|
||||
|
||||
async def run(self, event: Any) -> Any:
|
||||
if hasattr(event, "to_response"):
|
||||
event_response = event.to_response()
|
||||
if event_response.get("type") == "sources" and hasattr(event, "nodes"):
|
||||
await self._process_response_nodes(event.nodes)
|
||||
return event
|
||||
|
||||
async def _process_response_nodes(self, source_nodes: List[NodeWithScore]) -> None:
|
||||
try:
|
||||
LlamaCloudFileService.download_files_from_nodes(
|
||||
source_nodes, self.background_tasks
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_default(
|
||||
cls, background_tasks: BackgroundTasks
|
||||
) -> "LlamaCloudFileDownload":
|
||||
return cls(background_tasks=background_tasks)
|
||||
@@ -0,0 +1,32 @@
|
||||
from typing import Any
|
||||
|
||||
from llama_index.core.agent.workflow.workflow_events import ToolCallResult
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
from llama_index.server.api.models import SourceNodesEvent
|
||||
|
||||
|
||||
class SourceNodesFromToolCall(EventCallback):
|
||||
"""
|
||||
Extract source nodes from the query tool output.
|
||||
|
||||
Args:
|
||||
query_tool_name: The name of the tool that queries the index.
|
||||
default is "query_index"
|
||||
"""
|
||||
|
||||
def __init__(self, query_tool_name: str = "query_index"):
|
||||
self.query_tool_name = query_tool_name
|
||||
|
||||
def transform_tool_call_result(self, event: ToolCallResult) -> SourceNodesEvent:
|
||||
source_nodes = event.tool_output.raw_output.source_nodes
|
||||
return SourceNodesEvent(nodes=source_nodes)
|
||||
|
||||
async def run(self, event: Any) -> Any:
|
||||
if isinstance(event, ToolCallResult):
|
||||
if event.tool_name == self.query_tool_name:
|
||||
return event, self.transform_tool_call_result(event)
|
||||
return event
|
||||
|
||||
@classmethod
|
||||
def from_default(cls, *args: Any, **kwargs: Any) -> "SourceNodesFromToolCall":
|
||||
return cls()
|
||||
@@ -0,0 +1,69 @@
|
||||
import logging
|
||||
from typing import Any, AsyncGenerator, List, Optional
|
||||
|
||||
from llama_index.core.workflow.handler import WorkflowHandler
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""
|
||||
Streams events from a workflow handler through a chain of callbacks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow_handler: WorkflowHandler,
|
||||
callbacks: Optional[List[EventCallback]] = None,
|
||||
):
|
||||
self.workflow_handler = workflow_handler
|
||||
self.callbacks = callbacks or []
|
||||
self.accumulated_text = ""
|
||||
|
||||
async def cancel_run(self) -> None:
|
||||
"""Cancel the workflow handler."""
|
||||
await self.workflow_handler.cancel_run()
|
||||
|
||||
async def stream_events(self) -> AsyncGenerator[Any, None]:
|
||||
"""Stream events through the processor chain."""
|
||||
try:
|
||||
async for event in self.workflow_handler.stream_events():
|
||||
events_to_process = [event]
|
||||
for callback in self.callbacks:
|
||||
next_events: list[Any] = []
|
||||
for evt in events_to_process:
|
||||
callback_output = await callback.run(evt)
|
||||
if isinstance(callback_output, (list, tuple)):
|
||||
next_events.extend(callback_output)
|
||||
elif callback_output is not None:
|
||||
next_events.append(callback_output)
|
||||
events_to_process = next_events
|
||||
|
||||
# Yield all processed events
|
||||
for evt in events_to_process:
|
||||
yield evt
|
||||
|
||||
# After all events are processed, call on_complete for each callback
|
||||
for callback in self.callbacks:
|
||||
result = await callback.on_complete(self.accumulated_text)
|
||||
if result:
|
||||
yield result
|
||||
|
||||
except Exception:
|
||||
# Make sure to cancel the workflow on error
|
||||
await self.workflow_handler.cancel_run()
|
||||
raise
|
||||
|
||||
def accumulate_text(self, text: str) -> None:
|
||||
"""Accumulate text from the workflow handler."""
|
||||
self.accumulated_text += text
|
||||
|
||||
@classmethod
|
||||
def from_default(
|
||||
cls,
|
||||
handler: WorkflowHandler,
|
||||
callbacks: Optional[List[EventCallback]] = None,
|
||||
) -> "StreamHandler":
|
||||
"""Create a new instance with the given workflow handler and callbacks."""
|
||||
return cls(workflow_handler=handler, callbacks=callbacks)
|
||||
@@ -0,0 +1,45 @@
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
from llama_index.server.api.models import ChatRequest
|
||||
from llama_index.server.services.suggest_next_question import (
|
||||
SuggestNextQuestionsService,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class SuggestNextQuestions(EventCallback):
|
||||
"""Processor for generating next question suggestions."""
|
||||
|
||||
def __init__(
|
||||
self, chat_request: ChatRequest, logger: Optional[logging.Logger] = None
|
||||
):
|
||||
self.chat_request = chat_request
|
||||
self.accumulated_text = ""
|
||||
if logger:
|
||||
self.logger = logger
|
||||
else:
|
||||
self.logger = logging.getLogger("uvicorn")
|
||||
|
||||
async def on_complete(self, final_response: str) -> Any:
|
||||
if final_response == "":
|
||||
self.logger.warning(
|
||||
"SuggestNextQuestions is enabled but final response is empty, make sure your content generator accumulates text"
|
||||
)
|
||||
return None
|
||||
|
||||
questions = await SuggestNextQuestionsService.run(
|
||||
self.chat_request.messages, final_response
|
||||
)
|
||||
if questions:
|
||||
return {
|
||||
"type": "suggested_questions",
|
||||
"data": questions,
|
||||
}
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_default(cls, chat_request: ChatRequest) -> "SuggestNextQuestions":
|
||||
return cls(chat_request=chat_request)
|
||||
@@ -0,0 +1,136 @@
|
||||
import logging
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
from llama_index.core.types import ChatMessage, MessageRole
|
||||
from llama_index.core.workflow import Event
|
||||
from llama_index.server.settings import server_settings
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class ChatConfig(BaseModel):
|
||||
next_question_suggestions: bool = Field(
|
||||
default=True,
|
||||
description="Whether to suggest next questions",
|
||||
)
|
||||
|
||||
|
||||
class ChatAPIMessage(BaseModel):
|
||||
role: MessageRole
|
||||
content: str
|
||||
|
||||
def to_llamaindex_message(self) -> ChatMessage:
|
||||
return ChatMessage(role=self.role, content=self.content)
|
||||
|
||||
|
||||
class ChatRequest(BaseModel):
|
||||
messages: List[ChatAPIMessage]
|
||||
data: Optional[Any] = None
|
||||
config: Optional[ChatConfig] = ChatConfig()
|
||||
|
||||
@field_validator("messages")
|
||||
def validate_messages(cls, v: List[ChatAPIMessage]) -> List[ChatAPIMessage]:
|
||||
if v[-1].role != MessageRole.USER:
|
||||
raise ValueError("Last message must be from user")
|
||||
return v
|
||||
|
||||
|
||||
class AgentRunEventType(Enum):
|
||||
TEXT = "text"
|
||||
PROGRESS = "progress"
|
||||
|
||||
|
||||
class AgentRunEvent(Event):
|
||||
name: str
|
||||
msg: str
|
||||
event_type: AgentRunEventType = AgentRunEventType.TEXT
|
||||
data: Optional[dict] = None
|
||||
|
||||
def to_response(self) -> dict:
|
||||
return {
|
||||
"type": "agent",
|
||||
"data": {
|
||||
"agent": self.name,
|
||||
"type": self.event_type.value,
|
||||
"text": self.msg,
|
||||
"data": self.data,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class SourceNodesEvent(Event):
|
||||
nodes: List[NodeWithScore]
|
||||
|
||||
def to_response(self) -> dict:
|
||||
return {
|
||||
"type": "sources",
|
||||
"data": {
|
||||
"nodes": [
|
||||
SourceNodes.from_source_node(node).model_dump()
|
||||
for node in self.nodes
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
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) -> "SourceNodes":
|
||||
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],
|
||||
data_dir: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
url_prefix = server_settings.file_server_url_prefix
|
||||
if data_dir is None:
|
||||
data_dir = "data"
|
||||
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]
|
||||
) -> List["SourceNodes"]:
|
||||
return [cls.from_source_node(node) for node in source_nodes]
|
||||
@@ -0,0 +1,140 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
from typing import AsyncGenerator, Callable, Union
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from llama_index.core.agent.workflow.workflow_events import AgentStream
|
||||
from llama_index.core.workflow import StopEvent, Workflow
|
||||
from llama_index.server.api.callbacks import (
|
||||
SourceNodesFromToolCall,
|
||||
SuggestNextQuestions,
|
||||
)
|
||||
from llama_index.server.api.callbacks.base import EventCallback
|
||||
from llama_index.server.api.callbacks.llamacloud import LlamaCloudFileDownload
|
||||
from llama_index.server.api.callbacks.stream_handler import StreamHandler
|
||||
from llama_index.server.api.models import ChatRequest
|
||||
from llama_index.server.api.utils.vercel_stream import VercelStreamResponse
|
||||
from llama_index.server.services.llamacloud import LlamaCloudFileService
|
||||
|
||||
|
||||
def chat_router(
|
||||
workflow_factory: Callable[..., Workflow],
|
||||
logger: logging.Logger,
|
||||
) -> APIRouter:
|
||||
router = APIRouter(prefix="/chat")
|
||||
|
||||
@router.post("")
|
||||
async def chat(
|
||||
request: ChatRequest,
|
||||
background_tasks: BackgroundTasks,
|
||||
) -> StreamingResponse:
|
||||
try:
|
||||
user_message = request.messages[-1].to_llamaindex_message()
|
||||
chat_history = [
|
||||
message.to_llamaindex_message() for message in request.messages[:-1]
|
||||
]
|
||||
# detect if the workflow factory has chat_request as a parameter
|
||||
factory_sig = inspect.signature(workflow_factory)
|
||||
if "chat_request" in factory_sig.parameters:
|
||||
workflow = workflow_factory(chat_request=request)
|
||||
else:
|
||||
workflow = workflow_factory()
|
||||
workflow_handler = workflow.run(
|
||||
user_msg=user_message.content,
|
||||
chat_history=chat_history,
|
||||
)
|
||||
|
||||
callbacks: list[EventCallback] = [
|
||||
SourceNodesFromToolCall(),
|
||||
LlamaCloudFileDownload(background_tasks),
|
||||
]
|
||||
if request.config and request.config.next_question_suggestions:
|
||||
callbacks.append(SuggestNextQuestions(request))
|
||||
stream_handler = StreamHandler(
|
||||
workflow_handler=workflow_handler,
|
||||
callbacks=callbacks,
|
||||
)
|
||||
|
||||
return VercelStreamResponse(
|
||||
content_generator=_stream_content(stream_handler, request, logger),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
if LlamaCloudFileService.is_configured():
|
||||
|
||||
@router.get("/config/llamacloud")
|
||||
async def chat_llama_cloud_config() -> dict:
|
||||
if not os.getenv("LLAMA_CLOUD_API_KEY"):
|
||||
raise HTTPException(
|
||||
status_code=500, detail="LlamaCloud API KEY is not configured"
|
||||
)
|
||||
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,
|
||||
}
|
||||
|
||||
return router
|
||||
|
||||
|
||||
async def _stream_content(
|
||||
handler: StreamHandler,
|
||||
request: ChatRequest,
|
||||
logger: logging.Logger,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
async def _text_stream(
|
||||
event: Union[AgentStream, StopEvent],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
if isinstance(event, AgentStream):
|
||||
yield event.delta
|
||||
elif isinstance(event, StopEvent):
|
||||
if isinstance(event.result, str):
|
||||
yield event.result
|
||||
elif isinstance(event.result, AsyncGenerator):
|
||||
async for chunk in event.result:
|
||||
if isinstance(chunk, str):
|
||||
yield chunk
|
||||
elif hasattr(chunk, "delta") and chunk.delta:
|
||||
yield chunk.delta
|
||||
|
||||
stream_started = False
|
||||
try:
|
||||
async for event in handler.stream_events():
|
||||
if not stream_started:
|
||||
# Start the stream with an empty message
|
||||
stream_started = True
|
||||
yield VercelStreamResponse.convert_text("")
|
||||
|
||||
# Handle different types of events
|
||||
if isinstance(event, (AgentStream, StopEvent)):
|
||||
async for chunk in _text_stream(event):
|
||||
handler.accumulate_text(chunk)
|
||||
yield VercelStreamResponse.convert_text(chunk)
|
||||
elif isinstance(event, dict):
|
||||
yield VercelStreamResponse.convert_data(event)
|
||||
elif hasattr(event, "to_response"):
|
||||
event_response = event.to_response()
|
||||
yield VercelStreamResponse.convert_data(event_response)
|
||||
else:
|
||||
yield VercelStreamResponse.convert_data(event.model_dump())
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.warning("Client cancelled the request!")
|
||||
await handler.cancel_run()
|
||||
except Exception as e:
|
||||
logger.error(f"Error in stream response: {e}")
|
||||
yield VercelStreamResponse.convert_error(str(e))
|
||||
await handler.cancel_run()
|
||||
@@ -0,0 +1,44 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, AsyncGenerator, Union
|
||||
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class VercelStreamResponse(StreamingResponse):
|
||||
"""
|
||||
Converts preprocessed events into Vercel-compatible streaming response format.
|
||||
"""
|
||||
|
||||
TEXT_PREFIX = "0:"
|
||||
DATA_PREFIX = "8:"
|
||||
ERROR_PREFIX = "3:"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
content_generator: AsyncGenerator[str, None],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(content_generator, *args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def convert_text(cls, token: str) -> str:
|
||||
"""Convert text event to Vercel format."""
|
||||
# Escape newlines and double quotes to avoid breaking the stream
|
||||
token = json.dumps(token)
|
||||
return f"{cls.TEXT_PREFIX}{token}\n"
|
||||
|
||||
@classmethod
|
||||
def convert_data(cls, data: Union[dict, str]) -> str:
|
||||
"""Convert data event to Vercel format."""
|
||||
data_str = json.dumps(data) if isinstance(data, dict) else data
|
||||
return f"{cls.DATA_PREFIX}[{data_str}]\n"
|
||||
|
||||
@classmethod
|
||||
def convert_error(cls, error: str) -> str:
|
||||
"""Convert error event to Vercel format."""
|
||||
error_str = json.dumps(error)
|
||||
return f"{cls.ERROR_PREFIX}{error_str}\n"
|
||||
@@ -0,0 +1,55 @@
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
CHAT_UI_VERSION = "0.0.6"
|
||||
|
||||
|
||||
def download_chat_ui(
|
||||
logger: Optional[logging.Logger] = None, target_path: str = ".ui"
|
||||
) -> None:
|
||||
if logger is None:
|
||||
logger = logging.getLogger("uvicorn")
|
||||
path = Path(target_path)
|
||||
temp_dir = _download_package(_get_download_link(CHAT_UI_VERSION))
|
||||
_copy_ui_files(temp_dir, path)
|
||||
logger.info("Chat UI downloaded and copied to static folder")
|
||||
|
||||
|
||||
def _get_download_link(version: str) -> str:
|
||||
"""Get the download link for the chat UI from the npm registry."""
|
||||
return f"https://registry.npmjs.org/@llamaindex/server/-/server-{version}.tgz"
|
||||
|
||||
|
||||
def _download_package(url: str) -> Path:
|
||||
"""Download tar.gz file and extract all files into a temporary directory."""
|
||||
import io
|
||||
import tarfile
|
||||
import tempfile
|
||||
|
||||
response = requests.get(url, headers={"User-Agent": "Mozilla/5.0"})
|
||||
content = response.content
|
||||
|
||||
temp_dir = Path(tempfile.mkdtemp())
|
||||
|
||||
with tarfile.open(fileobj=io.BytesIO(content), mode="r:gz") as tar:
|
||||
tar.extractall(path=temp_dir)
|
||||
|
||||
return temp_dir
|
||||
|
||||
|
||||
def _copy_ui_files(temp_dir: Path, target_path: Path) -> None:
|
||||
"""Copy files from the .next directory to the static directory."""
|
||||
target_path.mkdir(parents=True, exist_ok=True)
|
||||
next_dir = temp_dir / "package/dist/static"
|
||||
|
||||
if next_dir.exists():
|
||||
for item in next_dir.iterdir():
|
||||
dest = target_path / item.name
|
||||
if item.is_dir():
|
||||
shutil.copytree(item, dest, dirs_exist_ok=True)
|
||||
else:
|
||||
shutil.copy2(item, dest)
|
||||
@@ -0,0 +1,184 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from llama_index.core.workflow import Workflow
|
||||
from llama_index.server.api.routers.chat import chat_router
|
||||
from llama_index.server.chat_ui import download_chat_ui
|
||||
from llama_index.server.settings import server_settings
|
||||
|
||||
|
||||
class LlamaIndexServer(FastAPI):
|
||||
workflow_factory: Callable[..., Workflow]
|
||||
include_ui: Optional[bool]
|
||||
starter_questions: Optional[list[str]]
|
||||
verbose: bool = False
|
||||
ui_path: str = ".ui"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow_factory: Callable[..., Workflow],
|
||||
logger: Optional[logging.Logger] = None,
|
||||
use_default_routers: Optional[bool] = True,
|
||||
env: Optional[str] = None,
|
||||
include_ui: Optional[bool] = None,
|
||||
starter_questions: Optional[list[str]] = None,
|
||||
server_url: Optional[str] = None,
|
||||
api_prefix: Optional[str] = None,
|
||||
verbose: bool = False,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
Initialize the LlamaIndexServer.
|
||||
|
||||
Args:
|
||||
workflow_factory: A factory function that creates a workflow instance for each request.
|
||||
logger: The logger to use.
|
||||
use_default_routers: Whether to use the default routers (chat, mount `data` and `output` directories).
|
||||
env: The environment to run the server in.
|
||||
include_ui: Whether to show an chat UI in the root path.
|
||||
starter_questions: A list of starter questions to display in the chat UI.
|
||||
server_url: The URL of the server.
|
||||
api_prefix: The prefix for the API endpoints.
|
||||
verbose: Whether to show verbose logs.
|
||||
"""
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self.workflow_factory = workflow_factory
|
||||
self.logger = logger or logging.getLogger("uvicorn")
|
||||
self.verbose = verbose
|
||||
self.include_ui = include_ui # Store the explicitly passed value first
|
||||
self.starter_questions = starter_questions
|
||||
self.use_default_routers = use_default_routers or True
|
||||
|
||||
# Update the settings
|
||||
if server_url:
|
||||
server_settings.set_url(server_url)
|
||||
if api_prefix:
|
||||
server_settings.set_api_prefix(api_prefix)
|
||||
|
||||
if self.use_default_routers:
|
||||
self.add_default_routers()
|
||||
|
||||
if str(env).lower() == "dev":
|
||||
self.allow_cors("*")
|
||||
if self.include_ui is None:
|
||||
self.include_ui = True
|
||||
if self.include_ui is None:
|
||||
self.include_ui = False
|
||||
|
||||
if self.include_ui:
|
||||
self.mount_ui()
|
||||
|
||||
@property
|
||||
def _ui_config(self) -> dict:
|
||||
config = {
|
||||
"CHAT_API": f"{server_settings.api_url}/chat",
|
||||
"STARTER_QUESTIONS": self.starter_questions,
|
||||
}
|
||||
is_llamacloud_configured = os.getenv("LLAMA_CLOUD_API_KEY") is not None
|
||||
if is_llamacloud_configured:
|
||||
config["LLAMA_CLOUD_API"] = (
|
||||
f"{server_settings.api_url}/chat/config/llamacloud"
|
||||
)
|
||||
return config
|
||||
|
||||
# Default routers
|
||||
def add_default_routers(self) -> None:
|
||||
self.add_chat_router()
|
||||
self.mount_data_dir()
|
||||
self.mount_output_dir()
|
||||
|
||||
def add_chat_router(self) -> None:
|
||||
"""
|
||||
Add the chat router.
|
||||
"""
|
||||
self.include_router(
|
||||
chat_router(
|
||||
self.workflow_factory,
|
||||
self.logger,
|
||||
),
|
||||
prefix=server_settings.api_prefix,
|
||||
)
|
||||
|
||||
def mount_ui(self) -> None:
|
||||
"""
|
||||
Mount the UI.
|
||||
"""
|
||||
# Check if the static folder exists
|
||||
if self.include_ui:
|
||||
if not os.path.exists(self.ui_path):
|
||||
self.logger.warning(
|
||||
f"UI files not found, downloading UI to {self.ui_path}"
|
||||
)
|
||||
download_chat_ui(logger=self.logger, target_path=self.ui_path)
|
||||
self._mount_static_files(directory=self.ui_path, path="/", html=True)
|
||||
self._override_ui_config()
|
||||
|
||||
def _override_ui_config(self) -> None:
|
||||
"""
|
||||
Override the UI config by writing a complete configuration file.
|
||||
"""
|
||||
try:
|
||||
config_path = os.path.join(self.ui_path, "config.js")
|
||||
if not os.path.exists(config_path):
|
||||
self.logger.error("Config file not found")
|
||||
return
|
||||
config_content = (
|
||||
f"window.LLAMAINDEX = {json.dumps(self._ui_config, indent=2)};"
|
||||
)
|
||||
with open(config_path, "w") as f:
|
||||
f.write(config_content)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error overriding UI config: {e}")
|
||||
|
||||
def mount_data_dir(self, data_dir: str = "data") -> None:
|
||||
"""
|
||||
Mount the data directory.
|
||||
"""
|
||||
self._mount_static_files(
|
||||
directory=data_dir,
|
||||
path=f"{server_settings.api_prefix}/files/data",
|
||||
html=True,
|
||||
)
|
||||
|
||||
def mount_output_dir(self, output_dir: str = "output") -> None:
|
||||
"""
|
||||
Mount the output directory.
|
||||
"""
|
||||
self._mount_static_files(
|
||||
directory=output_dir,
|
||||
path=f"{server_settings.api_prefix}/files/output",
|
||||
html=True,
|
||||
)
|
||||
|
||||
def _mount_static_files(
|
||||
self, directory: str, path: str, html: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Mount static files from a directory if it exists.
|
||||
"""
|
||||
if os.path.exists(directory):
|
||||
self.logger.info(f"Mounting static files '{directory}' at '{path}'")
|
||||
self.mount(
|
||||
path,
|
||||
StaticFiles(directory=directory, check_dir=False, html=html),
|
||||
name=f"{directory}-static",
|
||||
)
|
||||
|
||||
def allow_cors(self, origin: str = "*") -> None:
|
||||
"""
|
||||
Allow CORS for a specific origin.
|
||||
"""
|
||||
self.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[origin],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
@@ -0,0 +1,117 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from llama_index.server.settings import server_settings
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRIVATE_STORE_PATH = str(Path("output", "uploaded"))
|
||||
TOOL_STORE_PATH = str(Path("output", "tools"))
|
||||
LLAMA_CLOUD_STORE_PATH = str(Path("output", "llamacloud"))
|
||||
|
||||
|
||||
class DocumentFile(BaseModel):
|
||||
id: str
|
||||
name: str # Stored file name
|
||||
type: Optional[str] = None
|
||||
size: Optional[int] = None
|
||||
url: Optional[str] = None
|
||||
path: Optional[str] = Field(
|
||||
None,
|
||||
description="The stored file path. Used internally in the server.",
|
||||
exclude=True,
|
||||
)
|
||||
refs: Optional[List[str]] = Field(
|
||||
None, description="The document ids in the index."
|
||||
)
|
||||
|
||||
|
||||
class FileService:
|
||||
"""
|
||||
To store the files uploaded by the user.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def save_file(
|
||||
cls,
|
||||
content: Union[bytes, str],
|
||||
file_name: str,
|
||||
save_dir: Optional[str] = None,
|
||||
) -> DocumentFile:
|
||||
"""
|
||||
Save the content to a file in the local file server (accessible via URL).
|
||||
|
||||
Args:
|
||||
content (bytes | str): The content to save, either bytes or string.
|
||||
file_name (str): The original name of the file.
|
||||
save_dir (Optional[str]): The relative path from the current working directory. Defaults to the `output/uploaded` directory.
|
||||
|
||||
Returns:
|
||||
The metadata of the saved file.
|
||||
"""
|
||||
if save_dir is None:
|
||||
save_dir = os.path.join("output", "uploaded")
|
||||
|
||||
file_id = str(uuid.uuid4())
|
||||
name, extension = os.path.splitext(file_name)
|
||||
extension = extension.lstrip(".")
|
||||
sanitized_name = _sanitize_file_name(name)
|
||||
if extension == "":
|
||||
raise ValueError("File is not supported!")
|
||||
new_file_name = f"{sanitized_name}_{file_id}.{extension}"
|
||||
|
||||
file_path = os.path.join(save_dir, new_file_name)
|
||||
|
||||
if isinstance(content, str):
|
||||
content = content.encode()
|
||||
|
||||
try:
|
||||
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
||||
with open(file_path, "wb") as file:
|
||||
file.write(content)
|
||||
except PermissionError as e:
|
||||
logger.error(f"Permission denied when writing to file {file_path}: {e!s}")
|
||||
raise
|
||||
except OSError as e:
|
||||
logger.error(f"IO error occurred when writing to file {file_path}: {e!s}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error when writing to file {file_path}: {e!s}")
|
||||
raise
|
||||
|
||||
logger.info(f"Saved file to {file_path}")
|
||||
|
||||
file_size = os.path.getsize(file_path)
|
||||
file_url = (
|
||||
f"{server_settings.file_server_url_prefix}/{save_dir}/{new_file_name}"
|
||||
)
|
||||
return DocumentFile(
|
||||
id=file_id,
|
||||
name=new_file_name,
|
||||
type=extension,
|
||||
size=file_size,
|
||||
path=file_path,
|
||||
url=file_url,
|
||||
refs=None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_file_url(cls, file_name: str, save_dir: Optional[str] = None) -> str:
|
||||
"""
|
||||
Get the URL of a file.
|
||||
"""
|
||||
if save_dir is None:
|
||||
save_dir = os.path.join("output", "uploaded")
|
||||
return f"{server_settings.file_server_url_prefix}/{save_dir}/{file_name}"
|
||||
|
||||
|
||||
def _sanitize_file_name(file_name: str) -> str:
|
||||
"""
|
||||
Sanitize the file name by replacing all non-alphanumeric characters with underscores.
|
||||
"""
|
||||
return re.sub(r"[^a-zA-Z0-9.]", "_", file_name)
|
||||
@@ -0,0 +1,11 @@
|
||||
from .file import LlamaCloudFileService
|
||||
from .generate import load_to_llamacloud
|
||||
from .index import LlamaCloudIndex, get_client, get_index
|
||||
|
||||
__all__ = [
|
||||
"LlamaCloudFileService",
|
||||
"LlamaCloudIndex",
|
||||
"get_client",
|
||||
"get_index",
|
||||
"load_to_llamacloud",
|
||||
]
|
||||
@@ -0,0 +1,184 @@
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import typing
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
import requests
|
||||
from fastapi import BackgroundTasks
|
||||
from llama_cloud import ManagedIngestionStatus, PipelineFileCreateCustomMetadataValue
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
from llama_index.server.api.models import SourceNodes
|
||||
from llama_index.server.services.llamacloud.index import get_client
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class LlamaCloudFile(BaseModel):
|
||||
file_name: str
|
||||
pipeline_id: str
|
||||
|
||||
def __eq__(self, other: Any) -> bool:
|
||||
if not isinstance(other, LlamaCloudFile):
|
||||
return NotImplemented
|
||||
return (
|
||||
self.file_name == other.file_name and self.pipeline_id == other.pipeline_id
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
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]],
|
||||
wait_for_processing: bool = True,
|
||||
) -> str:
|
||||
client = get_client()
|
||||
file = client.files.upload_file(project_id=project_id, upload_file=upload_file)
|
||||
file_id = file.id
|
||||
files = [
|
||||
{
|
||||
"file_id": file_id,
|
||||
"custom_metadata": {"file_id": file_id, **(custom_metadata or {})},
|
||||
}
|
||||
]
|
||||
files = client.pipelines.add_files_to_pipeline_api(pipeline_id, request=files)
|
||||
|
||||
if not wait_for_processing:
|
||||
return file_id
|
||||
|
||||
# Wait 2s for the file to be processed
|
||||
max_attempts = 20
|
||||
attempt = 0
|
||||
while attempt < max_attempts:
|
||||
result = client.pipelines.get_pipeline_file_status(
|
||||
file_id=file_id, pipeline_id=pipeline_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,
|
||||
) -> None:
|
||||
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
|
||||
) -> None:
|
||||
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(cls.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"), # type: ignore
|
||||
pipeline_id=node.metadata.get("pipeline_id"), # type: ignore
|
||||
)
|
||||
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) -> None:
|
||||
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")
|
||||
|
||||
@classmethod
|
||||
def is_configured(cls) -> bool:
|
||||
try:
|
||||
return os.environ.get("LLAMA_CLOUD_API_KEY") is not None
|
||||
except Exception:
|
||||
return False
|
||||
@@ -0,0 +1,56 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from llama_index.core.readers import SimpleDirectoryReader
|
||||
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
|
||||
from llama_index.server.services.llamacloud.file import LlamaCloudFileService
|
||||
|
||||
|
||||
def load_to_llamacloud(
|
||||
index: LlamaCloudIndex,
|
||||
data_dir: Optional[str] = None,
|
||||
recursive: Optional[bool] = None,
|
||||
logger: Optional[logging.Logger] = None,
|
||||
) -> None:
|
||||
if logger is None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger()
|
||||
|
||||
logger.info("Generate index for the provided data")
|
||||
|
||||
# use SimpleDirectoryReader to retrieve the files to process
|
||||
reader = SimpleDirectoryReader(
|
||||
data_dir or "data",
|
||||
recursive=recursive or True,
|
||||
)
|
||||
files_to_process = reader.input_files
|
||||
|
||||
# add each file to the LlamaCloud pipeline
|
||||
error_files = []
|
||||
for input_file in tqdm(
|
||||
files_to_process,
|
||||
desc="Processing files",
|
||||
unit="file",
|
||||
):
|
||||
with open(input_file, "rb") as f:
|
||||
logger.debug(
|
||||
f"Adding file {input_file} to pipeline {index.name} in project {index.project_name}"
|
||||
)
|
||||
try:
|
||||
LlamaCloudFileService.add_file_to_pipeline(
|
||||
index.project.id,
|
||||
index.pipeline.id,
|
||||
f,
|
||||
custom_metadata={},
|
||||
wait_for_processing=False,
|
||||
)
|
||||
except Exception as e:
|
||||
error_files.append(input_file)
|
||||
logger.error(f"Error adding file {input_file}: {e}")
|
||||
|
||||
if error_files:
|
||||
logger.error(f"Failed to add the following files: {error_files}")
|
||||
|
||||
logger.info("Finished generating the index")
|
||||
@@ -0,0 +1,163 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from llama_cloud import PipelineType
|
||||
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.core.settings import Settings
|
||||
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
|
||||
from llama_index.server.api.models import ChatRequest
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from llama_cloud.client import LlamaCloud
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class LlamaCloudConfig(BaseModel):
|
||||
# Private attributes
|
||||
api_key: str = Field(
|
||||
exclude=True, # Exclude from the model representation
|
||||
)
|
||||
base_url: Optional[str] = Field(
|
||||
exclude=True,
|
||||
)
|
||||
organization_id: Optional[str] = Field(
|
||||
exclude=True,
|
||||
)
|
||||
# Configuration attributes, can be set by the user
|
||||
pipeline: str = Field(
|
||||
description="The name of the pipeline to use",
|
||||
)
|
||||
project: str = Field(
|
||||
description="The name of the LlamaCloud project",
|
||||
)
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
if "api_key" not in kwargs:
|
||||
kwargs["api_key"] = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
if "base_url" not in kwargs:
|
||||
kwargs["base_url"] = os.getenv("LLAMA_CLOUD_BASE_URL")
|
||||
if "organization_id" not in kwargs:
|
||||
kwargs["organization_id"] = os.getenv("LLAMA_CLOUD_ORGANIZATION_ID")
|
||||
if "pipeline" not in kwargs:
|
||||
kwargs["pipeline"] = os.getenv("LLAMA_CLOUD_INDEX_NAME")
|
||||
if "project" not in kwargs:
|
||||
kwargs["project"] = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Validate and throw error if the env variables are not set before starting the app
|
||||
@field_validator("pipeline", "project", "api_key", mode="before")
|
||||
@classmethod
|
||||
def validate_fields(cls, value: Any) -> Any:
|
||||
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_factory=lambda: 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,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_default(cls, chat_request: Optional[ChatRequest] = None) -> "IndexConfig":
|
||||
default_config = cls()
|
||||
if chat_request is not None and chat_request.data is not None:
|
||||
llamacloud_config = chat_request.data.get("llamaCloudPipeline")
|
||||
if llamacloud_config is not None:
|
||||
default_config.llama_cloud_pipeline_config.pipeline = llamacloud_config[
|
||||
"pipeline"
|
||||
]
|
||||
default_config.llama_cloud_pipeline_config.project = llamacloud_config[
|
||||
"project"
|
||||
]
|
||||
return default_config
|
||||
|
||||
|
||||
def get_index(
|
||||
chat_request: Optional[ChatRequest] = None,
|
||||
create_if_missing: bool = False,
|
||||
) -> Optional[LlamaCloudIndex]:
|
||||
config = IndexConfig.from_default(chat_request)
|
||||
# Check whether the index exists
|
||||
try:
|
||||
index = LlamaCloudIndex(**config.to_index_kwargs())
|
||||
return index
|
||||
except ValueError:
|
||||
logger.warning("Index not found")
|
||||
if create_if_missing:
|
||||
logger.info("Creating index")
|
||||
_create_index(config)
|
||||
return LlamaCloudIndex(**config.to_index_kwargs())
|
||||
return None
|
||||
|
||||
|
||||
def get_client() -> "LlamaCloud":
|
||||
config = LlamaCloudConfig()
|
||||
return llama_cloud_get_client(**config.to_client_kwargs())
|
||||
|
||||
|
||||
def _create_index(
|
||||
config: IndexConfig,
|
||||
) -> None:
|
||||
client = get_client()
|
||||
pipeline_name = config.llama_cloud_pipeline_config.pipeline
|
||||
|
||||
pipelines = client.pipelines.search_pipelines(
|
||||
pipeline_name=pipeline_name,
|
||||
pipeline_type=PipelineType.MANAGED.value,
|
||||
)
|
||||
if len(pipelines) == 0:
|
||||
from llama_index.embeddings.openai import OpenAIEmbedding
|
||||
|
||||
if not isinstance(Settings.embed_model, OpenAIEmbedding):
|
||||
raise ValueError(
|
||||
"Creating a new pipeline with a non-OpenAI embedding model is not supported."
|
||||
)
|
||||
client.pipelines.upsert_pipeline(
|
||||
request={
|
||||
"name": pipeline_name,
|
||||
"embedding_config": {
|
||||
"type": "OPENAI_EMBEDDING",
|
||||
"component": {
|
||||
"api_key": os.getenv("OPENAI_API_KEY"), # editable
|
||||
"model_name": os.getenv("EMBEDDING_MODEL"),
|
||||
},
|
||||
},
|
||||
"transform_config": {
|
||||
"mode": "auto",
|
||||
"config": {
|
||||
"chunk_size": Settings.chunk_size, # editable
|
||||
"chunk_overlap": Settings.chunk_overlap, # editable
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from llama_index.core.prompts import PromptTemplate
|
||||
from llama_index.core.settings import Settings
|
||||
from llama_index.server.api.models import ChatAPIMessage
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class SuggestNextQuestionsService:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the conversation history.
|
||||
"""
|
||||
|
||||
prompt = PromptTemplate(
|
||||
r"""
|
||||
You're a helpful assistant! Your task is to suggest the next questions that user might interested in to keep the conversation going.
|
||||
Here is the conversation history
|
||||
---------------------
|
||||
{conversation}
|
||||
---------------------
|
||||
Given the conversation history, please give me 3 questions that user might ask next!
|
||||
Your answer should be wrapped in three sticks without any index numbers and follows the following format:
|
||||
\`\`\`
|
||||
<question 1>
|
||||
<question 2>
|
||||
<question 3>
|
||||
\`\`\`
|
||||
"""
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_configured_prompt(cls) -> PromptTemplate:
|
||||
prompt = os.getenv("NEXT_QUESTION_PROMPT", None)
|
||||
if not prompt:
|
||||
return cls.prompt
|
||||
return PromptTemplate(prompt)
|
||||
|
||||
@classmethod
|
||||
async def suggest_next_questions_all_messages(
|
||||
cls,
|
||||
messages: List[ChatAPIMessage],
|
||||
) -> Optional[List[str]]:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the conversation history.
|
||||
"""
|
||||
prompt_template = cls.get_configured_prompt()
|
||||
|
||||
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)
|
||||
return cls._extract_questions(output.text)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error when generating next question: {e}")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _extract_questions(cls, text: str) -> Union[List[str], None]:
|
||||
content_match = re.search(r"```(.*?)```", text, re.DOTALL)
|
||||
content = content_match.group(1) if content_match else None
|
||||
if not content:
|
||||
return None
|
||||
return [q.strip() for q in content.split("\n") if q.strip()]
|
||||
|
||||
@classmethod
|
||||
async def run(
|
||||
cls,
|
||||
chat_history: List[ChatAPIMessage],
|
||||
response: str,
|
||||
) -> Optional[List[str]]:
|
||||
"""
|
||||
Suggest the next questions that user might ask based on the chat history and the last response.
|
||||
"""
|
||||
messages = [
|
||||
*chat_history,
|
||||
ChatAPIMessage(role="assistant", content=response), # type: ignore
|
||||
]
|
||||
return await cls.suggest_next_questions_all_messages(messages)
|
||||
@@ -0,0 +1,47 @@
|
||||
from pydantic import Field, validator
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class ServerSettings(BaseSettings):
|
||||
url: str = Field(
|
||||
default="",
|
||||
description="The deployment URL of the server, to be referenced by tools and file services",
|
||||
)
|
||||
api_prefix: str = Field(
|
||||
default="/api",
|
||||
description="The prefix for the API endpoints",
|
||||
)
|
||||
|
||||
@property
|
||||
def file_server_url_prefix(self) -> str:
|
||||
return f"{self.url}{self.api_prefix}/files"
|
||||
|
||||
@property
|
||||
def api_url(self) -> str:
|
||||
return f"{self.url}{self.api_prefix}"
|
||||
|
||||
@validator("url")
|
||||
def validate_url(cls, v: str) -> str:
|
||||
if v.endswith("/"):
|
||||
raise ValueError("URL must not end with a '/'")
|
||||
return v
|
||||
|
||||
@validator("api_prefix")
|
||||
def validate_api_prefix(cls, v: str) -> str:
|
||||
if not v.startswith("/"):
|
||||
raise ValueError("API prefix must start with a '/'")
|
||||
return v
|
||||
|
||||
def set_url(self, v: str) -> None:
|
||||
self.url = v
|
||||
self.validate_url(v) # type: ignore
|
||||
|
||||
def set_api_prefix(self, v: str) -> None:
|
||||
self.api_prefix = v
|
||||
self.validate_api_prefix(v) # type: ignore
|
||||
|
||||
class Config:
|
||||
env_file_encoding = "utf-8"
|
||||
|
||||
|
||||
server_settings = ServerSettings()
|
||||
@@ -0,0 +1,242 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from enum import Enum
|
||||
from io import BytesIO
|
||||
|
||||
from llama_index.core.tools.function_tool import FunctionTool
|
||||
|
||||
OUTPUT_DIR = "output/tools"
|
||||
|
||||
|
||||
class DocumentType(Enum):
|
||||
PDF = "pdf"
|
||||
HTML = "html"
|
||||
|
||||
|
||||
COMMON_STYLES = """
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
line-height: 1.3;
|
||||
color: #333;
|
||||
}
|
||||
h1, h2, h3, h4, h5, h6 {
|
||||
margin-top: 1em;
|
||||
margin-bottom: 0.5em;
|
||||
}
|
||||
p {
|
||||
margin-bottom: 0.7em;
|
||||
}
|
||||
code {
|
||||
background-color: #f4f4f4;
|
||||
padding: 2px 4px;
|
||||
border-radius: 4px;
|
||||
}
|
||||
pre {
|
||||
background-color: #f4f4f4;
|
||||
padding: 10px;
|
||||
border-radius: 4px;
|
||||
overflow-x: auto;
|
||||
}
|
||||
table {
|
||||
border-collapse: collapse;
|
||||
width: 100%;
|
||||
margin-bottom: 1em;
|
||||
}
|
||||
th, td {
|
||||
border: 1px solid #ddd;
|
||||
padding: 8px;
|
||||
text-align: left;
|
||||
}
|
||||
th {
|
||||
background-color: #f2f2f2;
|
||||
font-weight: bold;
|
||||
}
|
||||
"""
|
||||
|
||||
HTML_SPECIFIC_STYLES = """
|
||||
body {
|
||||
max-width: 800px;
|
||||
margin: 0 auto;
|
||||
padding: 20px;
|
||||
}
|
||||
"""
|
||||
|
||||
PDF_SPECIFIC_STYLES = """
|
||||
@page {
|
||||
size: letter;
|
||||
margin: 2cm;
|
||||
}
|
||||
body {
|
||||
font-size: 11pt;
|
||||
}
|
||||
h1 { font-size: 18pt; }
|
||||
h2 { font-size: 16pt; }
|
||||
h3 { font-size: 14pt; }
|
||||
h4, h5, h6 { font-size: 12pt; }
|
||||
pre, code {
|
||||
font-family: Courier, monospace;
|
||||
font-size: 0.9em;
|
||||
}
|
||||
"""
|
||||
|
||||
HTML_TEMPLATE = """
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<style>
|
||||
{common_styles}
|
||||
{specific_styles}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
{content}
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
|
||||
class DocumentGenerator:
|
||||
def __init__(self, file_server_url_prefix: str):
|
||||
if not file_server_url_prefix:
|
||||
raise ValueError("file_server_url_prefix is required")
|
||||
self.file_server_url_prefix = file_server_url_prefix
|
||||
|
||||
@classmethod
|
||||
def _generate_html_content(cls, original_content: str) -> str:
|
||||
"""
|
||||
Generate HTML content from the original markdown content.
|
||||
"""
|
||||
try:
|
||||
import markdown # type: ignore
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Failed to import required modules. Please install markdown."
|
||||
)
|
||||
|
||||
# Convert markdown to HTML with fenced code and table extensions
|
||||
return markdown.markdown(original_content, extensions=["fenced_code", "tables"])
|
||||
|
||||
@classmethod
|
||||
def _generate_pdf(cls, html_content: str) -> BytesIO:
|
||||
"""
|
||||
Generate a PDF from the HTML content.
|
||||
"""
|
||||
try:
|
||||
from xhtml2pdf import pisa
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Failed to import required modules. Please install xhtml2pdf."
|
||||
)
|
||||
|
||||
pdf_html = HTML_TEMPLATE.format(
|
||||
common_styles=COMMON_STYLES,
|
||||
specific_styles=PDF_SPECIFIC_STYLES,
|
||||
content=html_content,
|
||||
)
|
||||
|
||||
buffer = BytesIO()
|
||||
pdf = pisa.pisaDocument(
|
||||
BytesIO(pdf_html.encode("UTF-8")), buffer, encoding="UTF-8"
|
||||
)
|
||||
|
||||
if pdf.err:
|
||||
logging.error(f"PDF generation failed: {pdf.err}")
|
||||
raise ValueError("PDF generation failed")
|
||||
|
||||
buffer.seek(0)
|
||||
return buffer
|
||||
|
||||
@classmethod
|
||||
def _generate_html(cls, html_content: str) -> str:
|
||||
"""
|
||||
Generate a complete HTML document with the given HTML content.
|
||||
"""
|
||||
return HTML_TEMPLATE.format(
|
||||
common_styles=COMMON_STYLES,
|
||||
specific_styles=HTML_SPECIFIC_STYLES,
|
||||
content=html_content,
|
||||
)
|
||||
|
||||
def generate_document(
|
||||
self, original_content: str, document_type: str, file_name: str
|
||||
) -> str:
|
||||
"""
|
||||
To generate document as PDF or HTML file.
|
||||
Parameters:
|
||||
original_content: str (markdown style)
|
||||
document_type: str (pdf or html) specify the type of the file format based on the use case
|
||||
file_name: str (name of the document file) must be a valid file name, no extensions needed
|
||||
Returns:
|
||||
str (URL to the document file): A file URL ready to serve.
|
||||
"""
|
||||
try:
|
||||
doc_type = DocumentType(document_type.lower())
|
||||
except ValueError:
|
||||
raise ValueError(
|
||||
f"Invalid document type: {document_type}. Must be 'pdf' or 'html'."
|
||||
)
|
||||
# Always generate html content first
|
||||
html_content = self._generate_html_content(original_content)
|
||||
|
||||
# Based on the type of document, generate the corresponding file
|
||||
if doc_type == DocumentType.PDF:
|
||||
content = self._generate_pdf(html_content)
|
||||
file_extension = "pdf"
|
||||
elif doc_type == DocumentType.HTML:
|
||||
content = BytesIO(self._generate_html(html_content).encode("utf-8"))
|
||||
file_extension = "html"
|
||||
else:
|
||||
raise ValueError(f"Unexpected document type: {document_type}")
|
||||
|
||||
file_name = self._validate_file_name(file_name)
|
||||
file_path = os.path.join(OUTPUT_DIR, f"{file_name}.{file_extension}")
|
||||
|
||||
self._write_to_file(content, file_path)
|
||||
|
||||
return (
|
||||
f"{self.file_server_url_prefix}/{OUTPUT_DIR}/{file_name}.{file_extension}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _write_to_file(content: BytesIO, file_path: str) -> None:
|
||||
"""
|
||||
Write the content to a file.
|
||||
"""
|
||||
try:
|
||||
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
||||
with open(file_path, "wb") as file:
|
||||
file.write(content.getvalue())
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _validate_file_name(file_name: str) -> str:
|
||||
"""
|
||||
Validate the file name.
|
||||
"""
|
||||
# Don't allow directory traversal
|
||||
if os.path.isabs(file_name):
|
||||
raise ValueError("File name is not allowed.")
|
||||
# Don't allow special characters
|
||||
if re.match(r"^[a-zA-Z0-9_.-]+$", file_name):
|
||||
return file_name
|
||||
else:
|
||||
raise ValueError("File name is not allowed to contain special characters.")
|
||||
|
||||
@classmethod
|
||||
def _validate_packages(cls) -> None:
|
||||
try:
|
||||
import markdown # noqa: F401
|
||||
import xhtml2pdf # noqa: F401
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Failed to import required modules. Please install markdown and xhtml2pdf "
|
||||
"using `pip install markdown xhtml2pdf`"
|
||||
)
|
||||
|
||||
def to_tool(self) -> FunctionTool:
|
||||
self._validate_packages()
|
||||
return FunctionTool.from_defaults(self.generate_document)
|
||||
@@ -0,0 +1,3 @@
|
||||
from .query import get_query_engine_tool
|
||||
|
||||
__all__ = ["get_query_engine_tool"]
|
||||
@@ -0,0 +1,49 @@
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
from llama_index.core.base.base_query_engine import BaseQueryEngine
|
||||
from llama_index.core.tools.query_engine import QueryEngineTool
|
||||
from llama_index.core.indices.base import BaseIndex
|
||||
|
||||
|
||||
def create_query_engine(index: BaseIndex, **kwargs: Any) -> BaseQueryEngine:
|
||||
"""
|
||||
Create a query engine for the given index.
|
||||
|
||||
Args:
|
||||
index: The index to create a query engine for.
|
||||
params (optional): Additional parameters for the query engine, e.g: similarity_top_k
|
||||
"""
|
||||
top_k = int(os.getenv("TOP_K", 0))
|
||||
if top_k != 0 and kwargs.get("filters") is None:
|
||||
kwargs["similarity_top_k"] = top_k
|
||||
|
||||
return index.as_query_engine(**kwargs)
|
||||
|
||||
|
||||
def get_query_engine_tool(
|
||||
index: BaseIndex,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> QueryEngineTool:
|
||||
"""
|
||||
Get a query engine tool for the given index.
|
||||
|
||||
Args:
|
||||
index: The index to create a query engine for.
|
||||
name (optional): The name of the tool.
|
||||
description (optional): The description of the tool.
|
||||
"""
|
||||
if name is None:
|
||||
name = "query_index"
|
||||
if description is None:
|
||||
description = (
|
||||
"Use this tool to retrieve information about the text corpus from an index."
|
||||
)
|
||||
query_engine = create_query_engine(index, **kwargs)
|
||||
return QueryEngineTool.from_defaults(
|
||||
query_engine=query_engine,
|
||||
name=name,
|
||||
description=description,
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
from datetime import timedelta
|
||||
|
||||
from cachetools import TTLCache, cached # type: ignore
|
||||
|
||||
from llama_index.core.storage import StorageContext
|
||||
|
||||
|
||||
@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)
|
||||
@@ -0,0 +1,216 @@
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from llama_index.core.tools import FunctionTool
|
||||
from llama_index.server.services.file import DocumentFile, FileService
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class InterpreterExtraResult(BaseModel):
|
||||
type: str
|
||||
content: Optional[str] = None
|
||||
filename: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
|
||||
|
||||
class E2BToolOutput(BaseModel):
|
||||
is_error: bool
|
||||
logs: "Logs" # type: ignore # noqa: F821
|
||||
error_message: Optional[str] = None
|
||||
results: List[InterpreterExtraResult] = []
|
||||
retry_count: int = 0
|
||||
|
||||
|
||||
class E2BCodeInterpreter:
|
||||
output_dir = "output/tools"
|
||||
uploaded_files_dir = "output/uploaded"
|
||||
interpreter: Optional["Sandbox"] = None # type: ignore # noqa: F821
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
output_dir: Optional[str] = None,
|
||||
uploaded_files_dir: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
api_key: The API key for the E2B Code Interpreter.
|
||||
output_dir: The directory for the output files. Default is `output/tools`.
|
||||
uploaded_files_dir: The directory for the files to be uploaded to the sandbox. Default is `output/uploaded`.
|
||||
"""
|
||||
self._validate_package()
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"api_key is required to run code interpreter. Get it here: https://e2b.dev/docs/getting-started/api-key"
|
||||
)
|
||||
self.api_key = api_key
|
||||
self.output_dir = output_dir or "output/tools"
|
||||
self.uploaded_files_dir = uploaded_files_dir or "output/uploaded"
|
||||
|
||||
@classmethod
|
||||
def _validate_package(cls) -> None:
|
||||
try:
|
||||
from e2b_code_interpreter import Sandbox # noqa: F401
|
||||
from e2b_code_interpreter.models import Logs # noqa: F401
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"e2b_code_interpreter is not installed. Please install it using `pip install e2b-code-interpreter`."
|
||||
)
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""
|
||||
Kill the interpreter when the tool is no longer in use.
|
||||
"""
|
||||
if self.interpreter is not None:
|
||||
self.interpreter.kill()
|
||||
|
||||
def _init_interpreter(self, sandbox_files: List[str] = []) -> None:
|
||||
"""
|
||||
Lazily initialize the interpreter.
|
||||
"""
|
||||
from e2b_code_interpreter import Sandbox
|
||||
|
||||
logger.info(f"Initializing interpreter with {len(sandbox_files)} files")
|
||||
self.interpreter = Sandbox(api_key=self.api_key)
|
||||
if len(sandbox_files) > 0:
|
||||
for file_path in sandbox_files:
|
||||
file_name = os.path.basename(file_path)
|
||||
local_file_path = os.path.join(self.uploaded_files_dir, file_name)
|
||||
with open(local_file_path, "rb") as f:
|
||||
content = f.read()
|
||||
if self.interpreter and self.interpreter.files:
|
||||
self.interpreter.files.write(file_path, content)
|
||||
logger.info(f"Uploaded {len(sandbox_files)} files to sandbox")
|
||||
|
||||
def _save_to_disk(self, base64_data: str, ext: str) -> DocumentFile:
|
||||
buffer = base64.b64decode(base64_data)
|
||||
|
||||
# Output from e2b doesn't have a name. Create a random name for it.
|
||||
filename = f"e2b_file_{uuid.uuid4()}.{ext}"
|
||||
|
||||
return FileService.save_file(
|
||||
buffer, file_name=filename, save_dir=self.output_dir
|
||||
)
|
||||
|
||||
def _parse_result(self, result: Any) -> List[InterpreterExtraResult]:
|
||||
"""
|
||||
The result could include multiple formats (e.g. png, svg, etc.) but encoded in base64
|
||||
We save each result to disk and return saved file metadata (extension, filename, url).
|
||||
"""
|
||||
if not result:
|
||||
return []
|
||||
|
||||
output = []
|
||||
|
||||
try:
|
||||
formats = result.formats()
|
||||
results = [result[format] for format in formats]
|
||||
|
||||
for ext, data in zip(formats, results):
|
||||
if ext in ["png", "svg", "jpeg", "pdf"]:
|
||||
document_file = self._save_to_disk(data, ext)
|
||||
output.append(
|
||||
InterpreterExtraResult(
|
||||
type=ext,
|
||||
filename=document_file.name,
|
||||
url=document_file.url,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Try serialize data to string
|
||||
try:
|
||||
data = str(data)
|
||||
except Exception as e:
|
||||
data = f"Error when serializing data: {e}"
|
||||
output.append(
|
||||
InterpreterExtraResult(
|
||||
type=ext,
|
||||
content=data,
|
||||
)
|
||||
)
|
||||
except Exception as error:
|
||||
logger.exception(error, exc_info=True)
|
||||
logger.error("Error when parsing output from E2b interpreter tool", error)
|
||||
|
||||
return output
|
||||
|
||||
def interpret(
|
||||
self,
|
||||
code: str,
|
||||
sandbox_files: List[str] = [],
|
||||
retry_count: int = 0,
|
||||
) -> E2BToolOutput:
|
||||
"""
|
||||
Execute Python code in a Jupyter notebook cell. The tool will return the result, stdout, stderr, display_data, and error.
|
||||
If the code needs to use a file, ALWAYS pass the file path in the sandbox_files argument.
|
||||
You have a maximum of 3 retries to get the code to run successfully.
|
||||
|
||||
Parameters:
|
||||
code (str): The Python code to be executed in a single cell.
|
||||
sandbox_files (List[str]): List of local file paths to be used by the code. The tool will throw an error if a file is not found.
|
||||
retry_count (int): Number of times the tool has been retried.
|
||||
"""
|
||||
from e2b_code_interpreter.models import Logs
|
||||
|
||||
if retry_count > 2:
|
||||
return E2BToolOutput(
|
||||
is_error=True,
|
||||
logs=Logs(
|
||||
stdout="",
|
||||
stderr="",
|
||||
display_data="",
|
||||
error="",
|
||||
),
|
||||
error_message="Failed to execute the code after 3 retries. Explain the error to the user and suggest a fix.",
|
||||
retry_count=retry_count,
|
||||
)
|
||||
|
||||
if self.interpreter is None:
|
||||
self._init_interpreter(sandbox_files)
|
||||
|
||||
if self.interpreter:
|
||||
logger.info(
|
||||
f"\n{'=' * 50}\n> Running following AI-generated code:\n{code}\n{'=' * 50}"
|
||||
)
|
||||
exec = self.interpreter.run_code(code)
|
||||
|
||||
if exec.error:
|
||||
error_message = f"The code failed to execute successfully. Error: {exec.error}. Try to fix the code and run again."
|
||||
logger.error(error_message)
|
||||
# Calling the generated code caused an error. Kill the interpreter and return the error to the LLM so it can try to fix the error
|
||||
try:
|
||||
self.interpreter.kill() # type: ignore
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self.interpreter = None
|
||||
output = E2BToolOutput(
|
||||
is_error=True,
|
||||
logs=exec.logs,
|
||||
results=[],
|
||||
error_message=error_message,
|
||||
retry_count=retry_count + 1,
|
||||
)
|
||||
else:
|
||||
if len(exec.results) == 0:
|
||||
output = E2BToolOutput(is_error=False, logs=exec.logs, results=[])
|
||||
else:
|
||||
results = self._parse_result(exec.results[0])
|
||||
output = E2BToolOutput(
|
||||
is_error=False,
|
||||
logs=exec.logs,
|
||||
results=results,
|
||||
retry_count=retry_count + 1,
|
||||
)
|
||||
return output
|
||||
else:
|
||||
raise ValueError("Interpreter is not initialized.")
|
||||
|
||||
def to_tool(self) -> FunctionTool:
|
||||
self._validate_package()
|
||||
return FunctionTool.from_defaults(self.interpret)
|
||||
@@ -0,0 +1,253 @@
|
||||
import logging
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, AsyncGenerator, Optional
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from llama_index.core.base.llms.types import ChatMessage, ChatResponse
|
||||
from llama_index.core.llms.function_calling import FunctionCallingLLM
|
||||
from llama_index.core.tools import (
|
||||
BaseTool,
|
||||
FunctionTool,
|
||||
ToolOutput,
|
||||
ToolSelection,
|
||||
)
|
||||
from llama_index.core.workflow import Context
|
||||
from llama_index.server.api.models import AgentRunEvent, AgentRunEventType
|
||||
from llama_index.core.agent.workflow.workflow_events import ToolCall, ToolCallResult
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class ToolCallOutput(BaseModel):
|
||||
tool_call_id: str
|
||||
tool_output: ToolOutput
|
||||
|
||||
|
||||
class ContextAwareTool(FunctionTool, ABC):
|
||||
@abstractmethod
|
||||
async def acall(self, ctx: Context, input: Any) -> ToolOutput: # type: ignore
|
||||
pass
|
||||
|
||||
|
||||
class ChatWithToolsResponse(BaseModel):
|
||||
"""
|
||||
A tool call response from chat_with_tools.
|
||||
"""
|
||||
|
||||
tool_calls: Optional[list[ToolSelection]]
|
||||
tool_call_message: Optional[ChatMessage]
|
||||
generator: Optional[AsyncGenerator[ChatResponse | None, None]]
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
def is_calling_different_tools(self) -> bool:
|
||||
tool_names = {tool_call.tool_name for tool_call in self.tool_calls or []}
|
||||
return len(tool_names) > 1
|
||||
|
||||
def has_tool_calls(self) -> bool:
|
||||
return self.tool_calls is not None and len(self.tool_calls) > 0
|
||||
|
||||
def tool_name(self) -> str:
|
||||
if not self.has_tool_calls():
|
||||
raise ValueError("No tool calls")
|
||||
if self.is_calling_different_tools():
|
||||
raise ValueError("Calling different tools")
|
||||
return self.tool_calls[0].tool_name # type: ignore
|
||||
|
||||
async def full_response(self) -> str:
|
||||
assert self.generator is not None
|
||||
full_response = ""
|
||||
async for chunk in self.generator:
|
||||
content = chunk.delta # type: ignore
|
||||
if content:
|
||||
full_response += content
|
||||
return full_response
|
||||
|
||||
|
||||
async def chat_with_tools( # type: ignore
|
||||
llm: FunctionCallingLLM,
|
||||
tools: list[BaseTool],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> ChatWithToolsResponse:
|
||||
"""
|
||||
Request LLM to call tools or not.
|
||||
This function doesn't change the memory.
|
||||
"""
|
||||
generator = _tool_call_generator(llm, tools, chat_history)
|
||||
is_tool_call = await generator.__anext__()
|
||||
if is_tool_call:
|
||||
# Last chunk is the full response
|
||||
# Wait for the last chunk
|
||||
full_response = None
|
||||
async for chunk in generator:
|
||||
full_response = chunk
|
||||
assert isinstance(full_response, ChatResponse)
|
||||
return ChatWithToolsResponse(
|
||||
tool_calls=llm.get_tool_calls_from_response(full_response),
|
||||
tool_call_message=full_response.message,
|
||||
generator=None,
|
||||
)
|
||||
else:
|
||||
return ChatWithToolsResponse(
|
||||
tool_calls=None,
|
||||
tool_call_message=None,
|
||||
generator=generator, # type: ignore
|
||||
)
|
||||
|
||||
|
||||
async def call_tools(
|
||||
ctx: Context,
|
||||
agent_name: str,
|
||||
tools: list[BaseTool],
|
||||
tool_calls: list[ToolSelection],
|
||||
emit_agent_events: bool = True,
|
||||
) -> list[ToolCallOutput]:
|
||||
"""
|
||||
Call tools and return the tool call responses.
|
||||
"""
|
||||
if len(tool_calls) == 0:
|
||||
return []
|
||||
tools_by_name = {tool.metadata.get_name(): tool for tool in tools}
|
||||
if len(tool_calls) == 1:
|
||||
if emit_agent_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(
|
||||
name=agent_name,
|
||||
msg=f"{tool_calls[0].tool_name}: {tool_calls[0].tool_kwargs}",
|
||||
)
|
||||
)
|
||||
return [
|
||||
await call_tool(ctx, tools_by_name[tool_calls[0].tool_name], tool_calls[0])
|
||||
]
|
||||
# Multiple tool calls, show progress
|
||||
tool_call_outputs: list[ToolCallOutput] = []
|
||||
|
||||
progress_id = str(uuid.uuid4())
|
||||
total_steps = len(tool_calls)
|
||||
if emit_agent_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(
|
||||
name=agent_name,
|
||||
msg=f"Making {total_steps} tool calls",
|
||||
)
|
||||
)
|
||||
for i, tool_call in enumerate(tool_calls):
|
||||
tool = tools_by_name.get(tool_call.tool_name)
|
||||
if not tool:
|
||||
tool_call_outputs.append(
|
||||
ToolCallOutput(
|
||||
tool_call_id=tool_call.tool_id,
|
||||
tool_output=ToolOutput(
|
||||
is_error=True,
|
||||
content=f"Tool {tool_call.tool_name} does not exist",
|
||||
tool_name=tool_call.tool_name,
|
||||
raw_input=tool_call.tool_kwargs,
|
||||
raw_output={
|
||||
"error": f"Tool {tool_call.tool_name} does not exist",
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
tool_call_output = await call_tool(
|
||||
ctx,
|
||||
tool,
|
||||
tool_call,
|
||||
)
|
||||
if emit_agent_events:
|
||||
ctx.write_event_to_stream(
|
||||
AgentRunEvent(
|
||||
name=agent_name,
|
||||
msg=f"{tool_call.tool_name}: {tool_call.tool_kwargs}",
|
||||
event_type=AgentRunEventType.PROGRESS,
|
||||
data={
|
||||
"id": progress_id,
|
||||
"total": total_steps,
|
||||
"current": i,
|
||||
},
|
||||
)
|
||||
)
|
||||
tool_call_outputs.append(tool_call_output)
|
||||
return tool_call_outputs
|
||||
|
||||
|
||||
async def call_tool(
|
||||
ctx: Context,
|
||||
tool: BaseTool,
|
||||
tool_call: ToolSelection,
|
||||
) -> ToolCallOutput:
|
||||
ctx.write_event_to_stream(
|
||||
ToolCall(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_id=tool_call.tool_id,
|
||||
tool_kwargs=tool_call.tool_kwargs,
|
||||
)
|
||||
)
|
||||
try:
|
||||
if isinstance(tool, ContextAwareTool):
|
||||
if ctx is None:
|
||||
raise ValueError("Context is required for context aware tool")
|
||||
# inject context for calling an context aware tool
|
||||
output = await tool.acall(ctx=ctx, **tool_call.tool_kwargs)
|
||||
else:
|
||||
output = await tool.acall(**tool_call.tool_kwargs) # type: ignore
|
||||
except Exception as e:
|
||||
logger.error(f"Got error in tool {tool_call.tool_name}: {e!s}")
|
||||
output = ToolOutput(
|
||||
is_error=True,
|
||||
content=f"Error: {e!s}",
|
||||
tool_name=tool.metadata.get_name(),
|
||||
raw_input=tool_call.tool_kwargs,
|
||||
raw_output={
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
ctx.write_event_to_stream(
|
||||
ToolCallResult(
|
||||
tool_name=tool_call.tool_name,
|
||||
tool_kwargs=tool_call.tool_kwargs,
|
||||
tool_id=tool_call.tool_id,
|
||||
tool_output=output,
|
||||
return_direct=False,
|
||||
)
|
||||
)
|
||||
return ToolCallOutput(
|
||||
tool_call_id=tool_call.tool_id,
|
||||
tool_output=output,
|
||||
)
|
||||
|
||||
|
||||
async def _tool_call_generator(
|
||||
llm: FunctionCallingLLM,
|
||||
tools: list[BaseTool],
|
||||
chat_history: list[ChatMessage],
|
||||
) -> AsyncGenerator[ChatResponse | bool, None]:
|
||||
response_stream = await llm.astream_chat_with_tools(
|
||||
tools,
|
||||
chat_history=chat_history,
|
||||
allow_parallel_tool_calls=False,
|
||||
)
|
||||
|
||||
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 # type: ignore
|
||||
elif not yielded_indicator:
|
||||
# Yield the indicator for a tool call
|
||||
yield True
|
||||
yielded_indicator = True
|
||||
|
||||
full_response = chunk
|
||||
|
||||
if full_response:
|
||||
yield full_response # type: ignore
|
||||
Generated
+6100
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,64 @@
|
||||
[build-system]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
requires = ["poetry-core"]
|
||||
|
||||
[tool.codespell]
|
||||
check-filenames = true
|
||||
check-hidden = true
|
||||
# Feel free to un-skip examples, and experimental, you will just need to
|
||||
# work through many typos (--write-changes and --interactive will help)
|
||||
skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
|
||||
|
||||
[tool.mypy]
|
||||
disallow_untyped_defs = true
|
||||
# Remove venv skip when integrated with pre-commit
|
||||
exclude = ["_static", "build", "examples", "notebooks", "venv"]
|
||||
ignore_missing_imports = true
|
||||
namespace_packages = true
|
||||
explicit_package_bases = true
|
||||
python_version = "3.10"
|
||||
|
||||
[tool.poetry]
|
||||
authors = ["Your Name <you@example.com>"]
|
||||
description = "llama-index fastapi server"
|
||||
exclude = ["**/BUILD"]
|
||||
license = "MIT"
|
||||
name = "llama-index-server"
|
||||
packages = [{include = "llama_index/"}]
|
||||
readme = "README.md"
|
||||
version = "0.1.7"
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.9,<4.0"
|
||||
fastapi = {extras = ["standard"], version = "^0.115.11"}
|
||||
cachetools = "^5.5.2"
|
||||
requests = "^2.32.3"
|
||||
pydantic-settings = "^2.8.1"
|
||||
llama-index-core = "0.12.25"
|
||||
llama-index-readers-file = "^0.4.6"
|
||||
llama-index-indices-managed-llama-cloud = "0.6.3"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
|
||||
codespell = {extras = ["toml"], version = ">=v2.2.6"}
|
||||
e2b-code-interpreter = "^1.1.1"
|
||||
ipython = "8.10.0"
|
||||
jupyter = "^1.0.0"
|
||||
markdown = "^3.7"
|
||||
mypy = "1.15.0"
|
||||
pre-commit = "3.2.0"
|
||||
pylint = "2.15.10"
|
||||
pytest = "^8.3.5"
|
||||
pytest-asyncio = "^0.25.3"
|
||||
pytest-mock = "3.11.1"
|
||||
ruff = "0.0.292"
|
||||
tree-sitter-languages = "^1.8.0"
|
||||
types-Deprecated = ">=0.1.0"
|
||||
types-PyYAML = "^6.0.12.12"
|
||||
types-protobuf = "^4.24.0.4"
|
||||
types-redis = "4.5.5.0"
|
||||
types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991
|
||||
types-setuptools = "67.1.0.0"
|
||||
xhtml2pdf = "^0.2.17"
|
||||
pytest-cov = "^6.0.0"
|
||||
llama-cloud = "^0.1.17"
|
||||
@@ -0,0 +1,149 @@
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from llama_index.core.workflow import StopEvent, Workflow
|
||||
from llama_index.core.workflow.handler import WorkflowHandler
|
||||
from llama_index.server.api.models import ChatAPIMessage, ChatRequest
|
||||
from llama_index.server.api.routers.chat import chat_router
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def logger():
|
||||
return logging.getLogger("test")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chat_request():
|
||||
"""Create a simple chat request with one user message."""
|
||||
return ChatRequest(
|
||||
messages=[ChatAPIMessage(role="user", content="Hello, how are you?")]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_workflow():
|
||||
"""Create a mock workflow that returns a simple response."""
|
||||
workflow = MagicMock(spec=Workflow)
|
||||
handler = AsyncMock(spec=WorkflowHandler)
|
||||
|
||||
# Setup the handler to stream a simple response event
|
||||
async def mock_stream_events():
|
||||
yield StopEvent(result="I'm doing well, thank you for asking!")
|
||||
|
||||
handler.stream_events.return_value = mock_stream_events()
|
||||
workflow.run.return_value = handler
|
||||
|
||||
return workflow
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def workflow_factory(mock_workflow):
|
||||
"""Create a factory function that returns our mock workflow."""
|
||||
|
||||
def factory(verbose=False):
|
||||
return mock_workflow
|
||||
|
||||
return factory
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_chat_router(chat_request, workflow_factory, logger):
|
||||
"""Test that the chat router handles a request correctly."""
|
||||
# Create a FastAPI app and mount our router
|
||||
app = FastAPI()
|
||||
router = chat_router(workflow_factory, logger)
|
||||
app.include_router(router)
|
||||
|
||||
# Make a request to the chat endpoint
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=app), base_url="http://test"
|
||||
) as client:
|
||||
response = await client.post("/chat", json=chat_request.model_dump())
|
||||
|
||||
# Check response status
|
||||
assert response.status_code == 200
|
||||
|
||||
# For streaming responses we don't check the content-type header directly
|
||||
# Instead, check that we get the expected content in the response body
|
||||
|
||||
# The response is a stream, so we need to collect the chunks
|
||||
content = response.content.decode()
|
||||
|
||||
# Verify content structure follows expected format
|
||||
assert "0:" in content # Text prefix for VercelStreamResponse
|
||||
# Verify if the response contains the expected message
|
||||
assert "I'm doing well" in content
|
||||
|
||||
# Verify the mock workflow was called correctly
|
||||
mock_workflow = workflow_factory()
|
||||
mock_workflow.run.assert_called_once()
|
||||
|
||||
# Verify the workflow was called with the correct arguments
|
||||
call_args = mock_workflow.run.call_args[1]
|
||||
assert call_args["user_msg"] == "Hello, how are you?"
|
||||
assert isinstance(call_args["chat_history"], list)
|
||||
assert len(call_args["chat_history"]) == 0 # No history for first message
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_chat_with_agent_workflow(logger):
|
||||
"""Test that the chat router works with a workflow that mimics an agent workflow."""
|
||||
# Create a simple workflow that mimics an agent workflow
|
||||
mock_workflow = MagicMock(spec=Workflow)
|
||||
handler = AsyncMock(spec=WorkflowHandler)
|
||||
|
||||
# Setup the handler to stream a simple response about weather
|
||||
async def mock_stream_events():
|
||||
yield StopEvent(
|
||||
result="The weather in New York is sunny. I used the weather tool to get this information."
|
||||
)
|
||||
|
||||
handler.stream_events.return_value = mock_stream_events()
|
||||
mock_workflow.run.return_value = handler
|
||||
|
||||
# Create a factory function that returns our mock workflow
|
||||
def workflow_factory(verbose=False):
|
||||
return mock_workflow
|
||||
|
||||
# Create a FastAPI app and mount our router
|
||||
app = FastAPI()
|
||||
router = chat_router(workflow_factory, logger)
|
||||
app.include_router(router)
|
||||
|
||||
# Create a chat request asking about weather
|
||||
chat_request = ChatRequest(
|
||||
messages=[
|
||||
ChatAPIMessage(role="user", content="What's the weather in New York?")
|
||||
]
|
||||
)
|
||||
|
||||
# Make a request to the chat endpoint
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=app), base_url="http://test"
|
||||
) as client:
|
||||
response = await client.post("/chat", json=chat_request.model_dump())
|
||||
|
||||
# Check response status
|
||||
assert response.status_code == 200
|
||||
|
||||
# The response is a stream, so we need to collect the chunks
|
||||
content = response.content.decode()
|
||||
|
||||
# Verify content structure follows expected format
|
||||
assert "0:" in content # Text prefix for VercelStreamResponse
|
||||
|
||||
# Verify the response content contains expected keywords
|
||||
assert "weather" in content and "New York" in content and "sunny" in content
|
||||
|
||||
# Verify the mock workflow was called correctly
|
||||
mock_workflow.run.assert_called_once()
|
||||
|
||||
# Verify the workflow was called with the correct arguments
|
||||
call_args = mock_workflow.run.call_args[1]
|
||||
assert call_args["user_msg"] == "What's the weather in New York?"
|
||||
assert isinstance(call_args["chat_history"], list)
|
||||
assert len(call_args["chat_history"]) == 0 # No history for first message
|
||||
@@ -0,0 +1,249 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from llama_index.core.agent.workflow.workflow_events import AgentStream
|
||||
from llama_index.core.workflow import StopEvent
|
||||
from llama_index.core.workflow.handler import WorkflowHandler
|
||||
from llama_index.server.api.models import ChatAPIMessage, ChatRequest
|
||||
from llama_index.server.api.routers.chat import _stream_content
|
||||
from llama_index.server.api.utils.vercel_stream import VercelStreamResponse
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def logger():
|
||||
return logging.getLogger("test")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chat_request():
|
||||
return ChatRequest(messages=[ChatAPIMessage(role="user", content="test message")])
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mock_workflow_handler():
|
||||
handler = AsyncMock(spec=WorkflowHandler)
|
||||
handler.accumulate_text = MagicMock()
|
||||
return handler
|
||||
|
||||
|
||||
class TestEventStream:
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_agent_stream(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.return_value = (
|
||||
self._mock_agent_stream_events()
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 3 # Empty start + 2 text chunks
|
||||
assert result[0] == VercelStreamResponse.convert_text("")
|
||||
assert result[1] == VercelStreamResponse.convert_text("Hello")
|
||||
assert result[2] == VercelStreamResponse.convert_text(" World")
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_stop_event_string(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.return_value = (
|
||||
self._mock_stop_event_string()
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 2 # Empty start + result string
|
||||
assert result[0] == VercelStreamResponse.convert_text("")
|
||||
assert result[1] == VercelStreamResponse.convert_text("Final answer")
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_stop_event_delta_objects(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.return_value = (
|
||||
self._mock_stop_event_delta_objects()
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 3 # Empty start + 2 delta chunks
|
||||
assert result[0] == VercelStreamResponse.convert_text("")
|
||||
assert result[1] == VercelStreamResponse.convert_text("Delta 1")
|
||||
assert result[2] == VercelStreamResponse.convert_text("Delta 2")
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_event_with_to_response(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.return_value = (
|
||||
self._mock_event_with_to_response()
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 2 # Empty start + event with to_response
|
||||
assert result[0] == VercelStreamResponse.convert_text("")
|
||||
assert result[1] == VercelStreamResponse.convert_data({"event_type": "test"})
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_event_with_model_dump(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.return_value = (
|
||||
self._mock_event_with_model_dump()
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 2 # Empty start + event with model_dump
|
||||
assert result[0] == VercelStreamResponse.convert_text("")
|
||||
assert result[1] == VercelStreamResponse.convert_data(None)
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_cancelled_error(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
mock_workflow_handler.stream_events.side_effect = asyncio.CancelledError()
|
||||
logger.warning = MagicMock()
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 0
|
||||
mock_workflow_handler.cancel_run.assert_called_once()
|
||||
logger.warning.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_stream_content_with_exception(
|
||||
self, mock_workflow_handler, chat_request, logger
|
||||
):
|
||||
# Setup
|
||||
error_message = "Test error"
|
||||
mock_workflow_handler.stream_events.side_effect = Exception(error_message)
|
||||
logger.error = MagicMock()
|
||||
|
||||
# Execute
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in _stream_content(
|
||||
mock_workflow_handler, chat_request, logger
|
||||
)
|
||||
]
|
||||
|
||||
# Assert
|
||||
assert len(result) == 1
|
||||
assert result[0] == VercelStreamResponse.convert_error(error_message)
|
||||
mock_workflow_handler.cancel_run.assert_called_once()
|
||||
logger.error.assert_called_once()
|
||||
|
||||
async def _mock_agent_stream_events(self):
|
||||
yield AgentStream(
|
||||
delta="Hello", response="", current_agent_name="", tool_calls=[], raw=""
|
||||
)
|
||||
yield AgentStream(
|
||||
delta=" World", response="", current_agent_name="", tool_calls=[], raw=""
|
||||
)
|
||||
|
||||
async def _mock_agent_stream_with_empty_deltas(self):
|
||||
yield AgentStream(
|
||||
delta=" ", # Empty delta with spaces - should be filtered
|
||||
response="",
|
||||
current_agent_name="",
|
||||
tool_calls=[],
|
||||
raw="",
|
||||
)
|
||||
yield AgentStream(
|
||||
delta="Valid delta",
|
||||
response="",
|
||||
current_agent_name="",
|
||||
tool_calls=[],
|
||||
raw="",
|
||||
)
|
||||
yield AgentStream(
|
||||
delta="\n", # Newline-only delta - should be filtered
|
||||
response="",
|
||||
current_agent_name="",
|
||||
tool_calls=[],
|
||||
raw="",
|
||||
)
|
||||
|
||||
async def _mock_stop_event_string(self):
|
||||
yield StopEvent(result="Final answer")
|
||||
|
||||
async def _mock_stop_event_delta_objects(self):
|
||||
async def generator():
|
||||
# Create proper objects with delta attribute that can be serialized
|
||||
class ObjectWithDelta:
|
||||
def __init__(self, delta_value) -> None:
|
||||
self.delta = delta_value
|
||||
|
||||
yield ObjectWithDelta("Delta 1")
|
||||
yield ObjectWithDelta("Delta 2")
|
||||
|
||||
yield StopEvent(result=generator())
|
||||
|
||||
async def _mock_dict_event(self):
|
||||
yield {"key": "value"}
|
||||
|
||||
async def _mock_event_with_to_response(self):
|
||||
event = MagicMock()
|
||||
event.to_response.return_value = {"event_type": "test"}
|
||||
yield event
|
||||
|
||||
async def _mock_event_with_model_dump(self):
|
||||
event = MagicMock()
|
||||
event.model_dump.return_value = {"name": "test_event"}
|
||||
# Override to_response to return None - this means convert_data(None) will be called
|
||||
event.to_response = MagicMock(return_value=None)
|
||||
# The model_dump value is ignored when to_response returns None
|
||||
yield event
|
||||
@@ -0,0 +1,205 @@
|
||||
import os
|
||||
import uuid
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
import pytest
|
||||
from llama_index.server.services.file import FileService, _sanitize_file_name
|
||||
|
||||
|
||||
class TestFileService:
|
||||
def test_sanitize_file_name(self):
|
||||
# Test with normal alphanumeric name
|
||||
assert _sanitize_file_name("test123") == "test123"
|
||||
|
||||
# Test with spaces
|
||||
assert _sanitize_file_name("test file") == "test_file"
|
||||
|
||||
# Test with special characters
|
||||
assert _sanitize_file_name("test@file!name") == "test_file_name"
|
||||
|
||||
# Test with path-like characters
|
||||
assert _sanitize_file_name("test/file/name") == "test_file_name"
|
||||
|
||||
# Test with dots (should be preserved)
|
||||
assert _sanitize_file_name("test.file.name") == "test.file.name"
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_string_content(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_getsize.return_value = 11 # Length of "Hello World"
|
||||
|
||||
# Execute
|
||||
result = FileService.save_file(
|
||||
content="Hello World", file_name="test.txt", save_dir="test_dir"
|
||||
)
|
||||
|
||||
# Assert
|
||||
expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
|
||||
mock_makedirs.assert_called_once_with(
|
||||
os.path.dirname(expected_path), exist_ok=True
|
||||
)
|
||||
mock_file_open.assert_called_once_with(expected_path, "wb")
|
||||
mock_file_open().write.assert_called_once_with(b"Hello World")
|
||||
|
||||
assert result.id == test_uuid
|
||||
assert result.name == f"test_{test_uuid}.txt"
|
||||
assert result.type == "txt"
|
||||
assert result.size == 11
|
||||
assert result.path == expected_path
|
||||
assert result.url.endswith(expected_path.replace(os.path.sep, "/"))
|
||||
assert result.refs is None
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_bytes_content(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_getsize.return_value = 11 # Length of "Hello World"
|
||||
|
||||
# Execute
|
||||
result = FileService.save_file(
|
||||
content=b"Hello World", file_name="test.txt", save_dir="test_dir"
|
||||
)
|
||||
|
||||
# Assert
|
||||
expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
|
||||
mock_makedirs.assert_called_once_with(
|
||||
os.path.dirname(expected_path), exist_ok=True
|
||||
)
|
||||
mock_file_open.assert_called_once_with(expected_path, "wb")
|
||||
mock_file_open().write.assert_called_once_with(b"Hello World")
|
||||
assert result.path == expected_path
|
||||
assert result.type == "txt"
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_with_special_characters(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_getsize.return_value = 11
|
||||
|
||||
# Execute
|
||||
result = FileService.save_file(
|
||||
content="Hello World", file_name="test@file!.txt", save_dir="test_dir"
|
||||
)
|
||||
|
||||
# Assert
|
||||
expected_path = os.path.join("test_dir", f"test_file__{test_uuid}.txt")
|
||||
mock_makedirs.assert_called_once_with(
|
||||
os.path.dirname(expected_path), exist_ok=True
|
||||
)
|
||||
mock_file_open.assert_called_once_with(expected_path, "wb")
|
||||
assert result.path == expected_path
|
||||
assert result.name == f"test_file__{test_uuid}.txt"
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_default_directory(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_getsize.return_value = 11
|
||||
|
||||
# Execute
|
||||
result = FileService.save_file(content="Hello World", file_name="test.txt")
|
||||
|
||||
# Assert
|
||||
expected_path = os.path.join("output", "uploaded", f"test_{test_uuid}.txt")
|
||||
mock_makedirs.assert_called_once_with(
|
||||
os.path.dirname(expected_path), exist_ok=True
|
||||
)
|
||||
assert result.path == expected_path
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.getenv")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open", new_callable=mock_open)
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_custom_url_prefix(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_getenv, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_getsize.return_value = 11
|
||||
mock_getenv.return_value = "/api/files"
|
||||
|
||||
# Execute
|
||||
result = FileService.save_file(
|
||||
content="Hello World", file_name="test.txt", save_dir="test_dir"
|
||||
)
|
||||
|
||||
# Assert
|
||||
expected_path = os.path.join("test_dir", f"test_{test_uuid}.txt")
|
||||
mock_makedirs.assert_called_once_with(
|
||||
os.path.dirname(expected_path), exist_ok=True
|
||||
)
|
||||
mock_file_open.assert_called_once_with(expected_path, "wb")
|
||||
assert result.path == expected_path
|
||||
# URL paths must use forward slashes, even on Windows
|
||||
expected_url = f"/api/files/test_dir/test_{test_uuid}.txt"
|
||||
assert result.url == expected_url
|
||||
|
||||
def test_save_file_no_extension(self):
|
||||
# Test that saving a file without extension raises ValueError
|
||||
with pytest.raises(ValueError, match="File is not supported!"):
|
||||
FileService.save_file(
|
||||
content="Hello World", file_name="test", save_dir="test_dir"
|
||||
)
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open")
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_permission_error(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_file_open.side_effect = PermissionError("Permission denied")
|
||||
|
||||
# Execute and Assert
|
||||
with pytest.raises(PermissionError):
|
||||
FileService.save_file(
|
||||
content="Hello World", file_name="test.txt", save_dir="test_dir"
|
||||
)
|
||||
|
||||
@patch("uuid.uuid4")
|
||||
@patch("os.path.getsize")
|
||||
@patch("builtins.open")
|
||||
@patch("os.makedirs")
|
||||
def test_save_file_io_error(
|
||||
self, mock_makedirs, mock_file_open, mock_getsize, mock_uuid
|
||||
):
|
||||
# Setup
|
||||
test_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
mock_uuid.return_value = uuid.UUID(test_uuid)
|
||||
mock_file_open.side_effect = OSError("IO Error")
|
||||
|
||||
# Execute and Assert
|
||||
with pytest.raises(IOError):
|
||||
FileService.save_file(
|
||||
content="Hello World", file_name="test.txt", save_dir="test_dir"
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from llama_index.core.agent.workflow import AgentWorkflow
|
||||
from llama_index.core.llms import MockLLM
|
||||
from llama_index.server import LlamaIndexServer
|
||||
|
||||
|
||||
def fetch_weather(city: str) -> str:
|
||||
"""Fetch the weather for a given city."""
|
||||
return f"The weather in {city} is sunny."
|
||||
|
||||
|
||||
def _agent_workflow() -> AgentWorkflow:
|
||||
# Use MockLLM instead of default OpenAI
|
||||
mock_llm = MockLLM()
|
||||
return AgentWorkflow.from_tools_or_functions(
|
||||
tools_or_functions=[fetch_weather],
|
||||
verbose=True,
|
||||
llm=mock_llm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def server() -> LlamaIndexServer:
|
||||
"""Fixture to create a LlamaIndexServer instance."""
|
||||
return LlamaIndexServer(
|
||||
workflow_factory=_agent_workflow,
|
||||
verbose=True,
|
||||
use_default_routers=True,
|
||||
mount_ui=False,
|
||||
env="dev",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_server_has_chat_route(server: LlamaIndexServer) -> None:
|
||||
"""Test that the server has the chat API route."""
|
||||
chat_route_exists = any(route.path == "/api/chat" for route in server.routes)
|
||||
assert chat_route_exists, "Chat API route not found in server routes"
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_server_swagger_docs(server: LlamaIndexServer) -> None:
|
||||
"""Test that the server serves Swagger UI docs."""
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=server), base_url="http://test"
|
||||
) as ac:
|
||||
response = await ac.get("/docs")
|
||||
assert response.status_code == 200
|
||||
assert "text/html" in response.headers["content-type"]
|
||||
assert "Swagger UI" in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_ui_is_downloaded(server: LlamaIndexServer) -> None:
|
||||
"""
|
||||
Test if the UI is downloaded and mounted correctly.
|
||||
"""
|
||||
import os
|
||||
import shutil
|
||||
|
||||
# Clean up any existing static directory first
|
||||
if os.path.exists(".ui"):
|
||||
shutil.rmtree(".ui")
|
||||
|
||||
# Create a new server with UI enabled
|
||||
ui_server = LlamaIndexServer(
|
||||
workflow_factory=_agent_workflow,
|
||||
verbose=True,
|
||||
use_default_routers=True,
|
||||
env="dev",
|
||||
include_ui=True,
|
||||
)
|
||||
|
||||
# Verify that static directory was created with index.html
|
||||
assert os.path.exists("./.ui"), "Static directory was not created"
|
||||
assert os.path.isdir("./.ui"), "Static path is not a directory"
|
||||
assert os.path.exists("./.ui/index.html"), "index.html was not downloaded"
|
||||
|
||||
# Check if the UI is mounted and accessible
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=ui_server), base_url="http://test"
|
||||
) as ac:
|
||||
response = await ac.get("/")
|
||||
assert response.status_code == 200
|
||||
assert "text/html" in response.headers["content-type"]
|
||||
|
||||
# Clean up after test
|
||||
shutil.rmtree("./.ui")
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_ui_is_accessible(server: LlamaIndexServer) -> None:
|
||||
"""
|
||||
Test if the UI is accessible.
|
||||
"""
|
||||
# Manually trigger UI mounting
|
||||
server.mount_ui()
|
||||
|
||||
async with AsyncClient(
|
||||
transport=ASGITransport(app=server), base_url="http://test"
|
||||
) as ac:
|
||||
response = await ac.get("/")
|
||||
assert response.status_code == 200
|
||||
assert "text/html" in response.headers["content-type"]
|
||||
@@ -0,0 +1,89 @@
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from llama_index.server.tools.document_generator import (
|
||||
OUTPUT_DIR,
|
||||
DocumentGenerator,
|
||||
)
|
||||
|
||||
|
||||
class TestDocumentGenerator:
|
||||
def test_validate_file_name(self) -> None:
|
||||
# Valid names
|
||||
assert (
|
||||
DocumentGenerator("/api/files")._validate_file_name("valid-name")
|
||||
== "valid-name"
|
||||
)
|
||||
|
||||
# Invalid names
|
||||
with pytest.raises(ValueError):
|
||||
DocumentGenerator("/api/files")._validate_file_name("/invalid/path")
|
||||
|
||||
@patch("os.makedirs")
|
||||
@patch("builtins.open")
|
||||
def test_write_to_file(self, mock_open, mock_makedirs): # type: ignore
|
||||
content = BytesIO(b"test")
|
||||
DocumentGenerator("/api/files")._write_to_file(content, "path/file.txt")
|
||||
|
||||
mock_makedirs.assert_called_once()
|
||||
mock_open.assert_called_once()
|
||||
mock_open.return_value.__enter__.return_value.write.assert_called_once_with(
|
||||
b"test"
|
||||
)
|
||||
|
||||
@patch("markdown.markdown")
|
||||
def test_html_generation(self, mock_markdown): # type: ignore
|
||||
mock_markdown.return_value = "<h1>Test</h1>"
|
||||
|
||||
# Test HTML content generation
|
||||
assert (
|
||||
DocumentGenerator("/api/files")._generate_html_content("# Test")
|
||||
== "<h1>Test</h1>"
|
||||
)
|
||||
|
||||
# Test full HTML generation
|
||||
html = DocumentGenerator("/api/files")._generate_html("<h1>Test</h1>")
|
||||
assert "<!DOCTYPE html>" in html
|
||||
assert "<h1>Test</h1>" in html
|
||||
|
||||
@patch("xhtml2pdf.pisa.pisaDocument")
|
||||
def test_pdf_generation(self, mock_pisa): # type: ignore
|
||||
# Success case
|
||||
mock_pisa.return_value = MagicMock(err=None)
|
||||
assert isinstance(
|
||||
DocumentGenerator("/api/files")._generate_pdf("test"), BytesIO
|
||||
)
|
||||
|
||||
# Error case
|
||||
mock_pisa.return_value = MagicMock(err="Error")
|
||||
with pytest.raises(ValueError):
|
||||
DocumentGenerator("/api/files")._generate_pdf("test")
|
||||
|
||||
@patch.multiple(
|
||||
DocumentGenerator,
|
||||
_generate_html_content=MagicMock(return_value="<h1>Test</h1>"),
|
||||
_generate_html=MagicMock(
|
||||
return_value="<html><body><h1>Test</h1></body></html>"
|
||||
),
|
||||
_generate_pdf=MagicMock(return_value=BytesIO(b"pdf")),
|
||||
_write_to_file=MagicMock(),
|
||||
)
|
||||
def test_generate_document(self): # type: ignore
|
||||
# HTML generation
|
||||
url = DocumentGenerator("/api/files").generate_document(
|
||||
"# Test", "html", "test-doc"
|
||||
)
|
||||
assert url == f"/api/files/{OUTPUT_DIR}/test-doc.html"
|
||||
|
||||
# PDF generation
|
||||
url = DocumentGenerator("/api/files").generate_document(
|
||||
"# Test", "pdf", "test-doc"
|
||||
)
|
||||
assert url == f"/api/files/{OUTPUT_DIR}/test-doc.pdf"
|
||||
|
||||
# Invalid type
|
||||
with pytest.raises(ValueError):
|
||||
DocumentGenerator("/api/files").generate_document(
|
||||
"# Test", "invalid", "test-doc"
|
||||
)
|
||||
@@ -0,0 +1,65 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from e2b_code_interpreter.models import Execution, Logs
|
||||
from llama_index.server.tools.interpreter import E2BCodeInterpreter
|
||||
|
||||
|
||||
class TestE2BCodeInterpreter:
|
||||
@pytest.fixture()
|
||||
def sandbox(self): # type: ignore
|
||||
"""Create a mock Sandbox with no API key requirement."""
|
||||
mock_sandbox = MagicMock()
|
||||
mock_sandbox.files = MagicMock()
|
||||
mock_sandbox.files.write = MagicMock()
|
||||
mock_sandbox.run_code = MagicMock()
|
||||
return mock_sandbox
|
||||
|
||||
@pytest.fixture()
|
||||
def code_interpreter(self, sandbox): # type: ignore
|
||||
"""Create E2BCodeInterpreter that uses the mock Sandbox."""
|
||||
interpreter = E2BCodeInterpreter(api_key="dummy_key")
|
||||
interpreter.interpreter = sandbox
|
||||
return interpreter
|
||||
|
||||
def test_interpret_success(self, code_interpreter, sandbox) -> None: # type: ignore
|
||||
"""Test successful code execution."""
|
||||
# Mock execution result
|
||||
mock_execution = Execution()
|
||||
mock_execution.error = None
|
||||
mock_execution.results = []
|
||||
mock_execution.logs = Logs(
|
||||
stdout="stdout", stderr="", display_data="", error=""
|
||||
)
|
||||
sandbox.run_code.return_value = mock_execution
|
||||
|
||||
# Run the code
|
||||
result = code_interpreter.interpret("print('hello')")
|
||||
|
||||
# Verify
|
||||
sandbox.run_code.assert_called_once_with("print('hello')")
|
||||
assert result.is_error is False
|
||||
assert result.logs == mock_execution.logs
|
||||
|
||||
def test_interpret_error(self, code_interpreter, sandbox) -> None: # type: ignore
|
||||
"""Test error in code execution."""
|
||||
# Mock execution result with error
|
||||
mock_execution = Execution()
|
||||
mock_execution.error = "Test error"
|
||||
mock_execution.logs = Logs(
|
||||
stdout="", stderr="error", display_data="", error="Test error"
|
||||
)
|
||||
sandbox.run_code.return_value = mock_execution
|
||||
|
||||
# Run the code
|
||||
result = code_interpreter.interpret("bad code")
|
||||
|
||||
# Verify
|
||||
assert result.is_error is True
|
||||
assert "Error: Test error" in result.error_message
|
||||
sandbox.kill.assert_called_once()
|
||||
|
||||
def test_to_tool(self, code_interpreter) -> None: # type: ignore
|
||||
"""Test tool conversion."""
|
||||
tool = code_interpreter.to_tool()
|
||||
assert tool.fn == code_interpreter.interpret
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "create-llama",
|
||||
"version": "0.3.27",
|
||||
"version": "0.5.0",
|
||||
"description": "Create LlamaIndex-powered apps with one command",
|
||||
"keywords": [
|
||||
"rag",
|
||||
|
||||
+2
-1
@@ -16,5 +16,6 @@ export const askQuestions = async (
|
||||
await askProQuestions(args);
|
||||
return args as unknown as QuestionResults;
|
||||
}
|
||||
return await askSimpleQuestions(args);
|
||||
const results = await askSimpleQuestions(args);
|
||||
return results;
|
||||
};
|
||||
|
||||
+22
-113
@@ -1,25 +1,12 @@
|
||||
import prompts from "prompts";
|
||||
import {
|
||||
AI_REPORTS,
|
||||
EXAMPLE_10K_SEC_FILES,
|
||||
EXAMPLE_FILE,
|
||||
EXAMPLE_GDPR,
|
||||
} from "../helpers/datasources";
|
||||
import { EXAMPLE_10K_SEC_FILES, EXAMPLE_FILE } from "../helpers/datasources";
|
||||
import { askModelConfig } from "../helpers/providers";
|
||||
import { getTools } from "../helpers/tools";
|
||||
import { ModelConfig, TemplateFramework } from "../helpers/types";
|
||||
import { PureQuestionArgs, QuestionResults } from "./types";
|
||||
import { askPostInstallAction, questionHandlers } from "./utils";
|
||||
|
||||
type AppType =
|
||||
| "rag"
|
||||
| "code_artifact"
|
||||
| "financial_report_agent"
|
||||
| "form_filling"
|
||||
| "extractor"
|
||||
| "contract_review"
|
||||
| "data_scientist"
|
||||
| "deep_research";
|
||||
type AppType = "agentic_rag" | "financial_report" | "deep_research";
|
||||
|
||||
type SimpleAnswers = {
|
||||
appType: AppType;
|
||||
@@ -35,53 +22,22 @@ export const askSimpleQuestions = async (
|
||||
{
|
||||
type: "select",
|
||||
name: "appType",
|
||||
message: "What app do you want to build?",
|
||||
hint: "🤖: Agent, 🔀: Workflow",
|
||||
message: "What use case do you want to build?",
|
||||
choices: [
|
||||
{
|
||||
title: "🤖 Agentic RAG",
|
||||
value: "rag",
|
||||
title: "Agentic RAG",
|
||||
value: "agentic_rag",
|
||||
description:
|
||||
"Chatbot that answers questions based on provided documents.",
|
||||
},
|
||||
{
|
||||
title: "🤖 Data Scientist",
|
||||
value: "data_scientist",
|
||||
title: "Financial Report",
|
||||
value: "financial_report",
|
||||
description:
|
||||
"Agent that analyzes data and generates visualizations by using a code interpreter.",
|
||||
},
|
||||
{
|
||||
title: "🤖 Code Artifact Agent",
|
||||
value: "code_artifact",
|
||||
description:
|
||||
"Agent that writes code, runs it in a sandbox, and shows the output in the chat UI.",
|
||||
},
|
||||
{
|
||||
title: "🤖 Information Extractor",
|
||||
value: "extractor",
|
||||
description:
|
||||
"Extracts information from documents and returns it as a structured JSON object.",
|
||||
},
|
||||
{
|
||||
title: "🔀 Financial Report Generator",
|
||||
value: "financial_report_agent",
|
||||
description:
|
||||
"Generates a financial report by analyzing the provided 10-K SEC data. Uses a code interpreter to create charts or to conduct further analysis.",
|
||||
},
|
||||
{
|
||||
title: "🔀 Financial 10k SEC Form Filler",
|
||||
value: "form_filling",
|
||||
description:
|
||||
"Extracts information from 10k SEC data and uses it to fill out a CSV form.",
|
||||
},
|
||||
{
|
||||
title: "🔀 Contract Reviewer",
|
||||
value: "contract_review",
|
||||
description:
|
||||
"Extracts and reviews contracts to ensure compliance with GDPR regulations",
|
||||
},
|
||||
{
|
||||
title: "🔀 Deep Researcher",
|
||||
title: "Deep Research",
|
||||
value: "deep_research",
|
||||
description:
|
||||
"Researches and analyzes provided documents from multiple perspectives, generating a comprehensive report with citations to support key findings and insights.",
|
||||
@@ -93,13 +49,10 @@ export const askSimpleQuestions = async (
|
||||
|
||||
let language: TemplateFramework = "fastapi";
|
||||
let llamaCloudKey = args.llamaCloudKey;
|
||||
|
||||
let useLlamaCloud = false;
|
||||
|
||||
if (
|
||||
appType !== "extractor" &&
|
||||
appType !== "contract_review" &&
|
||||
appType !== "deep_research"
|
||||
) {
|
||||
if (appType !== "extractor" && appType !== "contract_review") {
|
||||
const { language: newLanguage } = await prompts(
|
||||
{
|
||||
type: "select",
|
||||
@@ -170,80 +123,36 @@ const convertAnswers = async (
|
||||
};
|
||||
const lookup: Record<
|
||||
AppType,
|
||||
Pick<
|
||||
QuestionResults,
|
||||
"template" | "tools" | "frontend" | "dataSources" | "useCase"
|
||||
> & {
|
||||
Pick<QuestionResults, "template" | "tools" | "dataSources" | "useCase"> & {
|
||||
modelConfig?: ModelConfig;
|
||||
}
|
||||
> = {
|
||||
rag: {
|
||||
template: "streaming",
|
||||
tools: getTools(["weather"]),
|
||||
frontend: true,
|
||||
agentic_rag: {
|
||||
template: "llamaindexserver",
|
||||
dataSources: [EXAMPLE_FILE],
|
||||
},
|
||||
data_scientist: {
|
||||
template: "streaming",
|
||||
financial_report: {
|
||||
template: "llamaindexserver",
|
||||
dataSources: EXAMPLE_10K_SEC_FILES,
|
||||
tools: getTools(["interpreter", "document_generator"]),
|
||||
frontend: true,
|
||||
dataSources: [],
|
||||
modelConfig: MODEL_GPT4o,
|
||||
},
|
||||
code_artifact: {
|
||||
template: "streaming",
|
||||
tools: getTools(["artifact"]),
|
||||
frontend: true,
|
||||
dataSources: [],
|
||||
modelConfig: MODEL_GPT4o,
|
||||
},
|
||||
financial_report_agent: {
|
||||
template: "multiagent",
|
||||
useCase: "financial_report",
|
||||
tools: getTools(["document_generator", "interpreter"]),
|
||||
dataSources: EXAMPLE_10K_SEC_FILES,
|
||||
frontend: true,
|
||||
modelConfig: MODEL_GPT4o,
|
||||
},
|
||||
form_filling: {
|
||||
template: "multiagent",
|
||||
useCase: "form_filling",
|
||||
tools: getTools(["form_filling"]),
|
||||
dataSources: EXAMPLE_10K_SEC_FILES,
|
||||
frontend: true,
|
||||
modelConfig: MODEL_GPT4o,
|
||||
},
|
||||
extractor: {
|
||||
template: "reflex",
|
||||
useCase: "extractor",
|
||||
tools: [],
|
||||
frontend: false,
|
||||
dataSources: [EXAMPLE_FILE],
|
||||
},
|
||||
contract_review: {
|
||||
template: "reflex",
|
||||
useCase: "contract_review",
|
||||
tools: [],
|
||||
frontend: false,
|
||||
dataSources: [EXAMPLE_GDPR],
|
||||
},
|
||||
deep_research: {
|
||||
template: "multiagent",
|
||||
useCase: "deep_research",
|
||||
template: "llamaindexserver",
|
||||
dataSources: EXAMPLE_10K_SEC_FILES,
|
||||
tools: [],
|
||||
frontend: true,
|
||||
dataSources: [AI_REPORTS],
|
||||
modelConfig: MODEL_GPT4o,
|
||||
},
|
||||
};
|
||||
|
||||
const results = lookup[answers.appType];
|
||||
return {
|
||||
framework: answers.language,
|
||||
useCase: answers.appType,
|
||||
ui: "shadcn",
|
||||
llamaCloudKey: answers.llamaCloudKey,
|
||||
useLlamaParse: answers.useLlamaCloud,
|
||||
llamapack: "",
|
||||
vectorDb: answers.useLlamaCloud ? "llamacloud" : "none",
|
||||
observability: "none",
|
||||
...results,
|
||||
modelConfig:
|
||||
results.modelConfig ??
|
||||
@@ -252,6 +161,6 @@ const convertAnswers = async (
|
||||
askModels: args.askModels ?? false,
|
||||
framework: answers.language,
|
||||
})),
|
||||
frontend: answers.language === "nextjs" ? false : results.frontend,
|
||||
frontend: true,
|
||||
};
|
||||
};
|
||||
|
||||
@@ -16,7 +16,10 @@ class AnalysisDecision(BaseModel):
|
||||
description="Whether to continue research, write a report, or cancel the research after several retries"
|
||||
)
|
||||
research_questions: Optional[List[str]] = Field(
|
||||
description="Questions to research if continuing research. Maximum 3 questions. Set to null or empty if writing a report.",
|
||||
description="""
|
||||
If the decision is to research, provide a list of questions to research that related to the user request.
|
||||
Maximum 3 questions. Set to null or empty if writing a report or cancel the research.
|
||||
""",
|
||||
default_factory=list,
|
||||
)
|
||||
cancel_reason: Optional[str] = Field(
|
||||
@@ -29,23 +32,23 @@ async def plan_research(
|
||||
memory: SimpleComposableMemory,
|
||||
context_nodes: List[Node],
|
||||
user_request: str,
|
||||
total_questions: int,
|
||||
) -> AnalysisDecision:
|
||||
analyze_prompt = PromptTemplate(
|
||||
"""
|
||||
analyze_prompt = """
|
||||
You are a professor who is guiding a researcher to research a specific request/problem.
|
||||
Your task is to decide on a research plan for the researcher.
|
||||
|
||||
The possible actions are:
|
||||
+ Provide a list of questions for the researcher to investigate, with the purpose of clarifying the request.
|
||||
+ Write a report if the researcher has already gathered enough research on the topic and can resolve the initial request.
|
||||
+ Cancel the research if most of the answers from researchers indicate there is insufficient information to research the request. Do not attempt more than 3 research iterations or too many questions.
|
||||
|
||||
The workflow should be:
|
||||
+ Always begin by providing some initial questions for the researcher to investigate.
|
||||
+ Analyze the provided answers against the initial topic/request. If the answers are insufficient to resolve the initial request, provide additional questions for the researcher to investigate.
|
||||
+ If the answers are sufficient to resolve the initial request, instruct the researcher to write a report.
|
||||
<User request>
|
||||
{user_request}
|
||||
</User request>
|
||||
|
||||
Here are the context:
|
||||
<Collected information>
|
||||
{context_str}
|
||||
</Collected information>
|
||||
@@ -53,8 +56,29 @@ async def plan_research(
|
||||
<Conversation context>
|
||||
{conversation_context}
|
||||
</Conversation context>
|
||||
|
||||
{enhanced_prompt}
|
||||
|
||||
Now, provide your decision in the required format for this user request:
|
||||
<User request>
|
||||
{user_request}
|
||||
</User request>
|
||||
"""
|
||||
)
|
||||
# Manually craft the prompt to avoid LLM hallucination
|
||||
enhanced_prompt = ""
|
||||
if total_questions == 0:
|
||||
# Avoid writing a report without any research context
|
||||
enhanced_prompt = """
|
||||
|
||||
The student has no questions to research. Let start by asking some questions.
|
||||
"""
|
||||
elif total_questions > 6:
|
||||
# Avoid asking too many questions (when the data is not ready for writing a report)
|
||||
enhanced_prompt = f"""
|
||||
|
||||
The student has researched {total_questions} questions. Should cancel the research if the context is not enough to write a report.
|
||||
"""
|
||||
|
||||
conversation_context = "\n".join(
|
||||
[f"{message.role}: {message.content}" for message in memory.get_all()]
|
||||
)
|
||||
@@ -63,10 +87,11 @@ async def plan_research(
|
||||
)
|
||||
res = await Settings.llm.astructured_predict(
|
||||
output_cls=AnalysisDecision,
|
||||
prompt=analyze_prompt,
|
||||
prompt=PromptTemplate(template=analyze_prompt),
|
||||
user_request=user_request,
|
||||
context_str=context_str,
|
||||
conversation_context=conversation_context,
|
||||
enhanced_prompt=enhanced_prompt,
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
@@ -32,7 +32,6 @@ logger.setLevel(logging.INFO)
|
||||
|
||||
|
||||
def create_workflow(
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Workflow:
|
||||
@@ -45,7 +44,6 @@ def create_workflow(
|
||||
|
||||
return DeepResearchWorkflow(
|
||||
index=index,
|
||||
chat_history=chat_history,
|
||||
timeout=120.0,
|
||||
)
|
||||
|
||||
@@ -73,27 +71,29 @@ class DeepResearchWorkflow(Workflow):
|
||||
def __init__(
|
||||
self,
|
||||
index: BaseIndex,
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
stream: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.index = index
|
||||
self.context_nodes = []
|
||||
self.stream = stream
|
||||
self.chat_history = chat_history
|
||||
self.memory = SimpleComposableMemory.from_defaults(
|
||||
primary_memory=ChatMemoryBuffer.from_defaults(
|
||||
chat_history=chat_history,
|
||||
),
|
||||
primary_memory=ChatMemoryBuffer.from_defaults(),
|
||||
)
|
||||
|
||||
@step
|
||||
def retrieve(self, ctx: Context, ev: StartEvent) -> PlanResearchEvent:
|
||||
async def retrieve(self, ctx: Context, ev: StartEvent) -> PlanResearchEvent:
|
||||
"""
|
||||
Initiate the workflow: memory, tools, agent
|
||||
"""
|
||||
self.user_request = ev.get("input")
|
||||
self.stream = ev.get("stream", True)
|
||||
self.user_request = ev.get("user_msg")
|
||||
chat_history = ev.get("chat_history")
|
||||
if chat_history is not None:
|
||||
self.memory.put_messages(chat_history)
|
||||
|
||||
await ctx.set("total_questions", 0)
|
||||
|
||||
# Add user message to memory
|
||||
self.memory.put_messages(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
@@ -132,9 +132,7 @@ class DeepResearchWorkflow(Workflow):
|
||||
nodes=nodes,
|
||||
)
|
||||
)
|
||||
return PlanResearchEvent(
|
||||
context_nodes=self.context_nodes,
|
||||
)
|
||||
return PlanResearchEvent()
|
||||
|
||||
@step
|
||||
async def analyze(
|
||||
@@ -153,10 +151,12 @@ class DeepResearchWorkflow(Workflow):
|
||||
},
|
||||
)
|
||||
)
|
||||
total_questions = await ctx.get("total_questions")
|
||||
res = await plan_research(
|
||||
memory=self.memory,
|
||||
context_nodes=self.context_nodes,
|
||||
user_request=self.user_request,
|
||||
total_questions=total_questions,
|
||||
)
|
||||
if res.decision == "cancel":
|
||||
ctx.write_event_to_stream(
|
||||
@@ -172,6 +172,22 @@ class DeepResearchWorkflow(Workflow):
|
||||
result=res.cancel_reason,
|
||||
)
|
||||
elif res.decision == "write":
|
||||
# Writing a report without any research context is not allowed.
|
||||
# It's a LLM hallucination.
|
||||
if total_questions == 0:
|
||||
ctx.write_event_to_stream(
|
||||
DataEvent(
|
||||
type="deep_research_event",
|
||||
data={
|
||||
"event": "analyze",
|
||||
"state": "done",
|
||||
},
|
||||
)
|
||||
)
|
||||
return StopEvent(
|
||||
result="Sorry, I have a problem when analyzing the retrieved information. Please try again.",
|
||||
)
|
||||
|
||||
self.memory.put(
|
||||
message=ChatMessage(
|
||||
role=MessageRole.ASSISTANT,
|
||||
@@ -180,7 +196,11 @@ class DeepResearchWorkflow(Workflow):
|
||||
)
|
||||
ctx.send_event(ReportEvent())
|
||||
else:
|
||||
await ctx.set("n_questions", len(res.research_questions))
|
||||
total_questions += len(res.research_questions)
|
||||
await ctx.set("total_questions", total_questions) # For tracking
|
||||
await ctx.set(
|
||||
"waiting_questions", len(res.research_questions)
|
||||
) # For waiting questions to be answered
|
||||
self.memory.put(
|
||||
message=ChatMessage(
|
||||
role=MessageRole.ASSISTANT,
|
||||
@@ -270,7 +290,7 @@ class DeepResearchWorkflow(Workflow):
|
||||
"""
|
||||
Collect answers to all questions
|
||||
"""
|
||||
num_questions = await ctx.get("n_questions")
|
||||
num_questions = await ctx.get("waiting_questions")
|
||||
results = ctx.collect_events(
|
||||
ev,
|
||||
expected=[CollectAnswersEvent] * num_questions,
|
||||
@@ -284,7 +304,7 @@ class DeepResearchWorkflow(Workflow):
|
||||
content=f"<Question>{result.question}</Question>\n<Answer>{result.answer}</Answer>",
|
||||
)
|
||||
)
|
||||
await ctx.set("n_questions", 0)
|
||||
await ctx.set("waiting_questions", 0)
|
||||
self.memory.put(
|
||||
message=ChatMessage(
|
||||
role=MessageRole.ASSISTANT,
|
||||
@@ -298,7 +318,6 @@ class DeepResearchWorkflow(Workflow):
|
||||
"""
|
||||
Report the answers
|
||||
"""
|
||||
logger.info("Writing the report")
|
||||
res = await write_report(
|
||||
memory=self.memory,
|
||||
user_request=self.user_request,
|
||||
|
||||
+24
-21
@@ -1,13 +1,5 @@
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.engine.index import IndexConfig, get_index
|
||||
from app.engine.tools import ToolFactory
|
||||
from app.engine.tools.query_engine import get_query_engine_tool
|
||||
from app.workflows.events import AgentRunEvent
|
||||
from app.workflows.tools import (
|
||||
call_tools,
|
||||
chat_with_tools,
|
||||
)
|
||||
from llama_index.core import Settings
|
||||
from llama_index.core.base.llms.types import ChatMessage, MessageRole
|
||||
from llama_index.core.llms.function_calling import FunctionCallingLLM
|
||||
@@ -22,9 +14,17 @@ from llama_index.core.workflow import (
|
||||
step,
|
||||
)
|
||||
|
||||
from app.engine.index import IndexConfig, get_index
|
||||
from app.engine.tools import ToolFactory
|
||||
from app.engine.tools.query_engine import get_query_engine_tool
|
||||
from app.workflows.events import AgentRunEvent
|
||||
from app.workflows.tools import (
|
||||
call_tools,
|
||||
chat_with_tools,
|
||||
)
|
||||
|
||||
|
||||
def create_workflow(
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Workflow:
|
||||
@@ -45,7 +45,6 @@ def create_workflow(
|
||||
query_engine_tool=query_engine_tool,
|
||||
code_interpreter_tool=code_interpreter_tool,
|
||||
document_generator_tool=document_generator_tool,
|
||||
chat_history=chat_history,
|
||||
)
|
||||
|
||||
|
||||
@@ -91,6 +90,7 @@ class FinancialReportWorkflow(Workflow):
|
||||
It's good to using appropriate tools for the user request and always use the information from the tools, don't make up anything yourself.
|
||||
For the query engine tool, you should break down the user request into a list of queries and call the tool with the queries.
|
||||
"""
|
||||
stream: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -99,12 +99,10 @@ class FinancialReportWorkflow(Workflow):
|
||||
document_generator_tool: FunctionTool,
|
||||
llm: Optional[FunctionCallingLLM] = None,
|
||||
timeout: int = 360,
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
):
|
||||
super().__init__(timeout=timeout)
|
||||
self.system_prompt = system_prompt or self._default_system_prompt
|
||||
self.chat_history = chat_history or []
|
||||
self.query_engine_tool = query_engine_tool
|
||||
self.code_interpreter_tool = code_interpreter_tool
|
||||
self.document_generator_tool = document_generator_tool
|
||||
@@ -122,13 +120,19 @@ class FinancialReportWorkflow(Workflow):
|
||||
]
|
||||
self.llm: FunctionCallingLLM = llm or Settings.llm
|
||||
assert isinstance(self.llm, FunctionCallingLLM)
|
||||
self.memory = ChatMemoryBuffer.from_defaults(
|
||||
llm=self.llm, chat_history=self.chat_history
|
||||
)
|
||||
self.memory = ChatMemoryBuffer.from_defaults(llm=self.llm)
|
||||
|
||||
@step()
|
||||
async def prepare_chat_history(self, ctx: Context, ev: StartEvent) -> InputEvent:
|
||||
ctx.data["input"] = ev.input
|
||||
self.stream = ev.get("stream", True)
|
||||
user_msg = ev.get("user_msg")
|
||||
chat_history = ev.get("chat_history")
|
||||
|
||||
if chat_history is not None:
|
||||
self.memory.put_messages(chat_history)
|
||||
|
||||
# Add user message to memory
|
||||
self.memory.put(ChatMessage(role=MessageRole.USER, content=user_msg))
|
||||
|
||||
if self.system_prompt:
|
||||
system_msg = ChatMessage(
|
||||
@@ -136,9 +140,6 @@ class FinancialReportWorkflow(Workflow):
|
||||
)
|
||||
self.memory.put(system_msg)
|
||||
|
||||
# Add user input to memory
|
||||
self.memory.put(ChatMessage(role=MessageRole.USER, content=ev.input))
|
||||
|
||||
return InputEvent(input=self.memory.get())
|
||||
|
||||
@step()
|
||||
@@ -160,8 +161,10 @@ class FinancialReportWorkflow(Workflow):
|
||||
chat_history,
|
||||
)
|
||||
if not response.has_tool_calls():
|
||||
# If no tool call, return the response generator
|
||||
return StopEvent(result=response.generator)
|
||||
if self.stream:
|
||||
return StopEvent(result=response.generator)
|
||||
else:
|
||||
return StopEvent(result=await response.full_response())
|
||||
# calling different tools at the same time is not supported at the moment
|
||||
# add an error message to tell the AI to process step by step
|
||||
if response.is_calling_different_tools():
|
||||
|
||||
@@ -25,7 +25,6 @@ from app.workflows.tools import (
|
||||
|
||||
|
||||
def create_workflow(
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Workflow:
|
||||
@@ -45,7 +44,6 @@ def create_workflow(
|
||||
query_engine_tool=query_engine_tool,
|
||||
extractor_tool=extractor_tool, # type: ignore
|
||||
filling_tool=filling_tool, # type: ignore
|
||||
chat_history=chat_history,
|
||||
)
|
||||
|
||||
return workflow
|
||||
@@ -88,6 +86,7 @@ class FormFillingWorkflow(Workflow):
|
||||
Only use provided data - never make up any information yourself. Fill N/A if an answer is not found.
|
||||
If there is no query engine tool or the gathered information has many N/A values indicating the questions don't match the data, respond with a warning and ask the user to upload a different file or connect to a knowledge base.
|
||||
"""
|
||||
stream: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -96,12 +95,10 @@ class FormFillingWorkflow(Workflow):
|
||||
filling_tool: FunctionTool,
|
||||
llm: Optional[FunctionCallingLLM] = None,
|
||||
timeout: int = 360,
|
||||
chat_history: Optional[List[ChatMessage]] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
):
|
||||
super().__init__(timeout=timeout)
|
||||
self.system_prompt = system_prompt or self._default_system_prompt
|
||||
self.chat_history = chat_history or []
|
||||
self.query_engine_tool = query_engine_tool
|
||||
self.extractor_tool = extractor_tool
|
||||
self.filling_tool = filling_tool
|
||||
@@ -113,13 +110,18 @@ class FormFillingWorkflow(Workflow):
|
||||
self.llm: FunctionCallingLLM = llm or Settings.llm
|
||||
if not isinstance(self.llm, FunctionCallingLLM):
|
||||
raise ValueError("FormFillingWorkflow only supports FunctionCallingLLM.")
|
||||
self.memory = ChatMemoryBuffer.from_defaults(
|
||||
llm=self.llm, chat_history=self.chat_history
|
||||
)
|
||||
self.memory = ChatMemoryBuffer.from_defaults(llm=self.llm)
|
||||
|
||||
@step()
|
||||
async def start(self, ctx: Context, ev: StartEvent) -> InputEvent:
|
||||
ctx.data["input"] = ev.input
|
||||
self.stream = ev.get("stream", True)
|
||||
user_msg = ev.get("user_msg", "")
|
||||
chat_history = ev.get("chat_history", [])
|
||||
|
||||
if chat_history:
|
||||
self.memory.put_messages(chat_history)
|
||||
|
||||
self.memory.put(ChatMessage(role=MessageRole.USER, content=user_msg))
|
||||
|
||||
if self.system_prompt:
|
||||
system_msg = ChatMessage(
|
||||
@@ -127,12 +129,7 @@ class FormFillingWorkflow(Workflow):
|
||||
)
|
||||
self.memory.put(system_msg)
|
||||
|
||||
user_input = ev.input
|
||||
user_msg = ChatMessage(role=MessageRole.USER, content=user_input)
|
||||
self.memory.put(user_msg)
|
||||
|
||||
chat_history = self.memory.get()
|
||||
return InputEvent(input=chat_history)
|
||||
return InputEvent(input=self.memory.get())
|
||||
|
||||
@step()
|
||||
async def handle_llm_input( # type: ignore
|
||||
@@ -150,7 +147,10 @@ class FormFillingWorkflow(Workflow):
|
||||
chat_history,
|
||||
)
|
||||
if not response.has_tool_calls():
|
||||
return StopEvent(result=response.generator)
|
||||
if self.stream:
|
||||
return StopEvent(result=response.generator)
|
||||
else:
|
||||
return StopEvent(result=await response.full_response())
|
||||
# calling different tools at the same time is not supported at the moment
|
||||
# add an error message to tell the AI to process step by step
|
||||
if response.is_calling_different_tools():
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import { BaseToolWithCall } from "llamaindex";
|
||||
import { ToolsFactory } from "llamaindex/tools/ToolsFactory";
|
||||
import fs from "node:fs/promises";
|
||||
import path from "node:path";
|
||||
import { CodeGeneratorTool, CodeGeneratorToolParams } from "./code-generator";
|
||||
@@ -18,6 +17,7 @@ import { ImgGeneratorTool, ImgGeneratorToolParams } from "./img-gen";
|
||||
import { InterpreterTool, InterpreterToolParams } from "./interpreter";
|
||||
import { OpenAPIActionTool } from "./openapi-action";
|
||||
import { WeatherTool, WeatherToolParams } from "./weather";
|
||||
import { WikipediaTool, WikipediaToolParams } from "./wikipedia";
|
||||
|
||||
type ToolCreator = (config: unknown) => Promise<BaseToolWithCall[]>;
|
||||
|
||||
@@ -27,12 +27,13 @@ export async function createTools(toolConfig: {
|
||||
}): Promise<BaseToolWithCall[]> {
|
||||
// add local tools from the 'tools' folder (if configured)
|
||||
const tools = await createLocalTools(toolConfig.local);
|
||||
// add tools from LlamaIndexTS (if configured)
|
||||
tools.push(...(await ToolsFactory.createTools(toolConfig.llamahub)));
|
||||
return tools;
|
||||
}
|
||||
|
||||
const toolFactory: Record<string, ToolCreator> = {
|
||||
"wikipedia.WikipediaToolSpec": async (config: unknown) => {
|
||||
return [new WikipediaTool(config as WikipediaToolParams)];
|
||||
},
|
||||
weather: async (config: unknown) => {
|
||||
return [new WeatherTool(config as WeatherToolParams)];
|
||||
},
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
import type { JSONSchemaType } from "ajv";
|
||||
import type { BaseTool, ToolMetadata } from "llamaindex";
|
||||
import { default as wiki } from "wikipedia";
|
||||
|
||||
type WikipediaParameter = {
|
||||
query: string;
|
||||
lang?: string;
|
||||
};
|
||||
|
||||
export type WikipediaToolParams = {
|
||||
metadata?: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
|
||||
};
|
||||
|
||||
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<WikipediaParameter>> = {
|
||||
name: "wikipedia_tool",
|
||||
description: "A tool that uses a query engine to search Wikipedia.",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: {
|
||||
query: {
|
||||
type: "string",
|
||||
description: "The query to search for",
|
||||
},
|
||||
lang: {
|
||||
type: "string",
|
||||
description: "The language to search in",
|
||||
nullable: true,
|
||||
},
|
||||
},
|
||||
required: ["query"],
|
||||
},
|
||||
};
|
||||
|
||||
export class WikipediaTool implements BaseTool<WikipediaParameter> {
|
||||
private readonly DEFAULT_LANG = "en";
|
||||
metadata: ToolMetadata<JSONSchemaType<WikipediaParameter>>;
|
||||
|
||||
constructor(params?: WikipediaToolParams) {
|
||||
this.metadata = params?.metadata || DEFAULT_META_DATA;
|
||||
}
|
||||
|
||||
async loadData(
|
||||
page: string,
|
||||
lang: string = this.DEFAULT_LANG,
|
||||
): Promise<string> {
|
||||
wiki.setLang(lang);
|
||||
const pageResult = await wiki.page(page, { autoSuggest: false });
|
||||
const content = await pageResult.content();
|
||||
return content;
|
||||
}
|
||||
|
||||
async call({
|
||||
query,
|
||||
lang = this.DEFAULT_LANG,
|
||||
}: WikipediaParameter): Promise<string> {
|
||||
const searchResult = await wiki.search(query);
|
||||
if (searchResult.results.length === 0) return "No search results.";
|
||||
return await this.loadData(searchResult.results[0].title, lang);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
import {
|
||||
FILE_EXT_TO_READER,
|
||||
SimpleDirectoryReader,
|
||||
} from "llamaindex/readers/index";
|
||||
} from "@llamaindex/readers/directory";
|
||||
|
||||
export const DATA_DIR = "./data";
|
||||
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import { LlamaParseReader } from "llamaindex";
|
||||
import {
|
||||
FILE_EXT_TO_READER,
|
||||
SimpleDirectoryReader,
|
||||
} from "llamaindex/readers/index";
|
||||
} from "@llamaindex/readers/directory";
|
||||
import { LlamaParseReader } from "llamaindex";
|
||||
|
||||
export const DATA_DIR = "./data";
|
||||
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class EventCallback(ABC):
|
||||
"""
|
||||
Base class for event callbacks during event streaming.
|
||||
"""
|
||||
|
||||
async def run(self, event: Any) -> Any:
|
||||
"""
|
||||
Called for each event in the stream.
|
||||
Default behavior: pass through the event unchanged.
|
||||
"""
|
||||
return event
|
||||
|
||||
async def on_complete(self, final_response: str) -> Any:
|
||||
"""
|
||||
Called when the stream is complete.
|
||||
Default behavior: return None.
|
||||
"""
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def from_default(self, *args, **kwargs) -> "EventCallback":
|
||||
"""
|
||||
Create a new instance of the processor from default values.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,42 @@
|
||||
import logging
|
||||
from typing import Any, List
|
||||
|
||||
from fastapi import BackgroundTasks
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
|
||||
from app.api.callbacks.base import EventCallback
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class LlamaCloudFileDownload(EventCallback):
|
||||
"""
|
||||
Processor for handling LlamaCloud file downloads from source nodes.
|
||||
Only work if LlamaCloud service code is available.
|
||||
"""
|
||||
|
||||
def __init__(self, background_tasks: BackgroundTasks):
|
||||
self.background_tasks = background_tasks
|
||||
|
||||
async def run(self, event: Any) -> Any:
|
||||
if hasattr(event, "to_response"):
|
||||
event_response = event.to_response()
|
||||
if event_response.get("type") == "sources" and hasattr(event, "nodes"):
|
||||
await self._process_response_nodes(event.nodes)
|
||||
return event
|
||||
|
||||
async def _process_response_nodes(self, source_nodes: List[NodeWithScore]):
|
||||
try:
|
||||
from app.engine.service import LLamaCloudFileService # type: ignore
|
||||
|
||||
LLamaCloudFileService.download_files_from_nodes(
|
||||
source_nodes, self.background_tasks
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_default(
|
||||
cls, background_tasks: BackgroundTasks
|
||||
) -> "LlamaCloudFileDownload":
|
||||
return cls(background_tasks=background_tasks)
|
||||
@@ -0,0 +1,34 @@
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from app.api.callbacks.base import EventCallback
|
||||
from app.api.routers.models import ChatData
|
||||
from app.api.services.suggestion import NextQuestionSuggestion
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class SuggestNextQuestions(EventCallback):
|
||||
"""Processor for generating next question suggestions."""
|
||||
|
||||
def __init__(self, chat_data: ChatData):
|
||||
self.chat_data = chat_data
|
||||
self.accumulated_text = ""
|
||||
|
||||
async def on_complete(self, final_response: str) -> Any:
|
||||
if final_response == "":
|
||||
return None
|
||||
|
||||
questions = await NextQuestionSuggestion.suggest_next_questions(
|
||||
self.chat_data.messages, final_response
|
||||
)
|
||||
if questions:
|
||||
return {
|
||||
"type": "suggested_questions",
|
||||
"data": questions,
|
||||
}
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_default(cls, chat_data: ChatData) -> "SuggestNextQuestions":
|
||||
return cls(chat_data=chat_data)
|
||||
@@ -0,0 +1,66 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from llama_index.core.workflow.handler import WorkflowHandler
|
||||
|
||||
from app.api.callbacks.base import EventCallback
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""
|
||||
Streams events from a workflow handler through a chain of callbacks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow_handler: WorkflowHandler,
|
||||
callbacks: Optional[List[EventCallback]] = None,
|
||||
):
|
||||
self.workflow_handler = workflow_handler
|
||||
self.callbacks = callbacks or []
|
||||
self.accumulated_text = ""
|
||||
|
||||
def vercel_stream(self):
|
||||
"""Create a streaming response with Vercel format."""
|
||||
from app.api.routers.vercel_response import VercelStreamResponse
|
||||
|
||||
return VercelStreamResponse(stream_handler=self)
|
||||
|
||||
async def cancel_run(self):
|
||||
"""Cancel the workflow handler."""
|
||||
await self.workflow_handler.cancel_run()
|
||||
|
||||
async def stream_events(self):
|
||||
"""Stream events through the processor chain."""
|
||||
try:
|
||||
async for event in self.workflow_handler.stream_events():
|
||||
# Process the event through each processor
|
||||
for callback in self.callbacks:
|
||||
event = await callback.run(event)
|
||||
yield event
|
||||
|
||||
# After all events are processed, call on_complete for each callback
|
||||
for callback in self.callbacks:
|
||||
result = await callback.on_complete(self.accumulated_text)
|
||||
if result:
|
||||
yield result
|
||||
|
||||
except Exception as e:
|
||||
# Make sure to cancel the workflow on error
|
||||
await self.workflow_handler.cancel_run()
|
||||
raise e
|
||||
|
||||
async def accumulate_text(self, text: str):
|
||||
"""Accumulate text from the workflow handler."""
|
||||
self.accumulated_text += text
|
||||
|
||||
@classmethod
|
||||
def from_default(
|
||||
cls,
|
||||
handler: WorkflowHandler,
|
||||
callbacks: Optional[List[EventCallback]] = None,
|
||||
) -> "StreamHandler":
|
||||
"""Create a new instance with the given workflow handler and callbacks."""
|
||||
return cls(workflow_handler=handler, callbacks=callbacks)
|
||||
@@ -2,10 +2,12 @@ import logging
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, status
|
||||
|
||||
from app.api.callbacks.llamacloud import LlamaCloudFileDownload
|
||||
from app.api.callbacks.next_question import SuggestNextQuestions
|
||||
from app.api.callbacks.stream_handler import StreamHandler
|
||||
from app.api.routers.models import (
|
||||
ChatData,
|
||||
)
|
||||
from app.api.routers.vercel_response import VercelStreamResponse
|
||||
from app.engine.query_filter import generate_filters
|
||||
from app.workflows import create_workflow
|
||||
|
||||
@@ -29,19 +31,22 @@ async def chat(
|
||||
params = data.data or {}
|
||||
|
||||
workflow = create_workflow(
|
||||
chat_history=messages,
|
||||
params=params,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
event_handler = workflow.run(input=last_message_content, streaming=True)
|
||||
return VercelStreamResponse(
|
||||
request=request,
|
||||
chat_data=data,
|
||||
background_tasks=background_tasks,
|
||||
event_handler=event_handler,
|
||||
events=workflow.stream_events(),
|
||||
handler = workflow.run(
|
||||
user_msg=last_message_content,
|
||||
chat_history=messages,
|
||||
stream=True,
|
||||
)
|
||||
return StreamHandler.from_default(
|
||||
handler=handler,
|
||||
callbacks=[
|
||||
LlamaCloudFileDownload.from_default(background_tasks),
|
||||
SuggestNextQuestions.from_default(data),
|
||||
],
|
||||
).vercel_stream()
|
||||
except Exception as e:
|
||||
logger.exception("Error in chat engine", exc_info=True)
|
||||
raise HTTPException(
|
||||
|
||||
@@ -1,22 +1,20 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import AsyncGenerator, Awaitable, List
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from aiostream import stream
|
||||
from fastapi import BackgroundTasks, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
from llama_index.core.agent.workflow.workflow_events import AgentStream
|
||||
from llama_index.core.workflow import StopEvent
|
||||
|
||||
from app.api.routers.models import ChatData, Message
|
||||
from app.api.services.suggestion import NextQuestionSuggestion
|
||||
from app.api.callbacks.stream_handler import StreamHandler
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
class VercelStreamResponse(StreamingResponse):
|
||||
"""
|
||||
Base class to convert the response from the chat engine to the streaming format expected by Vercel
|
||||
Converts preprocessed events into Vercel-compatible streaming response format.
|
||||
"""
|
||||
|
||||
TEXT_PREFIX = "0:"
|
||||
@@ -25,136 +23,77 @@ class VercelStreamResponse(StreamingResponse):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request: Request,
|
||||
chat_data: ChatData,
|
||||
background_tasks: BackgroundTasks,
|
||||
stream_handler: StreamHandler,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self.request = request
|
||||
self.chat_data = chat_data
|
||||
self.background_tasks = background_tasks
|
||||
content = self.content_generator(*args, **kwargs)
|
||||
super().__init__(content=content)
|
||||
self.handler = stream_handler
|
||||
super().__init__(content=self.content_generator())
|
||||
|
||||
async def content_generator(self, event_handler, events):
|
||||
stream = self._create_stream(
|
||||
self.request, self.chat_data, event_handler, events
|
||||
)
|
||||
is_stream_started = False
|
||||
async def content_generator(self):
|
||||
"""Generate Vercel-formatted content from preprocessed events."""
|
||||
stream_started = False
|
||||
try:
|
||||
async with stream.stream() as streamer:
|
||||
async for output in streamer:
|
||||
if not is_stream_started:
|
||||
is_stream_started = True
|
||||
# Stream a blank message to start the stream
|
||||
yield self.convert_text("")
|
||||
async for event in self.handler.stream_events():
|
||||
if not stream_started:
|
||||
# Start the stream with an empty message
|
||||
stream_started = True
|
||||
yield self.convert_text("")
|
||||
|
||||
yield output
|
||||
except asyncio.CancelledError:
|
||||
logger.warning("Workflow has been cancelled!")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Unexpected error in content_generator: {str(e)}", exc_info=True
|
||||
)
|
||||
yield self.convert_error(
|
||||
"An unexpected error occurred while processing your request, preventing the creation of a final answer. Please try again."
|
||||
)
|
||||
finally:
|
||||
await event_handler.cancel_run()
|
||||
logger.info("The stream has been stopped!")
|
||||
|
||||
def _create_stream(
|
||||
self,
|
||||
request: Request,
|
||||
chat_data: ChatData,
|
||||
event_handler: Awaitable,
|
||||
events: AsyncGenerator,
|
||||
verbose: bool = True,
|
||||
):
|
||||
# Yield the text response
|
||||
async def _chat_response_generator():
|
||||
result = await event_handler
|
||||
final_response = ""
|
||||
|
||||
if isinstance(result, AsyncGenerator):
|
||||
async for token in result:
|
||||
final_response += str(token.delta)
|
||||
yield self.convert_text(token.delta)
|
||||
else:
|
||||
if hasattr(result, "response"):
|
||||
content = result.response.message.content
|
||||
if content:
|
||||
for token in content:
|
||||
final_response += str(token)
|
||||
yield self.convert_text(token)
|
||||
else:
|
||||
final_response += str(result)
|
||||
yield self.convert_text(result)
|
||||
|
||||
# Generate next questions if next question prompt is configured
|
||||
question_data = await self._generate_next_questions(
|
||||
chat_data.messages, final_response
|
||||
)
|
||||
if question_data:
|
||||
yield self.convert_data(question_data)
|
||||
|
||||
# Yield the events from the event handler
|
||||
async def _event_generator():
|
||||
async for event in events:
|
||||
event_response = event.to_response()
|
||||
if verbose:
|
||||
logger.debug(event_response)
|
||||
if event_response is not None:
|
||||
# Handle different types of events
|
||||
if isinstance(event, (AgentStream, StopEvent)):
|
||||
async for chunk in self._stream_text(event):
|
||||
await self.handler.accumulate_text(chunk)
|
||||
yield self.convert_text(chunk)
|
||||
elif isinstance(event, dict):
|
||||
yield self.convert_data(event)
|
||||
elif hasattr(event, "to_response"):
|
||||
event_response = event.to_response()
|
||||
yield self.convert_data(event_response)
|
||||
if event_response.get("type") == "sources":
|
||||
self._process_response_nodes(event.nodes, self.background_tasks)
|
||||
else:
|
||||
yield self.convert_data(event.model_dump())
|
||||
|
||||
combine = stream.merge(_chat_response_generator(), _event_generator())
|
||||
return combine
|
||||
except asyncio.CancelledError:
|
||||
logger.warning("Client cancelled the request!")
|
||||
await self.handler.cancel_run()
|
||||
except Exception as e:
|
||||
logger.error(f"Error in stream response: {e}")
|
||||
yield self.convert_error(str(e))
|
||||
await self.handler.cancel_run()
|
||||
|
||||
@staticmethod
|
||||
def _process_response_nodes(
|
||||
source_nodes: List[NodeWithScore],
|
||||
background_tasks: BackgroundTasks,
|
||||
):
|
||||
try:
|
||||
# Start background tasks to download documents from LlamaCloud if needed
|
||||
from app.engine.service import LLamaCloudFileService # type: ignore
|
||||
|
||||
LLamaCloudFileService.download_files_from_nodes(
|
||||
source_nodes, background_tasks
|
||||
)
|
||||
except ImportError:
|
||||
logger.debug(
|
||||
"LlamaCloud is not configured. Skipping post processing of nodes"
|
||||
)
|
||||
pass
|
||||
async def _stream_text(
|
||||
self, event: AgentStream | StopEvent
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Accept stream text from either AgentStream or StopEvent with string or AsyncGenerator result
|
||||
"""
|
||||
if isinstance(event, AgentStream):
|
||||
yield self.convert_text(event.delta)
|
||||
elif isinstance(event, StopEvent):
|
||||
if isinstance(event.result, str):
|
||||
yield event.result
|
||||
elif isinstance(event.result, AsyncGenerator):
|
||||
async for chunk in event.result:
|
||||
if isinstance(chunk, str):
|
||||
yield chunk
|
||||
elif hasattr(chunk, "delta"):
|
||||
yield chunk.delta
|
||||
|
||||
@classmethod
|
||||
def convert_text(cls, token: str):
|
||||
def convert_text(cls, token: str) -> str:
|
||||
"""Convert text event to Vercel format."""
|
||||
# Escape newlines and double quotes to avoid breaking the stream
|
||||
token = json.dumps(token)
|
||||
return f"{cls.TEXT_PREFIX}{token}\n"
|
||||
|
||||
@classmethod
|
||||
def convert_data(cls, data: dict):
|
||||
def convert_data(cls, data: dict) -> str:
|
||||
"""Convert data event to Vercel format."""
|
||||
data_str = json.dumps(data)
|
||||
return f"{cls.DATA_PREFIX}[{data_str}]\n"
|
||||
|
||||
@classmethod
|
||||
def convert_error(cls, error: str):
|
||||
def convert_error(cls, error: str) -> str:
|
||||
"""Convert error event to Vercel format."""
|
||||
error_str = json.dumps(error)
|
||||
return f"{cls.ERROR_PREFIX}{error_str}\n"
|
||||
|
||||
@staticmethod
|
||||
async def _generate_next_questions(chat_history: List[Message], response: str):
|
||||
questions = await NextQuestionSuggestion.suggest_next_questions(
|
||||
chat_history, response
|
||||
)
|
||||
if questions:
|
||||
return {
|
||||
"type": "suggested_questions",
|
||||
"data": questions,
|
||||
}
|
||||
return None
|
||||
|
||||
@@ -3,7 +3,6 @@ import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, AsyncGenerator, Callable, Optional
|
||||
|
||||
from app.workflows.events import AgentRunEvent, AgentRunEventType
|
||||
from llama_index.core.base.llms.types import ChatMessage, ChatResponse, MessageRole
|
||||
from llama_index.core.llms.function_calling import FunctionCallingLLM
|
||||
from llama_index.core.tools import (
|
||||
@@ -15,6 +14,8 @@ from llama_index.core.tools import (
|
||||
from llama_index.core.workflow import Context
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from app.workflows.events import AgentRunEvent, AgentRunEventType
|
||||
|
||||
logger = logging.getLogger("uvicorn")
|
||||
|
||||
|
||||
@@ -51,7 +52,9 @@ class ChatWithToolsResponse(BaseModel):
|
||||
assert self.generator is not None
|
||||
full_response = ""
|
||||
async for chunk in self.generator:
|
||||
full_response += chunk.message.content
|
||||
content = chunk.message.content
|
||||
if content:
|
||||
full_response += content
|
||||
return full_response
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
import {
|
||||
ALL_AVAILABLE_ANTHROPIC_MODELS,
|
||||
Anthropic,
|
||||
} from "@llamaindex/anthropic";
|
||||
import { HuggingFaceEmbedding } from "@llamaindex/huggingface";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
const embedModelMap: Record<string, string> = {
|
||||
"all-MiniLM-L6-v2": "Xenova/all-MiniLM-L6-v2",
|
||||
"all-mpnet-base-v2": "Xenova/all-mpnet-base-v2",
|
||||
};
|
||||
Settings.llm = new Anthropic({
|
||||
model: process.env.MODEL as keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS,
|
||||
});
|
||||
Settings.embedModel = new HuggingFaceEmbedding({
|
||||
modelType: embedModelMap[process.env.EMBEDDING_MODEL!],
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
import { OpenAI, OpenAIEmbedding } from "@llamaindex/openai";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
// Map Azure OpenAI model names to OpenAI model names (only for TS)
|
||||
const AZURE_OPENAI_MODEL_MAP: Record<string, string> = {
|
||||
"gpt-35-turbo": "gpt-3.5-turbo",
|
||||
"gpt-35-turbo-16k": "gpt-3.5-turbo-16k",
|
||||
"gpt-4o": "gpt-4o",
|
||||
"gpt-4": "gpt-4",
|
||||
"gpt-4-32k": "gpt-4-32k",
|
||||
"gpt-4-turbo": "gpt-4-turbo",
|
||||
"gpt-4-turbo-2024-04-09": "gpt-4-turbo",
|
||||
"gpt-4-vision-preview": "gpt-4-vision-preview",
|
||||
"gpt-4-1106-preview": "gpt-4-1106-preview",
|
||||
"gpt-4o-2024-05-13": "gpt-4o-2024-05-13",
|
||||
};
|
||||
|
||||
const azureConfig = {
|
||||
apiKey: process.env.AZURE_OPENAI_KEY,
|
||||
endpoint: process.env.AZURE_OPENAI_ENDPOINT,
|
||||
apiVersion:
|
||||
process.env.AZURE_OPENAI_API_VERSION || process.env.OPENAI_API_VERSION,
|
||||
};
|
||||
|
||||
Settings.llm = new OpenAI({
|
||||
model:
|
||||
AZURE_OPENAI_MODEL_MAP[process.env.MODEL ?? "gpt-35-turbo"] ??
|
||||
"gpt-3.5-turbo",
|
||||
maxTokens: process.env.LLM_MAX_TOKENS
|
||||
? Number(process.env.LLM_MAX_TOKENS)
|
||||
: undefined,
|
||||
azure: {
|
||||
...azureConfig,
|
||||
deployment: process.env.AZURE_OPENAI_LLM_DEPLOYMENT,
|
||||
},
|
||||
});
|
||||
|
||||
Settings.embedModel = new OpenAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL,
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: undefined,
|
||||
azure: {
|
||||
...azureConfig,
|
||||
deployment: process.env.AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import {
|
||||
Gemini,
|
||||
GEMINI_EMBEDDING_MODEL,
|
||||
GEMINI_MODEL,
|
||||
GeminiEmbedding,
|
||||
} from "@llamaindex/google";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
Settings.llm = new Gemini({
|
||||
model: process.env.MODEL as GEMINI_MODEL,
|
||||
});
|
||||
Settings.embedModel = new GeminiEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL as GEMINI_EMBEDDING_MODEL,
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
import { Groq } from "@llamaindex/groq";
|
||||
import { HuggingFaceEmbedding } from "@llamaindex/huggingface";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
const embedModelMap: Record<string, string> = {
|
||||
"all-MiniLM-L6-v2": "Xenova/all-MiniLM-L6-v2",
|
||||
"all-mpnet-base-v2": "Xenova/all-mpnet-base-v2",
|
||||
};
|
||||
|
||||
Settings.llm = new Groq({
|
||||
model: process.env.MODEL!,
|
||||
});
|
||||
|
||||
Settings.embedModel = new HuggingFaceEmbedding({
|
||||
modelType: embedModelMap[process.env.EMBEDDING_MODEL!],
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import {
|
||||
ALL_AVAILABLE_MISTRAL_MODELS,
|
||||
MistralAI,
|
||||
MistralAIEmbedding,
|
||||
MistralAIEmbeddingModelType,
|
||||
} from "@llamaindex/mistral";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
Settings.llm = new MistralAI({
|
||||
model: process.env.MODEL as keyof typeof ALL_AVAILABLE_MISTRAL_MODELS,
|
||||
});
|
||||
Settings.embedModel = new MistralAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL as MistralAIEmbeddingModelType,
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import { Ollama, OllamaEmbedding } from "@llamaindex/ollama";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
const config = {
|
||||
host: process.env.OLLAMA_BASE_URL ?? "http://127.0.0.1:11434",
|
||||
};
|
||||
Settings.llm = new Ollama({
|
||||
model: process.env.MODEL ?? "",
|
||||
config,
|
||||
});
|
||||
Settings.embedModel = new OllamaEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL ?? "",
|
||||
config,
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
import { OpenAI, OpenAIEmbedding } from "@llamaindex/openai";
|
||||
import { Settings } from "llamaindex";
|
||||
|
||||
export function setupProvider() {
|
||||
Settings.llm = new OpenAI({
|
||||
model: process.env.MODEL ?? "gpt-4o-mini",
|
||||
maxTokens: process.env.LLM_MAX_TOKENS
|
||||
? Number(process.env.LLM_MAX_TOKENS)
|
||||
: undefined,
|
||||
});
|
||||
Settings.embedModel = new OpenAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL,
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: undefined,
|
||||
});
|
||||
}
|
||||
@@ -1,179 +1,18 @@
|
||||
import {
|
||||
ALL_AVAILABLE_MISTRAL_MODELS,
|
||||
Anthropic,
|
||||
GEMINI_EMBEDDING_MODEL,
|
||||
GEMINI_MODEL,
|
||||
Gemini,
|
||||
GeminiEmbedding,
|
||||
Groq,
|
||||
MistralAI,
|
||||
MistralAIEmbedding,
|
||||
MistralAIEmbeddingModelType,
|
||||
OpenAI,
|
||||
OpenAIEmbedding,
|
||||
Settings,
|
||||
} from "llamaindex";
|
||||
import { HuggingFaceEmbedding } from "llamaindex/embeddings/HuggingFaceEmbedding";
|
||||
import { OllamaEmbedding } from "llamaindex/embeddings/OllamaEmbedding";
|
||||
import { ALL_AVAILABLE_ANTHROPIC_MODELS } from "llamaindex/llm/anthropic";
|
||||
import { Ollama } from "llamaindex/llm/ollama";
|
||||
import { Settings } from "llamaindex";
|
||||
import { setupProvider } from "./provider";
|
||||
|
||||
const CHUNK_SIZE = 512;
|
||||
const CHUNK_OVERLAP = 20;
|
||||
|
||||
export const initSettings = async () => {
|
||||
// HINT: you can delete the initialization code for unused model providers
|
||||
console.log(`Using '${process.env.MODEL_PROVIDER}' model provider`);
|
||||
|
||||
if (!process.env.MODEL || !process.env.EMBEDDING_MODEL) {
|
||||
throw new Error("'MODEL' and 'EMBEDDING_MODEL' env variables must be set.");
|
||||
}
|
||||
|
||||
switch (process.env.MODEL_PROVIDER) {
|
||||
case "ollama":
|
||||
initOllama();
|
||||
break;
|
||||
case "groq":
|
||||
initGroq();
|
||||
break;
|
||||
case "anthropic":
|
||||
initAnthropic();
|
||||
break;
|
||||
case "gemini":
|
||||
initGemini();
|
||||
break;
|
||||
case "mistral":
|
||||
initMistralAI();
|
||||
break;
|
||||
case "azure-openai":
|
||||
initAzureOpenAI();
|
||||
break;
|
||||
default:
|
||||
initOpenAI();
|
||||
break;
|
||||
}
|
||||
Settings.chunkSize = CHUNK_SIZE;
|
||||
Settings.chunkOverlap = CHUNK_OVERLAP;
|
||||
|
||||
setupProvider();
|
||||
};
|
||||
|
||||
function initOpenAI() {
|
||||
Settings.llm = new OpenAI({
|
||||
model: process.env.MODEL ?? "gpt-4o-mini",
|
||||
maxTokens: process.env.LLM_MAX_TOKENS
|
||||
? Number(process.env.LLM_MAX_TOKENS)
|
||||
: undefined,
|
||||
});
|
||||
Settings.embedModel = new OpenAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL,
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: undefined,
|
||||
});
|
||||
}
|
||||
|
||||
function initAzureOpenAI() {
|
||||
// Map Azure OpenAI model names to OpenAI model names (only for TS)
|
||||
const AZURE_OPENAI_MODEL_MAP: Record<string, string> = {
|
||||
"gpt-35-turbo": "gpt-3.5-turbo",
|
||||
"gpt-35-turbo-16k": "gpt-3.5-turbo-16k",
|
||||
"gpt-4o": "gpt-4o",
|
||||
"gpt-4": "gpt-4",
|
||||
"gpt-4-32k": "gpt-4-32k",
|
||||
"gpt-4-turbo": "gpt-4-turbo",
|
||||
"gpt-4-turbo-2024-04-09": "gpt-4-turbo",
|
||||
"gpt-4-vision-preview": "gpt-4-vision-preview",
|
||||
"gpt-4-1106-preview": "gpt-4-1106-preview",
|
||||
"gpt-4o-2024-05-13": "gpt-4o-2024-05-13",
|
||||
};
|
||||
|
||||
const azureConfig = {
|
||||
apiKey: process.env.AZURE_OPENAI_KEY,
|
||||
endpoint: process.env.AZURE_OPENAI_ENDPOINT,
|
||||
apiVersion:
|
||||
process.env.AZURE_OPENAI_API_VERSION || process.env.OPENAI_API_VERSION,
|
||||
};
|
||||
|
||||
Settings.llm = new OpenAI({
|
||||
model:
|
||||
AZURE_OPENAI_MODEL_MAP[process.env.MODEL ?? "gpt-35-turbo"] ??
|
||||
"gpt-3.5-turbo",
|
||||
maxTokens: process.env.LLM_MAX_TOKENS
|
||||
? Number(process.env.LLM_MAX_TOKENS)
|
||||
: undefined,
|
||||
azure: {
|
||||
...azureConfig,
|
||||
deployment: process.env.AZURE_OPENAI_LLM_DEPLOYMENT,
|
||||
},
|
||||
});
|
||||
|
||||
Settings.embedModel = new OpenAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL,
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: undefined,
|
||||
azure: {
|
||||
...azureConfig,
|
||||
deployment: process.env.AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function initOllama() {
|
||||
const config = {
|
||||
host: process.env.OLLAMA_BASE_URL ?? "http://127.0.0.1:11434",
|
||||
};
|
||||
Settings.llm = new Ollama({
|
||||
model: process.env.MODEL ?? "",
|
||||
config,
|
||||
});
|
||||
Settings.embedModel = new OllamaEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL ?? "",
|
||||
config,
|
||||
});
|
||||
}
|
||||
|
||||
function initGroq() {
|
||||
const embedModelMap: Record<string, string> = {
|
||||
"all-MiniLM-L6-v2": "Xenova/all-MiniLM-L6-v2",
|
||||
"all-mpnet-base-v2": "Xenova/all-mpnet-base-v2",
|
||||
};
|
||||
|
||||
Settings.llm = new Groq({
|
||||
model: process.env.MODEL!,
|
||||
});
|
||||
|
||||
Settings.embedModel = new HuggingFaceEmbedding({
|
||||
modelType: embedModelMap[process.env.EMBEDDING_MODEL!],
|
||||
});
|
||||
}
|
||||
|
||||
function initAnthropic() {
|
||||
const embedModelMap: Record<string, string> = {
|
||||
"all-MiniLM-L6-v2": "Xenova/all-MiniLM-L6-v2",
|
||||
"all-mpnet-base-v2": "Xenova/all-mpnet-base-v2",
|
||||
};
|
||||
Settings.llm = new Anthropic({
|
||||
model: process.env.MODEL as keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS,
|
||||
});
|
||||
Settings.embedModel = new HuggingFaceEmbedding({
|
||||
modelType: embedModelMap[process.env.EMBEDDING_MODEL!],
|
||||
});
|
||||
}
|
||||
|
||||
function initGemini() {
|
||||
Settings.llm = new Gemini({
|
||||
model: process.env.MODEL as GEMINI_MODEL,
|
||||
});
|
||||
Settings.embedModel = new GeminiEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL as GEMINI_EMBEDDING_MODEL,
|
||||
});
|
||||
}
|
||||
|
||||
function initMistralAI() {
|
||||
Settings.llm = new MistralAI({
|
||||
model: process.env.MODEL as keyof typeof ALL_AVAILABLE_MISTRAL_MODELS,
|
||||
});
|
||||
Settings.embedModel = new MistralAIEmbedding({
|
||||
model: process.env.EMBEDDING_MODEL as MistralAIEmbeddingModelType,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# flake8: noqa: E402
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
import logging
|
||||
|
||||
from app.index import get_index
|
||||
from app.settings import init_settings
|
||||
from llama_index.server.services.llamacloud.generate import (
|
||||
load_to_llamacloud,
|
||||
)
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger()
|
||||
|
||||
|
||||
def generate_datasource():
|
||||
init_settings()
|
||||
logger.info("Generate index for the provided data")
|
||||
|
||||
index = get_index(create_if_missing=True)
|
||||
if index is None:
|
||||
raise ValueError("Index not found and could not be created")
|
||||
|
||||
load_to_llamacloud(index, logger=logger)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
generate_datasource()
|
||||
@@ -0,0 +1,7 @@
|
||||
from llama_index.server.services.llamacloud import (
|
||||
LlamaCloudIndex,
|
||||
get_client,
|
||||
get_index,
|
||||
)
|
||||
|
||||
__all__ = ["LlamaCloudIndex", "get_client", "get_index"]
|
||||
@@ -0,0 +1,100 @@
|
||||
import * as dotenv from "dotenv";
|
||||
import "dotenv/config";
|
||||
import * as fs from "fs/promises";
|
||||
import { LLamaCloudFileService } from "llamaindex";
|
||||
import * as path from "path";
|
||||
import { getIndex } from "./app/data";
|
||||
import { initSettings } from "./app/settings";
|
||||
|
||||
dotenv.config();
|
||||
|
||||
const REQUIRED_ENV_VARS = [
|
||||
"LLAMA_CLOUD_INDEX_NAME",
|
||||
"LLAMA_CLOUD_PROJECT_NAME",
|
||||
"LLAMA_CLOUD_API_KEY",
|
||||
];
|
||||
|
||||
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(", ")}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function loadAndIndex() {
|
||||
const index = await getIndex();
|
||||
// ensure the index is available or create a new one
|
||||
await index.ensureIndex({
|
||||
verbose: true,
|
||||
embedding: {
|
||||
type: "OPENAI_EMBEDDING",
|
||||
component: {
|
||||
api_key: process.env.OPENAI_API_KEY,
|
||||
model_name: "text-embedding-3-small",
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
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")) {
|
||||
const buffer = await fs.readFile(filePath);
|
||||
const filename = path.basename(filePath);
|
||||
try {
|
||||
await LLamaCloudFileService.addFileToPipeline(
|
||||
projectId,
|
||||
pipelineId,
|
||||
new File([buffer], filename),
|
||||
);
|
||||
} catch (error) {
|
||||
if (
|
||||
error instanceof ReferenceError &&
|
||||
error.message.includes("File is not defined")
|
||||
) {
|
||||
throw new Error(
|
||||
"File class is not supported in the current Node.js version. Please use Node.js 20 or higher.",
|
||||
);
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`Successfully uploaded documents to LlamaCloud!`);
|
||||
}
|
||||
|
||||
(async () => {
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
@@ -0,0 +1,28 @@
|
||||
import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
|
||||
|
||||
type LlamaCloudDataSourceParams = {
|
||||
llamaCloudPipeline?: {
|
||||
project: string;
|
||||
pipeline: string;
|
||||
};
|
||||
};
|
||||
|
||||
export async function getIndex(params?: LlamaCloudDataSourceParams) {
|
||||
const { project, pipeline } = params?.llamaCloudPipeline ?? {};
|
||||
const projectName = project ?? process.env.LLAMA_CLOUD_PROJECT_NAME;
|
||||
const pipelineName = pipeline ?? process.env.LLAMA_CLOUD_INDEX_NAME;
|
||||
const apiKey = process.env.LLAMA_CLOUD_API_KEY;
|
||||
if (!projectName || !pipelineName || !apiKey) {
|
||||
throw new Error(
|
||||
"LlamaCloud is not configured. Please set project, pipeline, and api key in the params or as environment variables.",
|
||||
);
|
||||
}
|
||||
const index = new LlamaCloudIndex({
|
||||
organizationId: process.env.LLAMA_CLOUD_ORGANIZATION_ID,
|
||||
name: pipelineName,
|
||||
projectName,
|
||||
apiKey,
|
||||
baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
|
||||
});
|
||||
return index;
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { AstraDBVectorStore } from "@llamaindex/astra";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { AstraDBVectorStore } from "llamaindex/vector-store/AstraDBVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
@@ -37,8 +37,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { AstraDBVectorStore } from "@llamaindex/astra";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { AstraDBVectorStore } from "llamaindex/vector-store/AstraDBVectorStore";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { ChromaVectorStore } from "@llamaindex/chroma";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { ChromaVectorStore } from "llamaindex/vector-store/ChromaVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
@@ -30,8 +30,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { ChromaVectorStore } from "@llamaindex/chroma";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { ChromaVectorStore } from "llamaindex/vector-store/ChromaVectorStore";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
|
||||
@@ -57,8 +57,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { MilvusVectorStore } from "@llamaindex/milvus";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { MilvusVectorStore } from "llamaindex/vector-store/MilvusVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars, getMilvusClient } from "./shared";
|
||||
@@ -29,8 +29,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { MilvusVectorStore } from "@llamaindex/milvus";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { MilvusVectorStore } from "llamaindex/vector-store/MilvusVectorStore";
|
||||
import { checkRequiredEnvVars, getMilvusClient } from "./shared";
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { MongoDBAtlasVectorSearch } from "@llamaindex/mongodb";
|
||||
import * as dotenv from "dotenv";
|
||||
import { storageContextFromDefaults, VectorStoreIndex } from "llamaindex";
|
||||
import { MongoDBAtlasVectorSearch } from "llamaindex/vector-store/MongoDBAtlasVectorStore";
|
||||
import { MongoClient } from "mongodb";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
@@ -45,8 +45,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { MongoDBAtlasVectorSearch } from "@llamaindex/mongodb";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { MongoDBAtlasVectorSearch } from "llamaindex/vector-store/MongoDBAtlasVectorStore";
|
||||
import { MongoClient } from "mongodb";
|
||||
import { checkRequiredEnvVars, POPULATED_METADATA_FIELDS } from "./shared";
|
||||
|
||||
|
||||
@@ -37,7 +37,11 @@ async function generateDatasource() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
initSettings();
|
||||
await generateDatasource();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
initSettings();
|
||||
await generateDatasource();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { PGVectorStore } from "@llamaindex/postgres";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { PGVectorStore } from "llamaindex/vector-store/PGVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import {
|
||||
@@ -35,9 +35,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
process.exit(0);
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { PGVectorStore } from "@llamaindex/postgres";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { PGVectorStore } from "llamaindex/vector-store/PGVectorStore";
|
||||
import {
|
||||
PGVECTOR_SCHEMA,
|
||||
PGVECTOR_TABLE,
|
||||
@@ -14,6 +14,9 @@ export async function getDataSource(params?: any) {
|
||||
},
|
||||
schemaName: PGVECTOR_SCHEMA,
|
||||
tableName: PGVECTOR_TABLE,
|
||||
dimensions: process.env.EMBEDDING_DIM
|
||||
? parseInt(process.env.EMBEDDING_DIM)
|
||||
: undefined,
|
||||
});
|
||||
return await VectorStoreIndex.fromVectorStore(pgvs);
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { PineconeVectorStore } from "@llamaindex/pinecone";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { PineconeVectorStore } from "llamaindex/vector-store/PineconeVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
@@ -25,8 +25,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { PineconeVectorStore } from "@llamaindex/pinecone";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { PineconeVectorStore } from "llamaindex/vector-store/PineconeVectorStore";
|
||||
import { checkRequiredEnvVars } from "./shared";
|
||||
|
||||
export async function getDataSource(params?: any) {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
/* eslint-disable turbo/no-undeclared-env-vars */
|
||||
import { QdrantVectorStore } from "@llamaindex/qdrant";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex, storageContextFromDefaults } from "llamaindex";
|
||||
import { QdrantVectorStore } from "llamaindex/vector-store/QdrantVectorStore";
|
||||
import { getDocuments } from "./loader";
|
||||
import { initSettings } from "./settings";
|
||||
import { checkRequiredEnvVars, getQdrantClient } from "./shared";
|
||||
@@ -30,8 +30,12 @@ async function loadAndIndex() {
|
||||
}
|
||||
|
||||
(async () => {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
try {
|
||||
checkRequiredEnvVars();
|
||||
initSettings();
|
||||
await loadAndIndex();
|
||||
console.log("Finished generating storage.");
|
||||
} catch (error) {
|
||||
console.error("Error generating storage.", error);
|
||||
}
|
||||
})();
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { QdrantVectorStore } from "@llamaindex/qdrant";
|
||||
import * as dotenv from "dotenv";
|
||||
import { VectorStoreIndex } from "llamaindex";
|
||||
import { QdrantVectorStore } from "llamaindex/vector-store/QdrantVectorStore";
|
||||
import { checkRequiredEnvVars, getQdrantClient } from "./shared";
|
||||
|
||||
dotenv.config();
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user