Compare commits

..

21 Commits

Author SHA1 Message Date
Thuc Pham 3d3d94ca33 feat: limit display 2024-05-30 10:03:32 +07:00
Thuc Pham b0e25b3d7d fix lint 2024-05-29 08:03:11 +00:00
Thuc Pham ce2de2d0a8 update convert message content 2024-05-29 07:58:15 +00:00
Thuc Pham b3f3ce403f resolve conflict 2024-05-29 08:41:05 +07:00
Thuc Pham d473867a9a get resouces inside chat input 2024-05-29 08:26:17 +07:00
Thuc Pham 5068e28c8e Merge branch 'main' into feat/implement-csv-upload 2024-05-29 08:10:11 +07:00
Thuc Pham 04cc7ce451 suggest csv resources from frontend annotation data 2024-05-27 17:45:16 +07:00
Thuc Pham bafb1d7d71 remove multi modal prop 2024-05-27 11:12:20 +07:00
Marcus Schiesser 2eb2803284 Create bright-turkeys-melt.md 2024-05-24 19:03:59 +07:00
Thuc Pham beaf4b4e90 add timeout for streaming 2024-05-24 10:54:52 +07:00
Thuc Pham c24e60a427 Update pre-commit 2024-05-24 09:15:28 +07:00
Thuc Pham 14c0b2ea87 Update .gitignore 2024-05-24 09:15:13 +07:00
Thuc Pham f3c1e54e97 fix lint 2024-05-24 09:12:55 +07:00
Thuc Pham 667c122eed update icon & fix code review 2024-05-24 09:09:00 +07:00
Thuc Pham cc45225a4a fix lint 2024-05-24 08:51:23 +07:00
Thuc Pham 5d190195b6 remove dead code 2024-05-24 08:49:16 +07:00
Thuc Pham 2e47af86ce feat: implement csv upload 2024-05-24 08:43:51 +07:00
Thuc Pham 778bdfaa0b fix: typo 2024-05-24 08:43:51 +07:00
Thuc Pham 2e43ba1615 refactor: use local file system, use absolute resource url 2024-05-24 08:42:46 +07:00
Thuc Pham ff8bed62c7 build tool system prompt 2024-05-24 08:39:38 +07:00
Thuc Pham a9d2177dc3 feat: implement interpreter tool 2024-05-24 08:39:00 +07:00
80 changed files with 830 additions and 2280 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"create-llama": patch
---
Add CSV upload
+5
View File
@@ -0,0 +1,5 @@
---
"create-llama": patch
---
Add support E2B code interpreter tool for FastAPI
+1 -1
View File
@@ -17,7 +17,7 @@ jobs:
matrix:
node-version: [18, 20]
python-version: ["3.11"]
os: [macos-latest, windows-latest, ubuntu-22.04]
os: [macos-latest, windows-latest]
defaults:
run:
shell: bash
-52
View File
@@ -1,57 +1,5 @@
# create-llama
## 0.1.15
### Patch Changes
- 64732f0: Fix the issue of images not showing with the sandbox URL from OpenAI's models
- aeb6fef: use llamacloud for chat
## 0.1.14
### Patch Changes
- f2c3389: chore: update to llamaindex 0.4.3
- 5093b37: Remove non-working file selectors for Linux
## 0.1.13
### Patch Changes
- b3c969d: Add image generator tool
## 0.1.12
### Patch Changes
- aa69014: Fix NextJS for TS 5.2
## 0.1.11
### Patch Changes
- 48b96ff: Add DuckDuckGo search tool
- 9c9decb: Reuse function tool instances and improve e2b interpreter tool for Python
- 02ed277: Add Groq as a model provider
- 0748f2e: Remove hard-coded Gemini supported models
## 0.1.10
### Patch Changes
- 9112d08: Add OpenAPI tool for Typescript
- 8f03f8d: Add OLLAMA_REQUEST_TIMEOUT variable to config Ollama timeout (Python)
- 8f03f8d: Apply nest_asyncio for llama parse
## 0.1.9
### Patch Changes
- a42fa53: Add CSV upload
- 563b51d: Fix Vercel streaming (python) to stream data events instantly
- d60b3c5: Add E2B code interpreter tool for FastAPI
- 956538e: Add OpenAPI action tool for FastAPI
## 0.1.8
### Patch Changes
-14
View File
@@ -151,19 +151,5 @@ export async function createApp({
);
}
if (
dataSources.some((dataSource) => dataSource.type === "file") &&
process.platform === "linux"
) {
console.log(
yellow(
`You can add your own data files to ${terminalLink(
"data",
`file://${root}/data`,
)} folder manually.`,
),
);
}
console.log();
}
-33
View File
@@ -133,26 +133,6 @@ const getVectorDBEnvs = (
"Optional API key for authenticating requests to Qdrant.",
},
];
case "llamacloud":
return [
{
name: "LLAMA_CLOUD_INDEX_NAME",
description:
"The name of the LlamaCloud index to use (part of the LlamaCloud project).",
value: "test",
},
{
name: "LLAMA_CLOUD_PROJECT_NAME",
description: "The name of the LlamaCloud project.",
value: "Default",
},
{
name: "LLAMA_CLOUD_BASE_URL",
description:
"The base URL for the LlamaCloud API. Only change this for non-production environments",
value: "https://api.cloud.llamaindex.ai",
},
];
case "chroma":
const envs = [
{
@@ -205,10 +185,6 @@ const getModelEnvs = (modelConfig: ModelConfig): EnvVar[] => {
description: "Dimension of the embedding model to use.",
value: modelConfig.dimensions.toString(),
},
{
name: "CONVERSATION_STARTERS",
description: "The questions to help users get started (multi-line).",
},
...(modelConfig.provider === "openai"
? [
{
@@ -235,15 +211,6 @@ const getModelEnvs = (modelConfig: ModelConfig): EnvVar[] => {
},
]
: []),
...(modelConfig.provider === "groq"
? [
{
name: "GROQ_API_KEY",
description: "The Groq API key to use.",
value: modelConfig.apiKey,
},
]
: []),
...(modelConfig.provider === "gemini"
? [
{
-99
View File
@@ -1,99 +0,0 @@
import ciInfo from "ci-info";
import prompts from "prompts";
import { ModelConfigParams } from ".";
import { questionHandlers, toChoice } from "../../questions";
const MODELS = ["llama3-8b", "llama3-70b", "mixtral-8x7b"];
const DEFAULT_MODEL = MODELS[0];
// Use huggingface embedding models for now as Groq doesn't support embedding models
enum HuggingFaceEmbeddingModelType {
XENOVA_ALL_MINILM_L6_V2 = "all-MiniLM-L6-v2",
XENOVA_ALL_MPNET_BASE_V2 = "all-mpnet-base-v2",
}
type ModelData = {
dimensions: number;
};
const EMBEDDING_MODELS: Record<HuggingFaceEmbeddingModelType, ModelData> = {
[HuggingFaceEmbeddingModelType.XENOVA_ALL_MINILM_L6_V2]: {
dimensions: 384,
},
[HuggingFaceEmbeddingModelType.XENOVA_ALL_MPNET_BASE_V2]: {
dimensions: 768,
},
};
const DEFAULT_EMBEDDING_MODEL = Object.keys(EMBEDDING_MODELS)[0];
const DEFAULT_DIMENSIONS = Object.values(EMBEDDING_MODELS)[0].dimensions;
type GroqQuestionsParams = {
apiKey?: string;
askModels: boolean;
};
export async function askGroqQuestions({
askModels,
apiKey,
}: GroqQuestionsParams): Promise<ModelConfigParams> {
const config: ModelConfigParams = {
apiKey,
model: DEFAULT_MODEL,
embeddingModel: DEFAULT_EMBEDDING_MODEL,
dimensions: DEFAULT_DIMENSIONS,
isConfigured(): boolean {
if (config.apiKey) {
return true;
}
if (process.env["GROQ_API_KEY"]) {
return true;
}
return false;
},
};
if (!config.apiKey) {
const { key } = await prompts(
{
type: "text",
name: "key",
message:
"Please provide your Groq API key (or leave blank to use GROQ_API_KEY env variable):",
},
questionHandlers,
);
config.apiKey = key || process.env.GROQ_API_KEY;
}
// use default model values in CI or if user should not be asked
const useDefaults = ciInfo.isCI || !askModels;
if (!useDefaults) {
const { model } = await prompts(
{
type: "select",
name: "model",
message: "Which LLM model would you like to use?",
choices: MODELS.map(toChoice),
initial: 0,
},
questionHandlers,
);
config.model = model;
const { embeddingModel } = await prompts(
{
type: "select",
name: "embeddingModel",
message: "Which embedding model would you like to use?",
choices: Object.keys(EMBEDDING_MODELS).map(toChoice),
initial: 0,
},
questionHandlers,
);
config.embeddingModel = embeddingModel;
config.dimensions =
EMBEDDING_MODELS[
embeddingModel as HuggingFaceEmbeddingModelType
].dimensions;
}
return config;
}
-5
View File
@@ -4,7 +4,6 @@ import { questionHandlers } from "../../questions";
import { ModelConfig, ModelProvider } from "../types";
import { askAnthropicQuestions } from "./anthropic";
import { askGeminiQuestions } from "./gemini";
import { askGroqQuestions } from "./groq";
import { askOllamaQuestions } from "./ollama";
import { askOpenAIQuestions } from "./openai";
@@ -33,7 +32,6 @@ export async function askModelConfig({
title: "OpenAI",
value: "openai",
},
{ title: "Groq", value: "groq" },
{ title: "Ollama", value: "ollama" },
{ title: "Anthropic", value: "anthropic" },
{ title: "Gemini", value: "gemini" },
@@ -50,9 +48,6 @@ export async function askModelConfig({
case "ollama":
modelConfig = await askOllamaQuestions({ askModels });
break;
case "groq":
modelConfig = await askGroqQuestions({ askModels });
break;
case "anthropic":
modelConfig = await askAnthropicQuestions({ askModels });
break;
+2 -8
View File
@@ -118,12 +118,6 @@ const getAdditionalDependencies = (
version: "^2.9.9",
});
break;
case "llamacloud":
dependencies.push({
name: "llama-index-indices-managed-llama-cloud",
version: "^0.2.1",
});
break;
}
}
}
@@ -150,7 +144,7 @@ const getAdditionalDependencies = (
case "openai":
dependencies.push({
name: "llama-index-agent-openai",
version: "0.2.6",
version: "0.2.2",
});
break;
case "anthropic":
@@ -166,7 +160,7 @@ const getAdditionalDependencies = (
case "gemini":
dependencies.push({
name: "llama-index-llms-gemini",
version: "0.1.10",
version: "0.1.7",
});
dependencies.push({
name: "llama-index-embeddings-gemini",
+9 -86
View File
@@ -30,7 +30,7 @@ export type ToolDependencies = {
export const supportedTools: Tool[] = [
{
display: "Google Search",
display: "Google Search (configuration required after installation)",
name: "google.GoogleSearchToolSpec",
config: {
engine:
@@ -54,29 +54,6 @@ export const supportedTools: Tool[] = [
},
],
},
{
// For python app, we will use a local DuckDuckGo search tool (instead of DuckDuckGo search tool in LlamaHub)
// to get the same results as the TS app.
display: "DuckDuckGo Search",
name: "duckduckgo",
dependencies: [
{
name: "duckduckgo-search",
version: "6.1.7",
},
],
supportedFrameworks: ["fastapi", "nextjs", "express"],
type: ToolType.LOCAL,
envVars: [
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for DuckDuckGo search tool.",
value: `You are a DuckDuckGo search agent.
You can use the duckduckgo search tool to get information from the web to answer user questions.
For better results, you can specify the region parameter to get results from a specific region but it's optional.`,
},
],
},
{
display: "Wikipedia",
name: "wikipedia.WikipediaToolSpec",
@@ -130,61 +107,13 @@ For better results, you can specify the region parameter to get results from a s
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for code interpreter tool.",
value: `-You are a Python interpreter that can run any python code in a secure environment.
- The python code runs in a Jupyter notebook. Every time you call the 'interpreter' tool, the python code is executed in a separate cell.
- You are given tasks to complete and you run python code to solve them.
- It's okay to make multiple calls to interpreter tool. If you get an error or the result is not what you expected, you can call the tool again. Don't give up too soon!
- Plot visualizations using matplotlib or any other visualization library directly in the notebook.
- You can install any pip package (if it exists) by running a cell with pip install.`,
},
],
},
{
display: "OpenAPI action",
name: "openapi_action.OpenAPIActionToolSpec",
dependencies: [
{
name: "llama-index-tools-openapi",
version: "0.1.3",
},
{
name: "jsonschema",
version: "^4.22.0",
},
{
name: "llama-index-tools-requests",
version: "0.1.3",
},
],
config: {
openapi_uri: "The URL or file path of the OpenAPI schema",
},
supportedFrameworks: ["fastapi", "express", "nextjs"],
type: ToolType.LOCAL,
envVars: [
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for openapi action tool.",
value:
"You are an OpenAPI action agent. You help users to make requests to the provided OpenAPI schema.",
},
],
},
{
display: "Image Generator",
name: "img_gen",
supportedFrameworks: ["fastapi", "express", "nextjs"],
type: ToolType.LOCAL,
envVars: [
{
name: "STABILITY_API_KEY",
description:
"STABILITY_API_KEY key is required to run image generator. Get it here: https://platform.stability.ai/account/keys",
},
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for image generator tool.",
value: `You are an image generator agent. You help users to generate images using the Stability API.`,
value: `You are a Python interpreter.
- You are given tasks to complete and you run python code to solve them.
- The python code runs in a Jupyter notebook. Every time you call \`interpreter\` tool, the python code is executed in a separate cell. It's okay to make multiple calls to \`interpreter\`.
- Display visualizations using matplotlib or any other visualization library directly in the notebook. Shouldn't save the visualizations to a file, just return the base64 encoded data.
- You can install any pip package (if it exists) if you need to but the usual packages for data analysis are already preinstalled.
- You can run any python code you want in a secure environment.
- Use absolute url from result to display images or any other media.`,
},
],
},
@@ -213,15 +142,9 @@ export const getTools = (toolsName: string[]): Tool[] => {
return tools;
};
export const toolRequiresConfig = (tool: Tool): boolean => {
const hasConfig = Object.keys(tool.config || {}).length > 0;
const hasEmptyEnvVar = tool.envVars?.some((envVar) => !envVar.value) ?? false;
return hasConfig || hasEmptyEnvVar;
};
export const toolsRequireConfig = (tools?: Tool[]): boolean => {
if (tools) {
return tools?.some(toolRequiresConfig);
return tools?.some((tool) => Object.keys(tool.config || {}).length > 0);
}
return false;
};
+3 -9
View File
@@ -1,12 +1,7 @@
import { PackageManager } from "../helpers/get-pkg-manager";
import { Tool } from "./tools";
export type ModelProvider =
| "openai"
| "groq"
| "ollama"
| "anthropic"
| "gemini";
export type ModelProvider = "openai" | "ollama" | "anthropic" | "gemini";
export type ModelConfig = {
provider: ModelProvider;
apiKey?: string;
@@ -26,8 +21,7 @@ export type TemplateVectorDB =
| "milvus"
| "astra"
| "qdrant"
| "chroma"
| "llamacloud";
| "chroma";
export type TemplatePostInstallAction =
| "none"
| "VSCode"
@@ -37,7 +31,7 @@ export type TemplateDataSource = {
type: TemplateDataSourceType;
config: TemplateDataSourceConfig;
};
export type TemplateDataSourceType = "file" | "web" | "db" | "llamacloud";
export type TemplateDataSourceType = "file" | "web" | "db";
export type TemplateObservability = "none" | "opentelemetry";
// Config for both file and folder
export type FileSourceConfig = {
+2 -8
View File
@@ -1,7 +1,7 @@
import fs from "fs/promises";
import os from "os";
import path from "path";
import { bold, cyan, yellow } from "picocolors";
import { bold, cyan } from "picocolors";
import { assetRelocator, copy } from "../helpers/copy";
import { callPackageManager } from "../helpers/install";
import { templatesDir } from "./dir";
@@ -105,13 +105,7 @@ export const installTSTemplate = async ({
const enginePath = path.join(root, relativeEngineDestPath, "engine");
// copy vector db component
if (vectorDb === "llamacloud") {
console.log(
`\nUsing managed index from LlamaCloud. Ensure the ${yellow("LLAMA_CLOUD_* environment variables are set correctly.")}`,
);
} else {
console.log("\nUsing vector DB:", vectorDb ?? "none");
}
console.log("\nUsing vector DB:", vectorDb ?? "none", "\n");
await copy("**", enginePath, {
parents: true,
cwd: path.join(compPath, "vectordbs", "typescript", vectorDb ?? "none"),
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "create-llama",
"version": "0.1.15",
"version": "0.1.8",
"description": "Create LlamaIndex-powered apps with one command",
"keywords": [
"rag",
+50 -100
View File
@@ -16,11 +16,7 @@ import { templatesDir } from "./helpers/dir";
import { getAvailableLlamapackOptions } from "./helpers/llama-pack";
import { askModelConfig } from "./helpers/providers";
import { getProjectOptions } from "./helpers/repo";
import {
supportedTools,
toolRequiresConfig,
toolsRequireConfig,
} from "./helpers/tools";
import { supportedTools, toolsRequireConfig } from "./helpers/tools";
export type QuestionArgs = Omit<
InstallAppArgs,
@@ -123,13 +119,7 @@ export const getDataSourceChoices = (
framework: TemplateFramework,
selectedDataSource: TemplateDataSource[],
) => {
// If LlamaCloud is already selected, don't show any other options
if (selectedDataSource.find((s) => s.type === "llamacloud")) {
return [];
}
const choices = [];
if (selectedDataSource.length > 0) {
choices.push({
title: "No",
@@ -142,30 +132,24 @@ export const getDataSourceChoices = (
value: "none",
});
choices.push({
title:
process.platform !== "linux"
? "Use an example PDF"
: "Use an example PDF (you can add your own data files later)",
title: "Use an example PDF",
value: "exampleFile",
});
}
// Linux has many distros so we won't support file/folder picker for now
if (process.platform !== "linux") {
choices.push(
{
title: `Use local files (${supportedContextFileTypes.join(", ")})`,
value: "file",
},
{
title:
process.platform === "win32"
? "Use a local folder"
: "Use local folders",
value: "folder",
},
);
}
choices.push(
{
title: `Use local files (${supportedContextFileTypes.join(", ")})`,
value: "file",
},
{
title:
process.platform === "win32"
? "Use a local folder"
: "Use local folders",
value: "folder",
},
);
if (framework === "fastapi") {
choices.push({
@@ -177,13 +161,6 @@ export const getDataSourceChoices = (
value: "db",
});
}
if (!selectedDataSource.length) {
choices.push({
title: "Use managed index from LlamaCloud",
value: "llamacloud",
});
}
return choices;
};
@@ -497,11 +474,6 @@ export const askQuestions = async (
// continue asking user for data sources if none are initially provided
while (true) {
const firstQuestion = program.dataSources.length === 0;
const choices = getDataSourceChoices(
program.framework,
program.dataSources,
);
if (choices.length === 0) break;
const { selectedSource } = await prompts(
{
type: "select",
@@ -509,7 +481,10 @@ export const askQuestions = async (
message: firstQuestion
? "Which data source would you like to use?"
: "Would you like to add another data source?",
choices,
choices: getDataSourceChoices(
program.framework,
program.dataSources,
),
initial: firstQuestion ? 1 : 0,
},
questionHandlers,
@@ -606,76 +581,51 @@ export const askQuestions = async (
config: await prompts(dbPrompts, questionHandlers),
});
}
case "llamacloud": {
program.dataSources.push({
type: "llamacloud",
config: {},
});
program.dataSources.push(EXAMPLE_FILE);
break;
}
}
}
}
}
const isUsingLlamaCloud = program.dataSources.some(
(ds) => ds.type === "llamacloud",
);
// Asking for LlamaParse if user selected file or folder data source
if (
program.dataSources.some((ds) => ds.type === "file") &&
program.useLlamaParse === undefined
) {
if (ciInfo.isCI) {
program.useLlamaParse = getPrefOrDefault("useLlamaParse");
program.llamaCloudKey = getPrefOrDefault("llamaCloudKey");
} else {
const { useLlamaParse } = await prompts(
{
type: "toggle",
name: "useLlamaParse",
message:
"Would you like to use LlamaParse (improved parser for RAG - requires API key)?",
initial: false,
active: "yes",
inactive: "no",
},
questionHandlers,
);
program.useLlamaParse = useLlamaParse;
// Asking for LlamaParse if user selected file data source
if (isUsingLlamaCloud) {
// default to use LlamaParse if using LlamaCloud
program.useLlamaParse = preferences.useLlamaParse = true;
} else {
if (program.dataSources.some((ds) => ds.type === "file")) {
if (ciInfo.isCI) {
program.useLlamaParse = getPrefOrDefault("useLlamaParse");
} else {
const { useLlamaParse } = await prompts(
// Ask for LlamaCloud API key
if (useLlamaParse && program.llamaCloudKey === undefined) {
const { llamaCloudKey } = await prompts(
{
type: "toggle",
name: "useLlamaParse",
type: "text",
name: "llamaCloudKey",
message:
"Would you like to use LlamaParse (improved parser for RAG - requires API key)?",
initial: false,
active: "yes",
inactive: "no",
"Please provide your LlamaIndex Cloud API key (leave blank to skip):",
},
questionHandlers,
);
program.useLlamaParse = useLlamaParse;
preferences.useLlamaParse = useLlamaParse;
program.llamaCloudKey = llamaCloudKey;
}
}
}
// Ask for LlamaCloud API key when using a LlamaCloud index or LlamaParse
if (isUsingLlamaCloud || program.useLlamaParse) {
if (ciInfo.isCI) {
program.llamaCloudKey = getPrefOrDefault("llamaCloudKey");
} else {
// Ask for LlamaCloud API key
const { llamaCloudKey } = await prompts(
{
type: "text",
name: "llamaCloudKey",
message:
"Please provide your LlamaCloud API key (leave blank to skip):",
},
questionHandlers,
);
program.llamaCloudKey = preferences.llamaCloudKey =
llamaCloudKey || process.env.LLAMA_CLOUD_API_KEY;
}
}
if (isUsingLlamaCloud) {
// When using a LlamaCloud index, don't ask for vector database and use code in `llamacloud` folder for vector database
const vectorDb = "llamacloud";
program.vectorDb = vectorDb;
preferences.vectorDb = vectorDb;
} else if (program.dataSources.length > 0 && !program.vectorDb) {
if (program.dataSources.length > 0 && !program.vectorDb) {
if (ciInfo.isCI) {
program.vectorDb = getPrefOrDefault("vectorDb");
} else {
@@ -702,7 +652,7 @@ export const askQuestions = async (
t.supportedFrameworks?.includes(program.framework),
);
const toolChoices = options.map((tool) => ({
title: `${tool.display}${toolRequiresConfig(tool) ? " (needs configuration)" : ""}`,
title: tool.display,
value: tool.name,
}));
const { toolsName } = await prompts({
@@ -1,8 +1,7 @@
import os
import yaml
import json
import importlib
from cachetools import cached, LRUCache
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from llama_index.core.tools.function_tool import FunctionTool
@@ -19,6 +18,7 @@ class ToolFactory:
ToolType.LOCAL: "app.engine.tools",
}
@staticmethod
def load_tools(tool_type: str, tool_name: str, config: dict) -> list[FunctionTool]:
source_package = ToolFactory.TOOL_SOURCE_PACKAGE_MAP[tool_type]
try:
@@ -31,7 +31,7 @@ class ToolFactory:
return tool_spec.to_tool_list()
else:
module = importlib.import_module(f"{source_package}.{tool_name}")
tools = module.get_tools(**config)
tools = getattr(module, "tools")
if not all(isinstance(tool, FunctionTool) for tool in tools):
raise ValueError(
f"The module {module} does not contain valid tools"
@@ -1,36 +0,0 @@
from llama_index.core.tools.function_tool import FunctionTool
def duckduckgo_search(
query: str,
region: str = "wt-wt",
max_results: int = 10,
):
"""
Use this function to search for any query in DuckDuckGo.
Args:
query (str): The query to search in DuckDuckGo.
region Optional(str): The region to be used for the search in [country-language] convention, ex us-en, uk-en, ru-ru, etc...
max_results Optional(int): The maximum number of results to be returned. Default is 10.
"""
try:
from duckduckgo_search import DDGS
except ImportError:
raise ImportError(
"duckduckgo_search package is required to use this function."
"Please install it by running: `poetry add duckduckgo_search` or `pip install duckduckgo_search`"
)
params = {
"keywords": query,
"region": region,
"max_results": max_results,
}
results = []
with DDGS() as ddg:
results = list(ddg.text(**params))
return results
def get_tools(**kwargs):
return [FunctionTool.from_defaults(duckduckgo_search)]
@@ -1,108 +0,0 @@
import os
import uuid
import logging
import requests
from typing import Optional
from pydantic import BaseModel, Field
from llama_index.core.tools import FunctionTool
logger = logging.getLogger(__name__)
class ImageGeneratorToolOutput(BaseModel):
is_success: bool = Field(
...,
description="Whether the image generation was successful.",
)
image_url: Optional[str] = Field(
None,
description="The URL of the generated image.",
)
error_message: Optional[str] = Field(
None,
description="The error message if the image generation failed.",
)
class ImageGeneratorTool:
_IMG_OUTPUT_FORMAT = "webp"
_IMG_OUTPUT_DIR = "tool-output"
_IMG_GEN_API = "https://api.stability.ai/v2beta/stable-image/generate/core"
def __init__(self, api_key: str = None):
if not api_key:
api_key = os.getenv("STABILITY_API_KEY")
self._api_key = api_key
self.fileserver_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
if self._api_key is None:
raise ValueError(
"STABILITY_API_KEY key is required to run image generator. Get it here: https://platform.stability.ai/account/keys"
)
if self.fileserver_url_prefix is None:
raise ValueError("FILESERVER_URL_PREFIX is required.")
def _prepare_output_dir(self):
"""
Create the output directory if it doesn't exist
"""
if not os.path.exists(self._IMG_OUTPUT_DIR):
os.makedirs(self._IMG_OUTPUT_DIR, exist_ok=True)
def _save_image(self, image_data: bytes):
self._prepare_output_dir()
filename = f"{uuid.uuid4()}.{self._IMG_OUTPUT_FORMAT}"
output_path = os.path.join(self._IMG_OUTPUT_DIR, filename)
with open(output_path, "wb") as f:
f.write(image_data)
url = f"{os.getenv('FILESERVER_URL_PREFIX')}/{self._IMG_OUTPUT_DIR}/{filename}"
logger.info(f"Saved image to {output_path}.\nURL: {url}")
return url
def _call_stability_api(self, prompt: str):
headers = {
"authorization": f"Bearer {self._api_key}",
"accept": "image/*",
}
data = {
"prompt": prompt,
"output_format": self._IMG_OUTPUT_FORMAT,
}
response = requests.post(
self._IMG_GEN_API,
headers=headers,
files={"none": ""},
data=data,
)
response.raise_for_status()
return response
def generate_image(self, prompt: str) -> ImageGeneratorToolOutput:
"""
Use this tool to generate an image based on the prompt.
Args:
prompt (str): The prompt to generate the image from.
"""
try:
# Call the Stability API
response = self._call_stability_api(prompt)
# Save the image and get the URL
image_url = self._save_image(response.content)
return ImageGeneratorToolOutput(
is_success=True,
image_url=image_url,
)
except Exception as e:
logger.exception(e, exc_info=True)
return ImageGeneratorToolOutput(
is_success=False,
error_message=str(e),
)
def get_tools(**kwargs):
return [FunctionTool.from_defaults(ImageGeneratorTool(**kwargs).generate_image)]
@@ -3,7 +3,7 @@ import logging
import base64
import uuid
from pydantic import BaseModel
from typing import List, Tuple, Dict, Optional
from typing import List, Tuple, Dict
from llama_index.core.tools import FunctionTool
from e2b_code_interpreter import CodeInterpreter
from e2b_code_interpreter.models import Logs
@@ -14,9 +14,8 @@ logger = logging.getLogger(__name__)
class InterpreterExtraResult(BaseModel):
type: str
content: Optional[str] = None
filename: Optional[str] = None
url: Optional[str] = None
filename: str
url: str
class E2BToolOutput(BaseModel):
@@ -29,24 +28,9 @@ class E2BCodeInterpreter:
output_dir = "tool-output"
def __init__(self, api_key: str = None):
if api_key is None:
api_key = os.getenv("E2B_API_KEY")
filesever_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
if not api_key:
raise ValueError(
"E2B_API_KEY key is required to run code interpreter. Get it here: https://e2b.dev/docs/getting-started/api-key"
)
if not filesever_url_prefix:
raise ValueError(
"FILESERVER_URL_PREFIX is required to display file output from sandbox"
)
def __init__(self, api_key: str, filesever_url_prefix: str):
self.api_key = api_key
self.filesever_url_prefix = filesever_url_prefix
self.interpreter = CodeInterpreter(api_key=api_key)
def __del__(self):
self.interpreter.close()
def get_output_path(self, filename: str) -> str:
# if output directory doesn't exist, create it
@@ -88,56 +72,63 @@ class E2BCodeInterpreter:
try:
formats = result.formats()
results = [result[format] for format in formats]
base64_data_arr = [result[format] for format in formats]
for ext, data in zip(formats, results):
match ext:
case "png" | "svg" | "jpeg" | "pdf":
result = self.save_to_disk(data, ext)
filename = result["filename"]
output.append(
InterpreterExtraResult(
type=ext,
filename=filename,
url=self.get_file_url(filename),
)
)
case _:
output.append(
InterpreterExtraResult(
type=ext,
content=data,
)
for ext, base64_data in zip(formats, base64_data_arr):
if ext and base64_data:
result = self.save_to_disk(base64_data, ext)
filename = result["filename"]
output.append(
InterpreterExtraResult(
type=ext, filename=filename, url=self.get_file_url(filename)
)
)
except Exception as error:
logger.exception(error, exc_info=True)
logger.error("Error when parsing output from E2b interpreter tool", error)
logger.error("Error when saving data to disk", error)
return output
def interpret(self, code: str) -> E2BToolOutput:
"""
Execute python code in a Jupyter notebook cell, the toll will return result, stdout, stderr, display_data, and error.
with CodeInterpreter(api_key=self.api_key) as interpreter:
logger.info(
f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}"
)
exec = interpreter.notebook.exec_cell(code)
Parameters:
code (str): The python code to be executed in a single cell.
"""
logger.info(
f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}"
)
exec = self.interpreter.notebook.exec_cell(code)
if exec.error:
logger.error("Error when executing code", exec.error)
output = E2BToolOutput(is_error=True, logs=exec.logs, results=[])
else:
if len(exec.results) == 0:
output = E2BToolOutput(is_error=False, logs=exec.logs, results=[])
if exec.error:
output = E2BToolOutput(is_error=True, logs=[exec.error])
else:
results = self.parse_result(exec.results[0])
output = E2BToolOutput(is_error=False, logs=exec.logs, results=results)
return output
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
)
return output
def get_tools(**kwargs):
return [FunctionTool.from_defaults(E2BCodeInterpreter(**kwargs).interpret)]
def code_interpret(code: str) -> Dict:
"""
Execute python code in a Jupyter notebook cell and return any result, stdout, stderr, display_data, and error.
"""
api_key = os.getenv("E2B_API_KEY")
filesever_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
if not api_key:
raise ValueError(
"E2B_API_KEY key is required to run code interpreter. Get it here: https://e2b.dev/docs/getting-started/api-key"
)
if not filesever_url_prefix:
raise ValueError(
"FILESERVER_URL_PREFIX is required to display file output from sandbox"
)
interpreter = E2BCodeInterpreter(
api_key=api_key, filesever_url_prefix=filesever_url_prefix
)
output = interpreter.interpret(code)
return output.dict()
# Specify as functions tools to be loaded by the ToolFactory
tools = [FunctionTool.from_defaults(code_interpret)]
@@ -1,78 +0,0 @@
from typing import Dict, List, Tuple
from llama_index.tools.openapi import OpenAPIToolSpec
from llama_index.tools.requests import RequestsToolSpec
class OpenAPIActionToolSpec(OpenAPIToolSpec, RequestsToolSpec):
"""
A combination of OpenAPI and Requests tool specs that can parse OpenAPI specs and make requests.
openapi_uri: str: The file path or URL to the OpenAPI spec.
domain_headers: dict: Whitelist domains and the headers to use.
"""
spec_functions = OpenAPIToolSpec.spec_functions + RequestsToolSpec.spec_functions
# Cached parsed specs by URI
_specs: Dict[str, Tuple[Dict, List[str]]] = {}
def __init__(self, openapi_uri: str, domain_headers: dict = None, **kwargs):
if domain_headers is None:
domain_headers = {}
if openapi_uri not in self._specs:
openapi_spec, servers = self._load_openapi_spec(openapi_uri)
self._specs[openapi_uri] = (openapi_spec, servers)
else:
openapi_spec, servers = self._specs[openapi_uri]
# Add the servers to the domain headers if they are not already present
for server in servers:
if server not in domain_headers:
domain_headers[server] = {}
OpenAPIToolSpec.__init__(self, spec=openapi_spec)
RequestsToolSpec.__init__(self, domain_headers)
@staticmethod
def _load_openapi_spec(uri: str) -> Tuple[Dict, List[str]]:
"""
Load an OpenAPI spec from a URI.
Args:
uri (str): A file path or URL to the OpenAPI spec.
Returns:
List[Document]: A list of Document objects.
"""
import yaml
from urllib.parse import urlparse
if uri.startswith("http"):
import requests
response = requests.get(uri)
if response.status_code != 200:
raise ValueError(
"Could not initialize OpenAPIActionToolSpec: "
f"Failed to load OpenAPI spec from {uri}, status code: {response.status_code}"
)
spec = yaml.safe_load(response.text)
elif uri.startswith("file"):
filepath = urlparse(uri).path
with open(filepath, "r") as file:
spec = yaml.safe_load(file)
else:
raise ValueError(
"Could not initialize OpenAPIActionToolSpec: Invalid OpenAPI URI provided. "
"Only HTTP and file path are supported."
)
# Add the servers to the whitelist
try:
servers = [
urlparse(server["url"]).netloc for server in spec.get("servers", [])
]
except KeyError as e:
raise ValueError(
"Could not initialize OpenAPIActionToolSpec: Invalid OpenAPI spec provided. "
"Could not get `servers` from the spec."
) from e
return spec, servers
@@ -69,5 +69,4 @@ class OpenMeteoWeather:
return response.json()
def get_tools(**kwargs):
return [FunctionTool.from_defaults(OpenMeteoWeather.get_weather_information)]
tools = [FunctionTool.from_defaults(OpenMeteoWeather.get_weather_information)]
@@ -2,6 +2,7 @@ import { BaseToolWithCall, OpenAIAgent, QueryEngineTool } from "llamaindex";
import fs from "node:fs/promises";
import path from "node:path";
import { getDataSource } from "./index";
import { STORAGE_CACHE_DIR } from "./shared";
import { createTools } from "./tools";
export async function createChatEngine() {
@@ -16,7 +17,7 @@ export async function createChatEngine() {
queryEngine: index.asQueryEngine(),
metadata: {
name: "data_query_engine",
description: `A query engine for documents from your data source.`,
description: `A query engine for documents in storage folder: ${STORAGE_CACHE_DIR}`,
},
}),
);
@@ -1,61 +0,0 @@
import { JSONSchemaType } from "ajv";
import { search } from "duck-duck-scrape";
import { BaseTool, ToolMetadata } from "llamaindex";
export type DuckDuckGoParameter = {
query: string;
region?: string;
};
export type DuckDuckGoToolParams = {
metadata?: ToolMetadata<JSONSchemaType<DuckDuckGoParameter>>;
};
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<DuckDuckGoParameter>> = {
name: "duckduckgo",
description: "Use this function to search for any query in DuckDuckGo.",
parameters: {
type: "object",
properties: {
query: {
type: "string",
description: "The query to search in DuckDuckGo.",
},
region: {
type: "string",
description:
"Optional, The region to be used for the search in [country-language] convention, ex us-en, uk-en, ru-ru, etc...",
nullable: true,
},
},
required: ["query"],
},
};
type DuckDuckGoSearchResult = {
title: string;
description: string;
url: string;
};
export class DuckDuckGoSearchTool implements BaseTool<DuckDuckGoParameter> {
metadata: ToolMetadata<JSONSchemaType<DuckDuckGoParameter>>;
constructor(params: DuckDuckGoToolParams) {
this.metadata = params.metadata ?? DEFAULT_META_DATA;
}
async call(input: DuckDuckGoParameter) {
const { query, region } = input;
const options = region ? { region } : {};
const searchResults = await search(query, options);
return searchResults.results.map((result) => {
return {
title: result.title,
description: result.description,
url: result.url,
} as DuckDuckGoSearchResult;
});
}
}
@@ -1,112 +0,0 @@
import type { JSONSchemaType } from "ajv";
import { FormData } from "formdata-node";
import fs from "fs";
import got from "got";
import { BaseTool, ToolMetadata } from "llamaindex";
import path from "node:path";
import { Readable } from "stream";
export type ImgGeneratorParameter = {
prompt: string;
};
export type ImgGeneratorToolParams = {
metadata?: ToolMetadata<JSONSchemaType<ImgGeneratorParameter>>;
};
export type ImgGeneratorToolOutput = {
isSuccess: boolean;
imageUrl?: string;
errorMessage?: string;
};
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<ImgGeneratorParameter>> = {
name: "image_generator",
description: `Use this function to generate an image based on the prompt.`,
parameters: {
type: "object",
properties: {
prompt: {
type: "string",
description: "The prompt to generate the image",
},
},
required: ["prompt"],
},
};
export class ImgGeneratorTool implements BaseTool<ImgGeneratorParameter> {
readonly IMG_OUTPUT_FORMAT = "webp";
readonly IMG_OUTPUT_DIR = "tool-output";
readonly IMG_GEN_API =
"https://api.stability.ai/v2beta/stable-image/generate/core";
metadata: ToolMetadata<JSONSchemaType<ImgGeneratorParameter>>;
constructor(params?: ImgGeneratorToolParams) {
this.checkRequiredEnvVars();
this.metadata = params?.metadata || DEFAULT_META_DATA;
}
async call(input: ImgGeneratorParameter): Promise<ImgGeneratorToolOutput> {
return await this.generateImage(input.prompt);
}
private generateImage = async (
prompt: string,
): Promise<ImgGeneratorToolOutput> => {
try {
const buffer = await this.promptToImgBuffer(prompt);
const imageUrl = this.saveImage(buffer);
return { isSuccess: true, imageUrl };
} catch (error) {
console.error(error);
return {
isSuccess: false,
errorMessage: "Failed to generate image. Please try again.",
};
}
};
private promptToImgBuffer = async (prompt: string) => {
const form = new FormData();
form.append("prompt", prompt);
form.append("output_format", this.IMG_OUTPUT_FORMAT);
const buffer = await got
.post(this.IMG_GEN_API, {
// Not sure why it shows an type error when passing form to body
// Although I follow document: https://github.com/sindresorhus/got/blob/main/documentation/2-options.md#body
// Tt still works fine, so I make casting to unknown to avoid the typescript warning
// Found a similar issue: https://github.com/sindresorhus/got/discussions/1877
body: form as unknown as Buffer | Readable | string,
headers: {
Authorization: `Bearer ${process.env.STABILITY_API_KEY}`,
Accept: "image/*",
},
})
.buffer();
return buffer;
};
private saveImage = (buffer: Buffer) => {
const filename = `${crypto.randomUUID()}.${this.IMG_OUTPUT_FORMAT}`;
const outputPath = path.join(this.IMG_OUTPUT_DIR, filename);
fs.writeFileSync(outputPath, buffer);
const url = `${process.env.FILESERVER_URL_PREFIX}/${this.IMG_OUTPUT_DIR}/${filename}`;
console.log(`Saved image to ${outputPath}.\nURL: ${url}`);
return url;
};
private checkRequiredEnvVars = () => {
if (!process.env.STABILITY_API_KEY) {
throw new Error(
"STABILITY_API_KEY key is required to run image generator. Get it here: https://platform.stability.ai/account/keys",
);
}
if (!process.env.FILESERVER_URL_PREFIX) {
throw new Error(
"FILESERVER_URL_PREFIX is required to display file output after generation",
);
}
};
}
@@ -1,61 +1,42 @@
import { BaseToolWithCall } from "llamaindex";
import { ToolsFactory } from "llamaindex/tools/ToolsFactory";
import { DuckDuckGoSearchTool, DuckDuckGoToolParams } from "./duckduckgo";
import { ImgGeneratorTool, ImgGeneratorToolParams } from "./img-gen";
import { InterpreterTool, InterpreterToolParams } from "./interpreter";
import { OpenAPIActionTool } from "./openapi-action";
import { WeatherTool, WeatherToolParams } from "./weather";
type ToolCreator = (config: unknown) => Promise<BaseToolWithCall[]>;
type ToolCreator = (config: unknown) => BaseToolWithCall;
export async function createTools(toolConfig: {
local: Record<string, unknown>;
llamahub: any;
}): Promise<BaseToolWithCall[]> {
// add local tools from the 'tools' folder (if configured)
const tools = await createLocalTools(toolConfig.local);
const tools = createLocalTools(toolConfig.local);
// add tools from LlamaIndexTS (if configured)
tools.push(...(await ToolsFactory.createTools(toolConfig.llamahub)));
return tools;
}
const toolFactory: Record<string, ToolCreator> = {
weather: async (config: unknown) => {
return [new WeatherTool(config as WeatherToolParams)];
weather: (config: unknown) => {
return new WeatherTool(config as WeatherToolParams);
},
interpreter: async (config: unknown) => {
return [new InterpreterTool(config as InterpreterToolParams)];
},
"openapi_action.OpenAPIActionToolSpec": async (config: unknown) => {
const { openapi_uri, domain_headers } = config as {
openapi_uri: string;
domain_headers: Record<string, Record<string, string>>;
};
const openAPIActionTool = new OpenAPIActionTool(
openapi_uri,
domain_headers,
);
return await openAPIActionTool.toToolFunctions();
},
duckduckgo: async (config: unknown) => {
return [new DuckDuckGoSearchTool(config as DuckDuckGoToolParams)];
},
img_gen: async (config: unknown) => {
return [new ImgGeneratorTool(config as ImgGeneratorToolParams)];
interpreter: (config: unknown) => {
return new InterpreterTool(config as InterpreterToolParams);
},
};
async function createLocalTools(
function createLocalTools(
localConfig: Record<string, unknown>,
): Promise<BaseToolWithCall[]> {
): BaseToolWithCall[] {
const tools: BaseToolWithCall[] = [];
for (const [key, toolConfig] of Object.entries(localConfig)) {
Object.keys(localConfig).forEach((key) => {
if (key in toolFactory) {
const newTools = await toolFactory[key](toolConfig);
tools.push(...newTools);
const toolConfig = localConfig[key];
const tool = toolFactory[key](toolConfig);
tools.push(tool);
}
}
});
return tools;
}
@@ -34,9 +34,8 @@ type InterpreterExtraType =
export type InterpreterExtraResult = {
type: InterpreterExtraType;
content?: string;
filename?: string;
url?: string;
filename: string;
url: string;
};
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<InterpreterParameter>> = {
@@ -107,11 +106,8 @@ export class InterpreterTool implements BaseTool<InterpreterParameter> {
async call(input: InterpreterParameter): Promise<InterpreterToolOutput> {
const result = await this.codeInterpret(input.code);
return result;
}
async close() {
await this.codeInterpreter?.close();
return result;
}
private async getExtraResult(
@@ -122,34 +118,23 @@ export class InterpreterTool implements BaseTool<InterpreterParameter> {
try {
const formats = res.formats(); // formats available for the result. Eg: ['png', ...]
const results = formats.map((f) => res[f as keyof Result]); // get base64 data for each format
const base64DataArr = formats.map((f) => res[f as keyof Result]); // get base64 data for each format
// save base64 data to file and return the url
for (let i = 0; i < formats.length; i++) {
const ext = formats[i];
const data = results[i];
switch (ext) {
case "png":
case "jpeg":
case "svg":
case "pdf":
const { filename } = this.saveToDisk(data, ext);
output.push({
type: ext as InterpreterExtraType,
filename,
url: this.getFileUrl(filename),
});
break;
default:
output.push({
type: ext as InterpreterExtraType,
content: data,
});
break;
const base64Data = base64DataArr[i];
if (ext && base64Data) {
const { filename } = this.saveToDisk(base64Data, ext);
output.push({
type: ext as InterpreterExtraType,
filename,
url: this.getFileUrl(filename),
});
}
}
} catch (error) {
console.error("Error when parsing e2b response", error);
console.error("Error when saving data to disk", error);
}
return output;
@@ -1,164 +0,0 @@
import SwaggerParser from "@apidevtools/swagger-parser";
import { JSONSchemaType } from "ajv";
import got from "got";
import { FunctionTool, JSONValue, ToolMetadata } from "llamaindex";
interface DomainHeaders {
[key: string]: { [header: string]: string };
}
type Input = {
url: string;
params: object;
};
type APIInfo = {
description: string;
title: string;
};
export class OpenAPIActionTool {
// cache the loaded specs by URL
private static specs: Record<string, any> = {};
private readonly INVALID_URL_PROMPT =
"This url did not include a hostname or scheme. Please determine the complete URL and try again.";
private createLoadSpecMetaData = (info: APIInfo) => {
return {
name: "load_openapi_spec",
description: `Use this to retrieve the OpenAPI spec for the API named ${info.title} with the following description: ${info.description}. Call it before making any requests to the API.`,
};
};
private readonly createMethodCallMetaData = (
method: "POST" | "PATCH" | "GET",
info: APIInfo,
) => {
return {
name: `${method.toLowerCase()}_request`,
description: `Use this to call the ${method} method on the API named ${info.title}`,
parameters: {
type: "object",
properties: {
url: {
type: "string",
description: `The url to make the ${method} request against`,
},
params: {
type: "object",
description:
method === "GET"
? "the URL parameters to provide with the get request"
: `the key-value pairs to provide with the ${method} request`,
},
},
required: ["url"],
},
} as ToolMetadata<JSONSchemaType<Input>>;
};
constructor(
public openapi_uri: string,
public domainHeaders: DomainHeaders = {},
) {}
async loadOpenapiSpec(url: string): Promise<any> {
const api = await SwaggerParser.validate(url);
return {
servers: "servers" in api ? api.servers : "",
info: { description: api.info.description, title: api.info.title },
endpoints: api.paths,
};
}
async getRequest(input: Input): Promise<JSONValue> {
if (!this.validUrl(input.url)) {
return this.INVALID_URL_PROMPT;
}
try {
const data = await got
.get(input.url, {
headers: this.getHeadersForUrl(input.url),
searchParams: input.params as URLSearchParams,
})
.json();
return data as JSONValue;
} catch (error) {
return error as JSONValue;
}
}
async postRequest(input: Input): Promise<JSONValue> {
if (!this.validUrl(input.url)) {
return this.INVALID_URL_PROMPT;
}
try {
const res = await got.post(input.url, {
headers: this.getHeadersForUrl(input.url),
json: input.params,
});
return res.body as JSONValue;
} catch (error) {
return error as JSONValue;
}
}
async patchRequest(input: Input): Promise<JSONValue> {
if (!this.validUrl(input.url)) {
return this.INVALID_URL_PROMPT;
}
try {
const res = await got.patch(input.url, {
headers: this.getHeadersForUrl(input.url),
json: input.params,
});
return res.body as JSONValue;
} catch (error) {
return error as JSONValue;
}
}
public async toToolFunctions() {
if (!OpenAPIActionTool.specs[this.openapi_uri]) {
console.log(`Loading spec for URL: ${this.openapi_uri}`);
const spec = await this.loadOpenapiSpec(this.openapi_uri);
OpenAPIActionTool.specs[this.openapi_uri] = spec;
}
const spec = OpenAPIActionTool.specs[this.openapi_uri];
// TODO: read endpoints with parameters from spec and create one tool for each endpoint
// For now, we just create a tool for each HTTP method which does not work well for passing parameters
return [
FunctionTool.from(() => {
return spec;
}, this.createLoadSpecMetaData(spec.info)),
FunctionTool.from(
this.getRequest.bind(this),
this.createMethodCallMetaData("GET", spec.info),
),
FunctionTool.from(
this.postRequest.bind(this),
this.createMethodCallMetaData("POST", spec.info),
),
FunctionTool.from(
this.patchRequest.bind(this),
this.createMethodCallMetaData("PATCH", spec.info),
),
];
}
private validUrl(url: string): boolean {
const parsed = new URL(url);
return !!parsed.protocol && !!parsed.hostname;
}
private getDomain(url: string): string {
const parsed = new URL(url);
return parsed.hostname;
}
private getHeadersForUrl(url: string): { [header: string]: string } {
const domain = this.getDomain(url);
return this.domainHeaders[domain] || {};
}
}
@@ -8,14 +8,14 @@ export async function createChatEngine() {
`StorageContext is empty - call 'npm run generate' to generate the storage first`,
);
}
const retriever = index.asRetriever({
similarityTopK: process.env.TOP_K ? parseInt(process.env.TOP_K) : 3,
});
const retriever = index.asRetriever();
retriever.similarityTopK = process.env.TOP_K
? parseInt(process.env.TOP_K)
: 3;
return new ContextChatEngine({
chatModel: Settings.llm,
retriever,
// disable as a custom system prompt disables the generated context
// systemPrompt: process.env.SYSTEM_PROMPT,
systemPrompt: process.env.SYSTEM_PROMPT,
});
}
+5 -14
View File
@@ -23,12 +23,7 @@ def llama_parse_parser():
"LLAMA_CLOUD_API_KEY environment variable is not set. "
"Please set it in .env file or in your shell environment then run again!"
)
parser = LlamaParse(
result_type="markdown",
verbose=True,
language="en",
ignore_errors=False,
)
parser = LlamaParse(result_type="markdown", verbose=True, language="en")
return parser
@@ -37,19 +32,15 @@ def get_file_documents(config: FileLoaderConfig):
try:
reader = SimpleDirectoryReader(
config.data_dir, recursive=True, filename_as_id=True, raise_on_error=True
config.data_dir,
recursive=True,
filename_as_id=True,
)
if config.use_llama_parse:
# LlamaParse is async first,
# so we need to use nest_asyncio to run it in sync mode
import nest_asyncio
nest_asyncio.apply()
parser = llama_parse_parser()
reader.file_extractor = {".pdf": parser}
return reader.load_data()
except Exception as e:
except ValueError as e:
import sys, traceback
# Catch the error if the data dir is empty
@@ -14,8 +14,8 @@ export interface ChatInputProps {
/** Form submission handler to automatically reset input and append a user message */
handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void;
isLoading: boolean;
multiModal?: boolean;
messages: Message[];
setInput?: (input: string) => void;
}
export default function ChatInput(props: ChatInputProps) {
@@ -19,12 +19,8 @@ export default function ChatMessages({
isLoading?: boolean;
stop?: () => void;
reload?: () => void;
append?: (
message: Message | Omit<Message, "id">,
) => Promise<string | null | undefined>;
}) {
const scrollableChatContainerRef = useRef<HTMLDivElement>(null);
const lastMessage = messages[messages.length - 1];
const scrollToBottom = () => {
if (scrollableChatContainerRef.current) {
@@ -35,14 +31,14 @@ export default function ChatMessages({
useEffect(() => {
scrollToBottom();
}, [messages.length, lastMessage]);
}, [messages.length]);
return (
<div
className="flex-1 w-full max-w-5xl p-4 bg-white rounded-xl shadow-xl overflow-auto"
ref={scrollableChatContainerRef}
>
<div className="flex flex-col gap-5 divide-y">
<div className="w-full max-w-5xl p-4 bg-white rounded-xl shadow-xl">
<div
className="flex flex-col gap-5 divide-y h-[50vh] overflow-auto"
ref={scrollableChatContainerRef}
>
{messages.map((m: Message) => (
<ChatItem key={m.id} {...m} />
))}
@@ -1,30 +0,0 @@
"use client";
import { useEffect, useMemo, useState } from "react";
export interface ChatConfig {
chatAPI?: string;
starterQuestions?: string[];
}
export function useClientConfig() {
const API_ROUTE = "/api/chat/config";
const chatAPI = process.env.NEXT_PUBLIC_CHAT_API;
const [config, setConfig] = useState<ChatConfig>({
chatAPI,
});
const configAPI = useMemo(() => {
const backendOrigin = chatAPI ? new URL(chatAPI).origin : "";
return `${backendOrigin}${API_ROUTE}`;
}, [chatAPI]);
useEffect(() => {
fetch(configAPI)
.then((response) => response.json())
.then((data) => setConfig({ ...data, chatAPI }))
.catch((error) => console.error("Error fetching config", error));
}, [chatAPI, configAPI]);
return config;
}
@@ -1,45 +0,0 @@
from dotenv import load_dotenv
load_dotenv()
import os
import logging
from app.settings import init_settings
from app.engine.loaders import get_documents
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
def generate_datasource():
init_settings()
logger.info("Generate index for the provided data")
name = os.getenv("LLAMA_CLOUD_INDEX_NAME")
project_name = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
api_key = os.getenv("LLAMA_CLOUD_API_KEY")
base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
if name is None or project_name is None or api_key is None:
raise ValueError(
"Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
" to your environment variables or config them in .env file"
)
documents = get_documents()
LlamaCloudIndex.from_documents(
documents=documents,
name=name,
project_name=project_name,
api_key=api_key,
base_url=base_url,
)
logger.info("Finished generating the index")
if __name__ == "__main__":
generate_datasource()
@@ -1,28 +0,0 @@
import logging
import os
from llama_index.indices.managed.llama_cloud import LlamaCloudIndex
logger = logging.getLogger("uvicorn")
def get_index():
name = os.getenv("LLAMA_CLOUD_INDEX_NAME")
project_name = os.getenv("LLAMA_CLOUD_PROJECT_NAME")
api_key = os.getenv("LLAMA_CLOUD_API_KEY")
base_url = os.getenv("LLAMA_CLOUD_BASE_URL")
if name is None or project_name is None or api_key is None:
raise ValueError(
"Please set LLAMA_CLOUD_INDEX_NAME, LLAMA_CLOUD_PROJECT_NAME and LLAMA_CLOUD_API_KEY"
" to your environment variables or config them in .env file"
)
index = LlamaCloudIndex(
name=name,
project_name=project_name,
api_key=api_key,
base_url=base_url,
)
return index
@@ -1,26 +0,0 @@
import * as dotenv from "dotenv";
import { LlamaCloudIndex } from "llamaindex";
import { getDocuments } from "./loader";
import { initSettings } from "./settings";
import { checkRequiredEnvVars } from "./shared";
dotenv.config();
async function loadAndIndex() {
const documents = await getDocuments();
await LlamaCloudIndex.fromDocuments({
documents,
name: process.env.LLAMA_CLOUD_INDEX_NAME!,
projectName: process.env.LLAMA_CLOUD_PROJECT_NAME!,
apiKey: process.env.LLAMA_CLOUD_API_KEY,
baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
});
console.log(`Successfully created embeddings!`);
}
(async () => {
checkRequiredEnvVars();
initSettings();
await loadAndIndex();
console.log("Finished generating storage.");
})();
@@ -1,13 +0,0 @@
import { LlamaCloudIndex } from "llamaindex/cloud/LlamaCloudIndex";
import { checkRequiredEnvVars } from "./shared";
export async function getDataSource() {
checkRequiredEnvVars();
const index = new LlamaCloudIndex({
name: process.env.LLAMA_CLOUD_INDEX_NAME!,
projectName: process.env.LLAMA_CLOUD_PROJECT_NAME!,
apiKey: process.env.LLAMA_CLOUD_API_KEY,
baseUrl: process.env.LLAMA_CLOUD_BASE_URL,
});
return index;
}
@@ -1,22 +0,0 @@
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(", ")}`,
);
}
}
@@ -3,8 +3,5 @@
"rules": {
"max-params": ["error", 4],
"prefer-const": "error"
},
"parserOptions": {
"sourceType": "module"
}
}
+5 -14
View File
@@ -1,32 +1,23 @@
{
"name": "llama-index-express-streaming",
"version": "1.0.0",
"exports": "./index.js",
"types": "./index.d.ts",
"type": "module",
"engines": {
"node": ">=18"
},
"main": "dist/index.js",
"scripts": {
"format": "prettier --ignore-unknown --cache --check .",
"format:write": "prettier --ignore-unknown --write .",
"build": "tsup index.ts --format esm --dts",
"build": "tsup index.ts --format cjs --dts",
"start": "node dist/index.js",
"dev": "concurrently \"tsup index.ts --format esm --dts --watch\" \"nodemon --watch dist/index.js\""
"dev": "concurrently \"tsup index.ts --format cjs --dts --watch\" \"nodemon -q dist/index.js\""
},
"dependencies": {
"ai": "^3.0.21",
"cors": "^2.8.5",
"dotenv": "^16.3.1",
"duck-duck-scrape": "^2.2.5",
"express": "^4.18.2",
"llamaindex": "0.4.6",
"llamaindex": "0.3.13",
"pdf2json": "3.0.5",
"ajv": "^8.12.0",
"@e2b/code-interpreter": "^0.0.5",
"got": "^14.4.1",
"@apidevtools/swagger-parser": "^10.1.0",
"formdata-node": "^6.0.3"
"@e2b/code-interpreter": "^0.0.5"
},
"devDependencies": {
"@types/cors": "^2.8.16",
@@ -1,14 +0,0 @@
import { Request, Response } from "express";
export const chatConfig = async (_req: Request, res: Response) => {
let starterQuestions = undefined;
if (
process.env.CONVERSATION_STARTERS &&
process.env.CONVERSATION_STARTERS.trim()
) {
starterQuestions = process.env.CONVERSATION_STARTERS.trim().split("\n");
}
return res.status(200).json({
starterQuestions,
});
};
@@ -2,7 +2,11 @@ import { Message, StreamData, streamToResponse } from "ai";
import { Request, Response } from "express";
import { ChatMessage, Settings } from "llamaindex";
import { createChatEngine } from "./engine/chat";
import { LlamaIndexStream, convertMessageContent } from "./llamaindex-stream";
import {
DataParserOptions,
LlamaIndexStream,
convertMessageContent,
} from "./llamaindex-stream";
import { createCallbackManager, createStreamTimeout } from "./stream-helper";
export const chat = async (req: Request, res: Response) => {
@@ -10,7 +14,10 @@ export const chat = async (req: Request, res: Response) => {
const vercelStreamData = new StreamData();
const streamTimeout = createStreamTimeout(vercelStreamData);
try {
const { messages }: { messages: Message[] } = req.body;
const {
messages,
data,
}: { messages: Message[]; data: DataParserOptions | undefined } = req.body;
const userMessage = messages.pop();
if (!messages || !userMessage || userMessage.role !== "user") {
return res.status(400).json({
@@ -21,24 +28,8 @@ export const chat = async (req: Request, res: Response) => {
const chatEngine = await createChatEngine();
let annotations = userMessage.annotations;
if (!annotations) {
// the user didn't send any new annotations with the last message
// so use the annotations from the last user message that has annotations
// REASON: GPT4 doesn't consider MessageContentDetail from previous messages, only strings
annotations = messages
.slice()
.reverse()
.find(
(message) => message.role === "user" && message.annotations,
)?.annotations;
}
// Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
const userMessageContent = convertMessageContent(
userMessage.content,
annotations,
);
const userMessageContent = convertMessageContent(userMessage.content, data);
// Setup callbacks
const callbackManager = createCallbackManager(vercelStreamData);
@@ -53,7 +44,12 @@ export const chat = async (req: Request, res: Response) => {
});
// Return a stream, which can be consumed by the Vercel/AI client
const stream = LlamaIndexStream(response, vercelStreamData);
const stream = LlamaIndexStream(response, vercelStreamData, {
parserOptions: {
imageUrl: data?.imageUrl,
uploadedCsv: data?.uploadedCsv,
},
});
return streamToResponse(stream, res, {}, vercelStreamData);
} catch (error) {
@@ -4,7 +4,6 @@ import {
GEMINI_MODEL,
Gemini,
GeminiEmbedding,
Groq,
OpenAI,
OpenAIEmbedding,
Settings,
@@ -29,9 +28,6 @@ export const initSettings = async () => {
case "ollama":
initOllama();
break;
case "groq":
initGroq();
break;
case "anthropic":
initAnthropic();
break;
@@ -49,9 +45,7 @@ export const initSettings = async () => {
function initOpenAI() {
Settings.llm = new OpenAI({
model: process.env.MODEL ?? "gpt-3.5-turbo",
maxTokens: process.env.LLM_MAX_TOKENS
? Number(process.env.LLM_MAX_TOKENS)
: undefined,
maxTokens: 512,
});
Settings.embedModel = new OpenAIEmbedding({
model: process.env.EMBEDDING_MODEL,
@@ -89,27 +83,6 @@ function initAnthropic() {
});
}
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",
};
const modelMap: Record<string, string> = {
"llama3-8b": "llama3-8b-8192",
"llama3-70b": "llama3-70b-8192",
"mixtral-8x7b": "mixtral-8x7b-32768",
};
Settings.llm = new Groq({
model: modelMap[process.env.MODEL!],
});
Settings.embedModel = new HuggingFaceEmbedding({
modelType: embedModelMap[process.env.EMBEDDING_MODEL!],
});
}
function initGemini() {
Settings.llm = new Gemini({
model: process.env.MODEL as GEMINI_MODEL,
@@ -1,5 +1,4 @@
import {
JSONValue,
StreamData,
createCallbacksTransformer,
createStreamDataTransformer,
@@ -7,90 +6,103 @@ import {
type AIStreamCallbacksAndOptions,
} from "ai";
import {
EngineResponse,
MessageContent,
MessageContentDetail,
Metadata,
NodeWithScore,
Response,
ToolCallLLMMessageOptions,
} from "llamaindex";
import { CsvFile } from "./stream-helper";
import { AgentStreamChatResponse } from "llamaindex/agent/base";
import {
UploadedCsv,
appendCsvData,
appendImageData,
appendSourceData,
} from "./stream-helper";
export const convertMessageContent = (
content: string,
annotations?: JSONValue[],
): MessageContent => {
if (!annotations) return content;
return [
{
type: "text",
text: content,
},
...convertAnnotations(annotations),
];
type LlamaIndexResponse =
| AgentStreamChatResponse<ToolCallLLMMessageOptions>
| Response;
export type DataParserOptions = {
imageUrl?: string;
uploadedCsv?: UploadedCsv;
};
const convertAnnotations = (
annotations: JSONValue[],
): MessageContentDetail[] => {
const content: MessageContentDetail[] = [];
annotations.forEach((annotation: JSONValue) => {
// first skip invalid annotation
if (
!(
annotation &&
typeof annotation === "object" &&
"type" in annotation &&
"data" in annotation &&
annotation.data &&
typeof annotation.data === "object"
)
) {
console.log(
"Client sent invalid annotation. Missing data and type",
annotation,
);
return;
}
const { type, data } = annotation;
// convert image
if (type === "image" && "url" in data && typeof data.url === "string") {
content.push({
type: "image_url",
image_url: {
url: data.url,
},
});
}
// convert CSV files to text
if (type === "csv" && "csvFiles" in data && Array.isArray(data.csvFiles)) {
const rawContents = data.csvFiles.map((csv) => {
return "```csv\n" + (csv as CsvFile).content + "\n```";
});
const csvContent =
"Use data from following CSV raw contents:\n" +
rawContents.join("\n\n");
content.push({
type: "text",
text: csvContent,
});
}
});
export const convertMessageContent = (
textMessage: string,
additionalData?: DataParserOptions,
): MessageContent => {
if (!additionalData) return textMessage;
const content: MessageContent = [
{
type: "text",
text: textMessage,
},
];
if (additionalData?.imageUrl) {
content.push({
type: "image_url",
image_url: {
url: additionalData?.imageUrl,
},
});
}
if (additionalData?.uploadedCsv) {
const csvContent =
"Use the following CSV data:\n" +
"```csv\n" +
additionalData.uploadedCsv.content +
"\n```";
content.push({
type: "text",
text: `${csvContent}\n\n${textMessage}`,
});
}
return content;
};
function createParser(res: AsyncIterable<EngineResponse>, data: StreamData) {
function createParser(
res: AsyncIterable<LlamaIndexResponse>,
data: StreamData,
opts?: DataParserOptions,
) {
const it = res[Symbol.asyncIterator]();
const trimStartOfStream = trimStartOfStreamHelper();
let sourceNodes: NodeWithScore<Metadata>[] | undefined;
return new ReadableStream<string>({
start() {
appendImageData(data, opts?.imageUrl);
appendCsvData(data, opts?.uploadedCsv);
},
async pull(controller): Promise<void> {
const { value, done } = await it.next();
if (done) {
if (sourceNodes) {
appendSourceData(data, sourceNodes);
}
controller.close();
data.close();
return;
}
const text = trimStartOfStream(value.delta ?? "");
let delta;
if (value instanceof Response) {
// handle Response type
if (value.sourceNodes) {
// get source nodes from the first response
sourceNodes = value.sourceNodes;
}
delta = value.response ?? "";
} else {
// handle other types
delta = value.response.delta;
}
const text = trimStartOfStream(delta ?? "");
if (text) {
controller.enqueue(text);
}
@@ -99,13 +111,14 @@ function createParser(res: AsyncIterable<EngineResponse>, data: StreamData) {
}
export function LlamaIndexStream(
response: AsyncIterable<EngineResponse>,
response: AsyncIterable<LlamaIndexResponse>,
data: StreamData,
opts?: {
callbacks?: AIStreamCallbacksAndOptions;
parserOptions?: DataParserOptions;
},
): ReadableStream<Uint8Array> {
return createParser(response, data)
return createParser(response, data, opts?.parserOptions)
.pipeThrough(createCallbacksTransformer(opts?.callbacks))
.pipeThrough(createStreamDataTransformer());
}
@@ -7,6 +7,16 @@ import {
ToolOutput,
} from "llamaindex";
export function appendImageData(data: StreamData, imageUrl?: string) {
if (!imageUrl) return;
data.appendMessageAnnotation({
type: "image",
data: {
url: imageUrl,
},
});
}
function getNodeUrl(metadata: Metadata) {
const url = metadata["URL"];
if (url) return url;
@@ -84,9 +94,8 @@ export function createStreamTimeout(stream: StreamData) {
export function createCallbackManager(stream: StreamData) {
const callbackManager = new CallbackManager();
callbackManager.on("retrieve-end", (data) => {
const { nodes, query } = data.detail.payload;
appendSourceData(stream, nodes);
callbackManager.on("retrieve", (data) => {
const { nodes, query } = data.detail;
appendEventData(stream, `Retrieving context for query: '${query}'`);
appendEventData(
stream,
@@ -113,9 +122,16 @@ export function createCallbackManager(stream: StreamData) {
return callbackManager;
}
export type CsvFile = {
export type UploadedCsv = {
content: string;
filename: string;
filesize: number;
id: string;
};
export function appendCsvData(data: StreamData, uploadedCsv?: UploadedCsv) {
if (!uploadedCsv) return;
data.appendMessageAnnotation({
type: "csv",
data: uploadedCsv,
});
}
@@ -1,5 +1,4 @@
import express, { Router } from "express";
import { chatConfig } from "../controllers/chat-config.controller";
import { chatRequest } from "../controllers/chat-request.controller";
import { chat } from "../controllers/chat.controller";
import { initSettings } from "../controllers/engine/settings";
@@ -9,6 +8,5 @@ const llmRouter: Router = express.Router();
initSettings();
llmRouter.route("/").post(chat);
llmRouter.route("/request").post(chatRequest);
llmRouter.route("/config").get(chatConfig);
export default llmRouter;
@@ -1,114 +1,170 @@
import os
import logging
from aiostream import stream
from pydantic import BaseModel
from typing import List, Any, Optional, Dict, Tuple
from fastapi import APIRouter, Depends, HTTPException, Request, status
from llama_index.core.chat_engine.types import BaseChatEngine
from llama_index.core.llms import MessageRole
from llama_index.core.schema import NodeWithScore
from llama_index.core.llms import ChatMessage, MessageRole
from app.engine import get_chat_engine
from app.api.routers.vercel_response import VercelStreamResponse
from app.api.routers.events import EventCallbackHandler
from app.api.routers.models import (
ChatData,
ChatConfig,
SourceNodes,
Result,
Message,
)
from app.api.routers.messaging import EventCallbackHandler
from aiostream import stream
chat_router = r = APIRouter()
logger = logging.getLogger("uvicorn")
class _Message(BaseModel):
role: MessageRole
content: str
class _ChatData(BaseModel):
messages: List[_Message]
class Config:
json_schema_extra = {
"example": {
"messages": [
{
"role": "user",
"content": "What standards for letters exist?",
}
]
}
}
class _SourceNodes(BaseModel):
id: str
metadata: Dict[str, Any]
score: Optional[float]
text: str
url: Optional[str]
@classmethod
def from_source_node(cls, source_node: NodeWithScore):
metadata = source_node.node.metadata
url = metadata.get("URL")
if not url:
file_name = metadata.get("file_name")
url_prefix = os.getenv("FILESERVER_URL_PREFIX")
if not url_prefix:
logger.warning("Warning: FILESERVER_URL_PREFIX not set in environment variables")
if file_name and url_prefix:
url = f"{url_prefix}/data/{file_name}"
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 from_source_nodes(cls, source_nodes: List[NodeWithScore]):
return [cls.from_source_node(node) for node in source_nodes]
class _Result(BaseModel):
result: _Message
nodes: List[_SourceNodes]
async def parse_chat_data(data: _ChatData) -> Tuple[str, List[ChatMessage]]:
# check preconditions and get last message
if len(data.messages) == 0:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="No messages provided",
)
last_message = data.messages.pop()
if last_message.role != MessageRole.USER:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Last message must be from user",
)
# convert messages coming from the request to type ChatMessage
messages = [
ChatMessage(
role=m.role,
content=m.content,
)
for m in data.messages
]
return last_message.content, messages
# streaming endpoint - delete if not needed
@r.post("")
async def chat(
request: Request,
data: ChatData,
data: _ChatData,
chat_engine: BaseChatEngine = Depends(get_chat_engine),
):
last_message_content, messages = await parse_chat_data(data)
event_handler = EventCallbackHandler()
chat_engine.callback_manager.handlers.append(event_handler) # type: ignore
try:
last_message_content = data.get_last_message_content()
messages = data.get_history_messages()
event_handler = EventCallbackHandler()
chat_engine.callback_manager.handlers.append(event_handler) # type: ignore
async def content_generator():
# Yield the text response
async def _chat_response_generator():
response = await chat_engine.astream_chat(
last_message_content, messages
)
async for token in response.async_response_gen():
yield VercelStreamResponse.convert_text(token)
# the text_generator is the leading stream, once it's finished, also finish the event stream
event_handler.is_done = True
# Yield the source nodes
yield VercelStreamResponse.convert_data(
{
"type": "sources",
"data": {
"nodes": [
SourceNodes.from_source_node(node).dict()
for node in response.source_nodes
]
},
}
)
# Yield the events from the event handler
async def _event_generator():
async for event in event_handler.async_event_gen():
event_response = event.to_response()
if event_response is not None:
yield VercelStreamResponse.convert_data(event_response)
combine = stream.merge(_chat_response_generator(), _event_generator())
is_stream_started = False
async with combine.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 VercelStreamResponse.convert_text("")
yield output
if await request.is_disconnected():
break
return VercelStreamResponse(content=content_generator())
response = await chat_engine.astream_chat(last_message_content, messages)
except Exception as e:
logger.exception("Error in chat engine", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error in chat engine: {e}",
) from e
)
async def content_generator():
# Yield the text response
async def _text_generator():
async for token in response.async_response_gen():
yield VercelStreamResponse.convert_text(token)
# the text_generator is the leading stream, once it's finished, also finish the event stream
event_handler.is_done = True
# Yield the events from the event handler
async def _event_generator():
async for event in event_handler.async_event_gen():
event_response = event.to_response()
if event_response is not None:
yield VercelStreamResponse.convert_data(event_response)
combine = stream.merge(_text_generator(), _event_generator())
async with combine.stream() as streamer:
async for item in streamer:
if await request.is_disconnected():
break
yield item
# Yield the source nodes
yield VercelStreamResponse.convert_data(
{
"type": "sources",
"data": {
"nodes": [
_SourceNodes.from_source_node(node).dict()
for node in response.source_nodes
]
},
}
)
return VercelStreamResponse(content=content_generator())
# non-streaming endpoint - delete if not needed
@r.post("/request")
async def chat_request(
data: ChatData,
data: _ChatData,
chat_engine: BaseChatEngine = Depends(get_chat_engine),
) -> Result:
last_message_content = data.get_last_message_content()
messages = data.get_history_messages()
) -> _Result:
last_message_content, messages = await parse_chat_data(data)
response = await chat_engine.achat(last_message_content, messages)
return Result(
result=Message(role=MessageRole.ASSISTANT, content=response.response),
nodes=SourceNodes.from_source_nodes(response.source_nodes),
return _Result(
result=_Message(role=MessageRole.ASSISTANT, content=response.response),
nodes=_SourceNodes.from_source_nodes(response.source_nodes),
)
@r.get("/config")
async def chat_config() -> ChatConfig:
starter_questions = None
conversation_starters = os.getenv("CONVERSATION_STARTERS")
if conversation_starters and conversation_starters.strip():
starter_questions = conversation_starters.strip().split("\n")
return ChatConfig(starterQuestions=starter_questions)
@@ -1,170 +0,0 @@
import os
import logging
from pydantic import BaseModel, Field, validator
from pydantic.alias_generators import to_camel
from typing import List, Any, Optional, Dict
from llama_index.core.schema import NodeWithScore
from llama_index.core.llms import ChatMessage, MessageRole
logger = logging.getLogger("uvicorn")
class CsvFile(BaseModel):
content: str
filename: str
filesize: int
id: str
class AnnotationData(BaseModel):
csv_files: List[CsvFile] | None = Field(
default=None,
description="List of CSV files",
)
class Config:
json_schema_extra = {
"example": {
"csvFiles": [
{
"content": "Name, Age\nAlice, 25\nBob, 30",
"filename": "example.csv",
"filesize": 123,
"id": "123",
"type": "text/csv",
}
]
}
}
alias_generator = to_camel
class Annotation(BaseModel):
type: str
data: AnnotationData
def to_content(self) -> str:
if self.type == "csv":
csv_files = self.data.csv_files
if csv_files is not None and len(csv_files) > 0:
return "Use data from following CSV raw contents\n" + "\n".join(
[f"```csv\n{csv_file.content}\n```" for csv_file in csv_files]
)
raise ValueError(f"Unsupported annotation type: {self.type}")
class Message(BaseModel):
role: MessageRole
content: str
annotations: List[Annotation] | None = None
class ChatData(BaseModel):
messages: List[Message]
class Config:
json_schema_extra = {
"example": {
"messages": [
{
"role": "user",
"content": "What standards for letters exist?",
}
]
}
}
@validator("messages")
def messages_must_not_be_empty(cls, v):
if len(v) == 0:
raise ValueError("Messages must not be empty")
return v
def get_last_message_content(self) -> str:
"""
Get the content of the last message along with the data content if available. Fallback to use data content from previous messages
"""
if len(self.messages) == 0:
raise ValueError("There is not any message in the chat")
last_message = self.messages[-1]
message_content = last_message.content
for message in reversed(self.messages):
if message.role == MessageRole.USER and message.annotations is not None:
annotation_contents = (
annotation.to_content() for annotation in message.annotations
)
annotation_text = "\n".join(annotation_contents)
message_content = f"{message_content}\n{annotation_text}"
break
return message_content
def get_history_messages(self) -> List[Message]:
"""
Get the history messages
"""
return [
ChatMessage(role=message.role, content=message.content)
for message in self.messages[:-1]
]
def is_last_message_from_user(self) -> bool:
return self.messages[-1].role == MessageRole.USER
class SourceNodes(BaseModel):
id: str
metadata: Dict[str, Any]
score: Optional[float]
text: str
url: Optional[str]
@classmethod
def from_source_node(cls, source_node: NodeWithScore):
metadata = source_node.node.metadata
url = metadata.get("URL")
if not url:
file_name = metadata.get("file_name")
url_prefix = os.getenv("FILESERVER_URL_PREFIX")
if not url_prefix:
logger.warning(
"Warning: FILESERVER_URL_PREFIX not set in environment variables"
)
if file_name and url_prefix:
url = f"{url_prefix}/data/{file_name}"
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 from_source_nodes(cls, source_nodes: List[NodeWithScore]):
return [cls.from_source_node(node) for node in source_nodes]
class Result(BaseModel):
result: Message
nodes: List[SourceNodes]
class ChatConfig(BaseModel):
starter_questions: Optional[List[str]] = Field(
default=None,
description="List of starter questions",
)
class Config:
json_schema_extra = {
"example": {
"starterQuestions": [
"What standards for letters exist?",
"What are the requirements for a letter to be considered a letter?",
]
}
}
alias_generator = to_camel
@@ -1,51 +1,40 @@
import os
from typing import Dict
from llama_index.core.settings import Settings
def init_settings():
model_provider = os.getenv("MODEL_PROVIDER")
match model_provider:
case "openai":
init_openai()
case "groq":
init_groq()
case "ollama":
init_ollama()
case "anthropic":
init_anthropic()
case "gemini":
init_gemini()
case "azure-openai":
init_azure_openai()
case _:
raise ValueError(f"Invalid model provider: {model_provider}")
if model_provider == "openai":
init_openai()
elif model_provider == "ollama":
init_ollama()
elif model_provider == "anthropic":
init_anthropic()
elif model_provider == "gemini":
init_gemini()
else:
raise ValueError(f"Invalid model provider: {model_provider}")
Settings.chunk_size = int(os.getenv("CHUNK_SIZE", "1024"))
Settings.chunk_overlap = int(os.getenv("CHUNK_OVERLAP", "20"))
def init_ollama():
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.llms.ollama.base import DEFAULT_REQUEST_TIMEOUT, Ollama
base_url = os.getenv("OLLAMA_BASE_URL") or "http://127.0.0.1:11434"
request_timeout = float(
os.getenv("OLLAMA_REQUEST_TIMEOUT", DEFAULT_REQUEST_TIMEOUT)
)
Settings.embed_model = OllamaEmbedding(
base_url=base_url,
model_name=os.getenv("EMBEDDING_MODEL"),
)
Settings.llm = Ollama(
base_url=base_url, model=os.getenv("MODEL"), request_timeout=request_timeout
)
Settings.llm = Ollama(base_url=base_url, model=os.getenv("MODEL"))
def init_openai():
from llama_index.core.constants import DEFAULT_TEMPERATURE
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core.constants import DEFAULT_TEMPERATURE
max_tokens = os.getenv("LLM_MAX_TOKENS")
config = {
@@ -63,58 +52,9 @@ def init_openai():
Settings.embed_model = OpenAIEmbedding(**config)
def init_azure_openai():
from llama_index.core.constants import DEFAULT_TEMPERATURE
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.llms.azure_openai import AzureOpenAI
llm_deployment = os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT")
embedding_deployment = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
max_tokens = os.getenv("LLM_MAX_TOKENS")
api_key = os.getenv("AZURE_OPENAI_API_KEY")
llm_config = {
"api_key": api_key,
"deployment_name": llm_deployment,
"model": os.getenv("MODEL"),
"temperature": float(os.getenv("LLM_TEMPERATURE", DEFAULT_TEMPERATURE)),
"max_tokens": int(max_tokens) if max_tokens is not None else None,
}
Settings.llm = AzureOpenAI(**llm_config)
dimensions = os.getenv("EMBEDDING_DIM")
embedding_config = {
"api_key": api_key,
"deployment_name": embedding_deployment,
"model": os.getenv("EMBEDDING_MODEL"),
"dimensions": int(dimensions) if dimensions is not None else None,
}
Settings.embed_model = AzureOpenAIEmbedding(**embedding_config)
def init_groq():
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.groq import Groq
model_map: Dict[str, str] = {
"llama3-8b": "llama3-8b-8192",
"llama3-70b": "llama3-70b-8192",
"mixtral-8x7b": "mixtral-8x7b-32768",
}
embed_model_map: Dict[str, str] = {
"all-MiniLM-L6-v2": "sentence-transformers/all-MiniLM-L6-v2",
"all-mpnet-base-v2": "sentence-transformers/all-mpnet-base-v2",
}
Settings.llm = Groq(model=model_map[os.getenv("MODEL")])
Settings.embed_model = HuggingFaceEmbedding(
model_name=embed_model_map[os.getenv("EMBEDDING_MODEL")]
)
def init_anthropic():
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.anthropic import Anthropic
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
model_map: Dict[str, str] = {
"claude-3-opus": "claude-3-opus-20240229",
@@ -136,11 +76,21 @@ def init_anthropic():
def init_gemini():
from llama_index.embeddings.gemini import GeminiEmbedding
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
model_name = f"models/{os.getenv('MODEL')}"
embed_model_name = f"models/{os.getenv('EMBEDDING_MODEL')}"
model_map: Dict[str, str] = {
"gemini-1.5-pro-latest": "models/gemini-1.5-pro-latest",
"gemini-pro": "models/gemini-pro",
"gemini-pro-vision": "models/gemini-pro-vision",
}
Settings.llm = Gemini(model=model_name)
Settings.embed_model = GeminiEmbedding(model_name=embed_model_name)
embed_model_map: Dict[str, str] = {
"embedding-001": "models/embedding-001",
"text-embedding-004": "models/text-embedding-004",
}
Settings.llm = Gemini(model=model_map[os.getenv("MODEL")])
Settings.embed_model = GeminiEmbedding(
model_name=embed_model_map[os.getenv("EMBEDDING_MODEL")]
)
@@ -14,8 +14,8 @@ fastapi = "^0.109.1"
uvicorn = { extras = ["standard"], version = "^0.23.2" }
python-dotenv = "^1.0.0"
aiostream = "^0.5.2"
llama-index = "0.10.50"
llama-index-core = "0.10.50"
llama-index = "0.10.28"
llama-index-core = "0.10.28"
cachetools = "^5.3.3"
[build-system]
@@ -1,11 +0,0 @@
import { NextResponse } from "next/server";
/**
* This API is to get config from the backend envs and expose them to the frontend
*/
export async function GET() {
const config = {
starterQuestions: process.env.CONVERSATION_STARTERS?.trim().split("\n"),
};
return NextResponse.json(config, { status: 200 });
}
@@ -4,7 +4,6 @@ import {
GEMINI_MODEL,
Gemini,
GeminiEmbedding,
Groq,
OpenAI,
OpenAIEmbedding,
Settings,
@@ -29,9 +28,6 @@ export const initSettings = async () => {
case "ollama":
initOllama();
break;
case "groq":
initGroq();
break;
case "anthropic":
initAnthropic();
break;
@@ -49,9 +45,7 @@ export const initSettings = async () => {
function initOpenAI() {
Settings.llm = new OpenAI({
model: process.env.MODEL ?? "gpt-3.5-turbo",
maxTokens: process.env.LLM_MAX_TOKENS
? Number(process.env.LLM_MAX_TOKENS)
: undefined,
maxTokens: 512,
});
Settings.embedModel = new OpenAIEmbedding({
model: process.env.EMBEDDING_MODEL,
@@ -75,27 +69,6 @@ function initOllama() {
});
}
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",
};
const modelMap: Record<string, string> = {
"llama3-8b": "llama3-8b-8192",
"llama3-70b": "llama3-70b-8192",
"mixtral-8x7b": "mixtral-8x7b-32768",
};
Settings.llm = new Groq({
model: modelMap[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",
@@ -1,5 +1,4 @@
import {
JSONValue,
StreamData,
createCallbacksTransformer,
createStreamDataTransformer,
@@ -7,90 +6,103 @@ import {
type AIStreamCallbacksAndOptions,
} from "ai";
import {
EngineResponse,
MessageContent,
MessageContentDetail,
Metadata,
NodeWithScore,
Response,
ToolCallLLMMessageOptions,
} from "llamaindex";
import { CsvFile } from "./stream-helper";
import { AgentStreamChatResponse } from "llamaindex/agent/base";
import {
UploadedCsv,
appendCsvData,
appendImageData,
appendSourceData,
} from "./stream-helper";
export const convertMessageContent = (
content: string,
annotations?: JSONValue[],
): MessageContent => {
if (!annotations) return content;
return [
{
type: "text",
text: content,
},
...convertAnnotations(annotations),
];
type LlamaIndexResponse =
| AgentStreamChatResponse<ToolCallLLMMessageOptions>
| Response;
export type DataParserOptions = {
imageUrl?: string;
uploadedCsv?: UploadedCsv;
};
const convertAnnotations = (
annotations: JSONValue[],
): MessageContentDetail[] => {
const content: MessageContentDetail[] = [];
annotations.forEach((annotation: JSONValue) => {
// first skip invalid annotation
if (
!(
annotation &&
typeof annotation === "object" &&
"type" in annotation &&
"data" in annotation &&
annotation.data &&
typeof annotation.data === "object"
)
) {
console.log(
"Client sent invalid annotation. Missing data and type",
annotation,
);
return;
}
const { type, data } = annotation;
// convert image
if (type === "image" && "url" in data && typeof data.url === "string") {
content.push({
type: "image_url",
image_url: {
url: data.url,
},
});
}
// convert CSV files to text
if (type === "csv" && "csvFiles" in data && Array.isArray(data.csvFiles)) {
const rawContents = data.csvFiles.map((csv) => {
return "```csv\n" + (csv as CsvFile).content + "\n```";
});
const csvContent =
"Use data from following CSV raw contents:\n" +
rawContents.join("\n\n");
content.push({
type: "text",
text: csvContent,
});
}
});
export const convertMessageContent = (
textMessage: string,
additionalData?: DataParserOptions,
): MessageContent => {
if (!additionalData) return textMessage;
const content: MessageContent = [
{
type: "text",
text: textMessage,
},
];
if (additionalData?.imageUrl) {
content.push({
type: "image_url",
image_url: {
url: additionalData?.imageUrl,
},
});
}
if (additionalData?.uploadedCsv) {
const csvContent =
"Use the following CSV data:\n" +
"```csv\n" +
additionalData.uploadedCsv.content +
"\n```";
content.push({
type: "text",
text: `${csvContent}\n\n${textMessage}`,
});
}
return content;
};
function createParser(res: AsyncIterable<EngineResponse>, data: StreamData) {
function createParser(
res: AsyncIterable<LlamaIndexResponse>,
data: StreamData,
opts?: DataParserOptions,
) {
const it = res[Symbol.asyncIterator]();
const trimStartOfStream = trimStartOfStreamHelper();
let sourceNodes: NodeWithScore<Metadata>[] | undefined;
return new ReadableStream<string>({
start() {
appendImageData(data, opts?.imageUrl);
appendCsvData(data, opts?.uploadedCsv);
},
async pull(controller): Promise<void> {
const { value, done } = await it.next();
if (done) {
if (sourceNodes) {
appendSourceData(data, sourceNodes);
}
controller.close();
data.close();
return;
}
const text = trimStartOfStream(value.delta ?? "");
let delta;
if (value instanceof Response) {
// handle Response type
if (value.sourceNodes) {
// get source nodes from the first response
sourceNodes = value.sourceNodes;
}
delta = value.response ?? "";
} else {
// handle other types
delta = value.response.delta;
}
const text = trimStartOfStream(delta ?? "");
if (text) {
controller.enqueue(text);
}
@@ -99,13 +111,14 @@ function createParser(res: AsyncIterable<EngineResponse>, data: StreamData) {
}
export function LlamaIndexStream(
response: AsyncIterable<EngineResponse>,
response: AsyncIterable<LlamaIndexResponse>,
data: StreamData,
opts?: {
callbacks?: AIStreamCallbacksAndOptions;
parserOptions?: DataParserOptions;
},
): ReadableStream<Uint8Array> {
return createParser(response, data)
return createParser(response, data, opts?.parserOptions)
.pipeThrough(createCallbacksTransformer(opts?.callbacks))
.pipeThrough(createStreamDataTransformer());
}
@@ -4,7 +4,11 @@ import { ChatMessage, Settings } from "llamaindex";
import { NextRequest, NextResponse } from "next/server";
import { createChatEngine } from "./engine/chat";
import { initSettings } from "./engine/settings";
import { LlamaIndexStream, convertMessageContent } from "./llamaindex-stream";
import {
DataParserOptions,
LlamaIndexStream,
convertMessageContent,
} from "./llamaindex-stream";
import { createCallbackManager, createStreamTimeout } from "./stream-helper";
initObservability();
@@ -20,7 +24,10 @@ export async function POST(request: NextRequest) {
try {
const body = await request.json();
const { messages }: { messages: Message[] } = body;
const {
messages,
data,
}: { messages: Message[]; data: DataParserOptions | undefined } = body;
const userMessage = messages.pop();
if (!messages || !userMessage || userMessage.role !== "user") {
return NextResponse.json(
@@ -34,24 +41,8 @@ export async function POST(request: NextRequest) {
const chatEngine = await createChatEngine();
let annotations = userMessage.annotations;
if (!annotations) {
// the user didn't send any new annotations with the last message
// so use the annotations from the last user message that has annotations
// REASON: GPT4 doesn't consider MessageContentDetail from previous messages, only strings
annotations = messages
.slice()
.reverse()
.find(
(message) => message.role === "user" && message.annotations,
)?.annotations;
}
// Convert message content from Vercel/AI format to LlamaIndex/OpenAI format
const userMessageContent = convertMessageContent(
userMessage.content,
annotations,
);
const userMessageContent = convertMessageContent(userMessage.content, data);
// Setup callbacks
const callbackManager = createCallbackManager(vercelStreamData);
@@ -66,7 +57,12 @@ export async function POST(request: NextRequest) {
});
// Transform LlamaIndex stream to Vercel/AI format
const stream = LlamaIndexStream(response, vercelStreamData);
const stream = LlamaIndexStream(response, vercelStreamData, {
parserOptions: {
imageUrl: data?.imageUrl,
uploadedCsv: data?.uploadedCsv,
},
});
// Return a StreamingTextResponse, which can be consumed by the Vercel/AI client
return new StreamingTextResponse(stream, {}, vercelStreamData);
@@ -7,6 +7,16 @@ import {
ToolOutput,
} from "llamaindex";
export function appendImageData(data: StreamData, imageUrl?: string) {
if (!imageUrl) return;
data.appendMessageAnnotation({
type: "image",
data: {
url: imageUrl,
},
});
}
function getNodeUrl(metadata: Metadata) {
const url = metadata["URL"];
if (url) return url;
@@ -84,9 +94,8 @@ export function createStreamTimeout(stream: StreamData) {
export function createCallbackManager(stream: StreamData) {
const callbackManager = new CallbackManager();
callbackManager.on("retrieve-end", (data) => {
const { nodes, query } = data.detail.payload;
appendSourceData(stream, nodes);
callbackManager.on("retrieve", (data) => {
const { nodes, query } = data.detail;
appendEventData(stream, `Retrieving context for query: '${query}'`);
appendEventData(
stream,
@@ -113,9 +122,16 @@ export function createCallbackManager(stream: StreamData) {
return callbackManager;
}
export type CsvFile = {
export type UploadedCsv = {
content: string;
filename: string;
filesize: number;
id: string;
};
export function appendCsvData(data: StreamData, uploadedCsv?: UploadedCsv) {
if (!uploadedCsv) return;
data.appendMessageAnnotation({
type: "csv",
data: uploadedCsv,
});
}
@@ -2,10 +2,8 @@
import { useChat } from "ai/react";
import { ChatInput, ChatMessages } from "./ui/chat";
import { useClientConfig } from "./ui/chat/hooks/use-config";
export default function ChatSection() {
const { chatAPI } = useClientConfig();
const {
messages,
input,
@@ -14,10 +12,8 @@ export default function ChatSection() {
handleInputChange,
reload,
stop,
append,
setInput,
} = useChat({
api: chatAPI,
api: process.env.NEXT_PUBLIC_CHAT_API,
headers: {
"Content-Type": "application/json", // using JSON because of vercel/ai 2.2.26
},
@@ -29,13 +25,12 @@ export default function ChatSection() {
});
return (
<div className="space-y-4 w-full h-full flex flex-col">
<div className="space-y-4 max-w-5xl w-full">
<ChatMessages
messages={messages}
isLoading={isLoading}
reload={reload}
stop={stop}
append={append}
/>
<ChatInput
input={input}
@@ -43,8 +38,6 @@ export default function ChatSection() {
handleInputChange={handleInputChange}
isLoading={isLoading}
messages={messages}
append={append}
setInput={setInput}
/>
</div>
);
@@ -7,7 +7,7 @@ export default function Header() {
Get started by editing&nbsp;
<code className="font-mono font-bold">app/page.tsx</code>
</p>
<div className="fixed bottom-0 left-0 mb-4 flex h-auto w-full items-end justify-center bg-gradient-to-t from-white via-white dark:from-black dark:via-black lg:static lg:w-auto lg:bg-none lg:mb-0">
<div className="fixed bottom-0 left-0 flex h-48 w-full items-end justify-center bg-gradient-to-t from-white via-white dark:from-black dark:via-black lg:static lg:h-auto lg:w-auto lg:bg-none">
<a
href="https://www.llamaindex.ai/"
className="flex items-center justify-center font-nunito text-lg font-bold gap-2"
@@ -1,12 +1,12 @@
import { ChevronDown, ChevronRight, Loader2 } from "lucide-react";
import { useState } from "react";
import { Button } from "../../button";
import { Button } from "../button";
import {
Collapsible,
CollapsibleContent,
CollapsibleTrigger,
} from "../../collapsible";
import { EventData } from "../index";
} from "../collapsible";
import { EventData } from "./index";
export function ChatEvents({
data,
@@ -1,5 +1,5 @@
import Image from "next/image";
import { type ImageData } from "../index";
import { type ImageData } from "./index";
export function ChatImage({ data }: { data: ImageData }) {
return (
@@ -1,14 +1,12 @@
import { JSONValue } from "ai";
import { useState } from "react";
import { v4 as uuidv4 } from "uuid";
import { MessageAnnotation, MessageAnnotationType } from ".";
import { useEffect, useState } from "react";
import { CsvData, getInputResources } from ".";
import { Button } from "../button";
import FileUploader from "../file-uploader";
import { Input } from "../input";
import UploadCsvPreview from "../upload-csv-preview";
import UploadImagePreview from "../upload-image-preview";
import ChatResources from "./chat-resources";
import { ChatHandler } from "./chat.interface";
import { useCsv } from "./hooks/use-csv";
export default function ChatInput(
props: Pick<
@@ -20,97 +18,74 @@ export default function ChatInput(
| "handleSubmit"
| "handleInputChange"
| "messages"
| "setInput"
| "append"
>,
) {
const [imageUrl, setImageUrl] = useState<string | null>(null);
const { files: csvFiles, upload, remove, reset } = useCsv();
const [uploadedCsv, setUploadedCsv] = useState<CsvData>();
const [inputResources, setInputResources] = useState<
Array<CsvData & { selected: boolean }>
>([]);
const getAnnotations = () => {
if (!imageUrl && csvFiles.length === 0) return undefined;
const annotations: MessageAnnotation[] = [];
if (imageUrl) {
annotations.push({
type: MessageAnnotationType.IMAGE,
data: { url: imageUrl },
});
}
if (csvFiles.length > 0) {
annotations.push({
type: MessageAnnotationType.CSV,
data: {
csvFiles: csvFiles.map((file) => ({
id: file.id,
content: file.content,
filename: file.filename,
filesize: file.filesize,
})),
},
});
}
return annotations as JSONValue[];
};
// default submit function does not handle including annotations in the message
// so we need to use append function to submit new message with annotations
const handleSubmitWithAnnotations = (
e: React.FormEvent<HTMLFormElement>,
annotations: JSONValue[] | undefined,
) => {
e.preventDefault();
props.append!({
content: props.input,
role: "user",
createdAt: new Date(),
annotations,
});
props.setInput!("");
};
useEffect(() => {
const resources = getInputResources(props.messages);
setInputResources(
resources.csv.map((data) => ({ ...data, selected: true })),
);
}, [props.messages]);
const onSubmit = (e: React.FormEvent<HTMLFormElement>) => {
const annotations = getAnnotations();
if (annotations) {
handleSubmitWithAnnotations(e, annotations);
imageUrl && setImageUrl(null);
csvFiles.length && reset();
if (imageUrl) {
props.handleSubmit(e, {
data: { imageUrl: imageUrl },
});
setImageUrl(null);
return;
}
// if users upload a new csv file, we will send it to backend
if (uploadedCsv) {
props.handleSubmit(e, {
data: { uploadedCsv },
});
setUploadedCsv(undefined);
return;
}
// if users upload a new csv file, we can reuse provided csv resources
const attachCsv = inputResources.filter((r) => r.selected)[0];
if (attachCsv) {
props.handleSubmit(e, {
data: { uploadedCsv: attachCsv },
});
return;
}
props.handleSubmit(e);
};
const onRemovePreviewImage = () => setImageUrl(null);
const readContent = async (file: File): Promise<string> => {
const content = await new Promise<string>((resolve, reject) => {
const handleUploadImageFile = async (file: File) => {
const base64 = await new Promise<string>((resolve, reject) => {
const reader = new FileReader();
if (file.type.startsWith("image/")) {
reader.readAsDataURL(file);
} else {
reader.readAsText(file);
}
reader.readAsDataURL(file);
reader.onload = () => resolve(reader.result as string);
reader.onerror = (error) => reject(error);
});
return content;
};
const handleUploadImageFile = async (file: File) => {
const base64 = await readContent(file);
setImageUrl(base64);
};
const handleUploadCsvFile = async (file: File) => {
const content = await readContent(file);
const isSuccess = upload({
id: uuidv4(),
const content = await new Promise<string>((resolve, reject) => {
const reader = new FileReader();
reader.readAsText(file);
reader.onload = () => resolve(reader.result as string);
reader.onerror = (error) => reject(error);
});
setUploadedCsv({
content,
filename: file.name,
filesize: file.size,
});
if (!isSuccess) {
alert("File already exists in the list.");
}
};
const handleUploadFile = async (file: File) => {
@@ -119,10 +94,6 @@ export default function ChatInput(
return await handleUploadImageFile(file);
}
if (file.type === "text/csv") {
if (csvFiles.length > 0) {
alert("You can only upload one csv file at a time.");
return;
}
return await handleUploadCsvFile(file);
}
props.onFileUpload?.(file);
@@ -131,26 +102,33 @@ export default function ChatInput(
}
};
const removeResource = (index: number) => {
setInputResources((resources) => {
const newResources = [...resources];
newResources[index].selected = false;
return newResources;
});
};
return (
<form
onSubmit={onSubmit}
className="rounded-xl bg-white p-4 shadow-xl space-y-4 shrink-0"
className="rounded-xl bg-white p-4 shadow-xl space-y-4"
>
<ChatResources
isLoading={props.isLoading}
resources={inputResources}
removeResource={removeResource}
/>
{imageUrl && (
<UploadImagePreview url={imageUrl} onRemove={onRemovePreviewImage} />
)}
{csvFiles.length > 0 && (
<div className="flex gap-4 w-full overflow-auto py-2">
{csvFiles.map((csv) => {
return (
<UploadCsvPreview
key={csv.id}
csv={csv}
onRemove={() => remove(csv)}
/>
);
})}
</div>
{uploadedCsv && (
<UploadCsvPreview
filename={uploadedCsv.filename}
filesize={uploadedCsv.filesize}
onRemove={() => setUploadedCsv(undefined)}
/>
)}
<div className="flex w-full items-start justify-between gap-4 ">
<Input
@@ -165,7 +143,7 @@ export default function ChatInput(
onFileUpload={handleUploadFile}
onFileError={props.onFileError}
/>
<Button type="submit" disabled={props.isLoading || !props.input.trim()}>
<Button type="submit" disabled={props.isLoading}>
Send message
</Button>
</div>
@@ -2,8 +2,13 @@ import { Check, Copy } from "lucide-react";
import { Message } from "ai";
import { Fragment } from "react";
import { Button } from "../../button";
import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
import { Button } from "../button";
import ChatAvatar from "./chat-avatar";
import { ChatEvents } from "./chat-events";
import { ChatImage } from "./chat-image";
import { ChatSources } from "./chat-sources";
import ChatTools from "./chat-tools";
import CsvContent from "./csv-content";
import {
CsvData,
EventData,
@@ -13,14 +18,9 @@ import {
SourceData,
ToolData,
getAnnotationData,
} from "../index";
import ChatAvatar from "./chat-avatar";
import { ChatEvents } from "./chat-events";
import { ChatImage } from "./chat-image";
import { ChatSources } from "./chat-sources";
import ChatTools from "./chat-tools";
import CsvContent from "./csv-content";
} from "./index";
import Markdown from "./markdown";
import { useCopyToClipboard } from "./use-copy-to-clipboard";
type ContentDisplayConfig = {
order: number;
@@ -60,7 +60,7 @@ function ChatMessageContent({
const contents: ContentDisplayConfig[] = [
{
order: 1,
order: -4,
component: imageData[0] ? <ChatImage data={imageData[0]} /> : null,
},
{
@@ -71,7 +71,7 @@ function ChatMessageContent({
) : null,
},
{
order: 2,
order: -2,
component: csvData[0] ? <CsvContent data={csvData[0]} /> : null,
},
{
@@ -83,7 +83,7 @@ function ChatMessageContent({
component: <Markdown content={message.content} />,
},
{
order: 3,
order: 1,
component: sourceData[0] ? <ChatSources data={sourceData[0]} /> : null,
},
];
@@ -1,13 +0,0 @@
import UploadCsvPreview from "../../upload-csv-preview";
import { CsvData } from "../index";
export default function CsvContent({ data }: { data: CsvData }) {
if (!data.csvFiles.length) return null;
return (
<div className="flex gap-2 items-center">
{data.csvFiles.map((csv, index) => (
<UploadCsvPreview key={index} csv={csv} />
))}
</div>
);
}
@@ -1,19 +1,13 @@
import { Loader2 } from "lucide-react";
import { useEffect, useRef } from "react";
import { Button } from "../button";
import ChatActions from "./chat-actions";
import ChatMessage from "./chat-message";
import { ChatHandler } from "./chat.interface";
import { useClientConfig } from "./hooks/use-config";
export default function ChatMessages(
props: Pick<
ChatHandler,
"messages" | "isLoading" | "reload" | "stop" | "append"
>,
props: Pick<ChatHandler, "messages" | "isLoading" | "reload" | "stop">,
) {
const { starterQuestions } = useClientConfig();
const scrollableChatContainerRef = useRef<HTMLDivElement>(null);
const messageLength = props.messages.length;
const lastMessage = props.messages[messageLength - 1];
@@ -41,11 +35,11 @@ export default function ChatMessages(
}, [messageLength, lastMessage]);
return (
<div
className="flex-1 w-full rounded-xl bg-white p-4 shadow-xl relative overflow-y-auto"
ref={scrollableChatContainerRef}
>
<div className="flex flex-col gap-5 divide-y">
<div className="w-full rounded-xl bg-white p-4 shadow-xl pb-0">
<div
className="flex h-[50vh] flex-col gap-5 divide-y overflow-y-auto pb-4"
ref={scrollableChatContainerRef}
>
{props.messages.map((m, i) => {
const isLoadingMessage = i === messageLength - 1 && props.isLoading;
return (
@@ -62,33 +56,14 @@ export default function ChatMessages(
</div>
)}
</div>
{(showReload || showStop) && (
<div className="flex justify-end py-4">
<ChatActions
reload={props.reload}
stop={props.stop}
showReload={showReload}
showStop={showStop}
/>
</div>
)}
{!messageLength && starterQuestions?.length && props.append && (
<div className="absolute bottom-6 left-0 w-full">
<div className="grid grid-cols-2 gap-2 mx-20">
{starterQuestions.map((question, i) => (
<Button
variant="outline"
key={i}
onClick={() =>
props.append!({ role: "user", content: question })
}
>
{question}
</Button>
))}
</div>
</div>
)}
<div className="flex justify-end py-4">
<ChatActions
reload={props.reload}
stop={props.stop}
showReload={showReload}
showStop={showStop}
/>
</div>
</div>
);
}
@@ -0,0 +1,48 @@
import { Loader2, XIcon } from "lucide-react";
import Image from "next/image";
import { CsvData } from ".";
import SheetIcon from "../../ui/icons/sheet.svg";
export interface ChatResourcesProps {
isLoading: boolean;
resources: Array<CsvData & { selected: boolean }>;
removeResource: (index: number) => void;
}
export default function ChatResources(props: ChatResourcesProps) {
if (!props.resources.length) return null;
return (
<div className="flex gap-4 text-sm">
{props.resources.map((data, index) => {
if (!data.selected) return null;
const fileSizeInKB = Math.round((data.filesize / 1024) * 10) / 10;
return (
<div
className="border-2 border-green-700 py-2 px-3 rounded-lg flex gap-2 items-center"
key={data.filename}
>
<div className="h-4 w-4 shrink-0 rounded-md">
<Image
className="h-full w-auto"
priority
src={SheetIcon}
alt="SheetIcon"
/>
</div>
<span>
{data.filename} - {fileSizeInKB} KB
</span>
{props.isLoading ? (
<Loader2 className="h-4 w-4 animate-spin" />
) : (
<XIcon
className="w-4 h-4 cursor-pointer"
onClick={() => props.removeResource(index)}
/>
)}
</div>
);
})}
</div>
);
}
@@ -1,14 +1,10 @@
import { Check, Copy } from "lucide-react";
import { useMemo } from "react";
import { Button } from "../../button";
import {
HoverCard,
HoverCardContent,
HoverCardTrigger,
} from "../../hover-card";
import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
import { SourceData } from "../index";
import PdfDialog from "../widgets/PdfDialog";
import { Button } from "../button";
import { HoverCard, HoverCardContent, HoverCardTrigger } from "../hover-card";
import { SourceData } from "./index";
import { useCopyToClipboard } from "./use-copy-to-clipboard";
import PdfDialog from "./widgets/PdfDialog";
const SCORE_THRESHOLD = 0.3;
@@ -1,5 +1,5 @@
import { ToolData } from "../index";
import { WeatherCard, WeatherData } from "../widgets/WeatherCard";
import { ToolData } from "./index";
import { WeatherCard, WeatherData } from "./widgets/WeatherCard";
// TODO: If needed, add displaying more tool outputs here
export default function ChatTools({ data }: { data: ToolData }) {
@@ -15,11 +15,4 @@ export interface ChatHandler {
stop?: () => void;
onFileUpload?: (file: File) => Promise<void>;
onFileError?: (errMsg: string) => void;
setInput?: (input: string) => void;
append?: (
message: Message | Omit<Message, "id">,
ops?: {
data: any;
},
) => Promise<string | null | undefined>;
}
@@ -5,8 +5,8 @@ import { FC, memo } from "react";
import { Prism, SyntaxHighlighterProps } from "react-syntax-highlighter";
import { coldarkDark } from "react-syntax-highlighter/dist/cjs/styles/prism";
import { Button } from "../../button";
import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
import { Button } from "../button";
import { useCopyToClipboard } from "./use-copy-to-clipboard";
// TODO: Remove this when @type/react-syntax-highlighter is updated
const SyntaxHighlighter = Prism as unknown as FC<SyntaxHighlighterProps>;
@@ -0,0 +1,15 @@
import { CsvData } from ".";
const LIMIT_DISPLAY = 100; // Limit the display of CSV content to 100 characters
export default function CsvContent({ data }: { data: CsvData }) {
const summaryContent = data.content.slice(0, LIMIT_DISPLAY) + "...";
return (
<div className="space-y-2">
<h3 className="font-semibold">CSV Raw Content</h3>
<pre className="bg-secondary max-h-[200px] overflow-auto rounded-md p-4 block text-sm">
{summaryContent}
</pre>
</div>
);
}
@@ -1,30 +0,0 @@
"use client";
import { useEffect, useMemo, useState } from "react";
export interface ChatConfig {
chatAPI?: string;
starterQuestions?: string[];
}
export function useClientConfig() {
const API_ROUTE = "/api/chat/config";
const chatAPI = process.env.NEXT_PUBLIC_CHAT_API;
const [config, setConfig] = useState<ChatConfig>({
chatAPI,
});
const configAPI = useMemo(() => {
const backendOrigin = chatAPI ? new URL(chatAPI).origin : "";
return `${backendOrigin}${API_ROUTE}`;
}, [chatAPI]);
useEffect(() => {
fetch(configAPI)
.then((response) => response.json())
.then((data) => setConfig({ ...data, chatAPI }))
.catch((error) => console.error("Error fetching config", error));
}, [chatAPI, configAPI]);
return config;
}
@@ -1,33 +0,0 @@
"use client";
import { useState } from "react";
import { CsvFile } from "../index";
export function useCsv() {
const [files, setFiles] = useState<CsvFile[]>([]);
const csvEqual = (a: CsvFile, b: CsvFile) => {
if (a.id === b.id) return true;
if (a.filename === b.filename && a.filesize === b.filesize) return true;
return false;
};
const upload = (file: CsvFile) => {
const existedCsv = files.find((f) => csvEqual(f, file));
if (!existedCsv) {
setFiles((prev) => [...prev, file]);
return true;
}
return false;
};
const remove = (file: CsvFile) => {
setFiles((prev) => prev.filter((f) => f.id !== file.id));
};
const reset = () => {
setFiles([]);
};
return { files, upload, remove, reset };
}
@@ -1,4 +1,4 @@
import { JSONValue } from "ai";
import { JSONValue, Message } from "ai";
import ChatInput from "./chat-input";
import ChatMessages from "./chat-messages";
@@ -17,15 +17,10 @@ export type ImageData = {
url: string;
};
export type CsvFile = {
export type CsvData = {
content: string;
filename: string;
filesize: number;
id: string;
};
export type CsvData = {
csvFiles: CsvFile[];
};
export type SourceNode = {
@@ -77,3 +72,32 @@ export function getAnnotationData<T extends AnnotationData>(
): T[] {
return annotations.filter((a) => a.type === type).map((a) => a.data as T);
}
// this function is used to get the additional resources for a message
// it filters the annotations of a message and returns the unique resources
// currently only CSV resources are supported
export const getInputResources = (
messages: Message[],
): {
csv: Array<CsvData>;
} => {
const csvResources: CsvData[] = [];
messages.forEach((message) => {
if (message.annotations) {
const csvData = getAnnotationData<CsvData>(
message.annotations as MessageAnnotation[],
MessageAnnotationType.CSV,
);
csvData.forEach((data) => {
if (
csvResources.findIndex((r) => r.filename === data.filename) === -1
) {
csvResources.push(data);
}
});
}
});
return {
csv: csvResources,
};
};
@@ -17,30 +17,19 @@ const MemoizedReactMarkdown: FC<Options> = memo(
const preprocessLaTeX = (content: string) => {
// Replace block-level LaTeX delimiters \[ \] with $$ $$
const blockProcessedContent = content.replace(
/\\\[([\s\S]*?)\\\]/g,
/\\\[(.*?)\\\]/gs,
(_, equation) => `$$${equation}$$`,
);
// Replace inline LaTeX delimiters \( \) with $ $
const inlineProcessedContent = blockProcessedContent.replace(
/\\\[([\s\S]*?)\\\]/g,
/\\\((.*?)\\\)/gs,
(_, equation) => `$${equation}$`,
);
return inlineProcessedContent;
};
const preprocessMedia = (content: string) => {
// Remove `sandbox:` from the beginning of the URL
// to fix OpenAI's models issue appending `sandbox:` to the relative URL
return content.replace(/(sandbox|attachment|snt):/g, "");
};
const preprocessContent = (content: string) => {
return preprocessMedia(preprocessLaTeX(content));
};
export default function Markdown({ content }: { content: string }) {
const processedContent = preprocessContent(content);
const processedContent = preprocessLaTeX(content);
return (
<MemoizedReactMarkdown
className="prose dark:prose-invert prose-p:leading-relaxed prose-pre:p-0 break-words custom-markdown"
@@ -1,59 +1,20 @@
import { XCircleIcon } from "lucide-react";
import Image from "next/image";
import SheetIcon from "../ui/icons/sheet.svg";
import { Button } from "./button";
import { CsvFile } from "./chat";
import {
Drawer,
DrawerClose,
DrawerContent,
DrawerDescription,
DrawerHeader,
DrawerTitle,
DrawerTrigger,
} from "./drawer";
import { cn } from "./lib/utils";
export interface UploadCsvPreviewProps {
csv: CsvFile;
onRemove?: () => void;
}
export default function UploadCsvPreview(props: UploadCsvPreviewProps) {
const { filename, filesize, content } = props.csv;
export default function UploadCsvPreview({
filename,
filesize,
onRemove,
}: {
filename: string;
filesize: number;
onRemove: () => void;
}) {
const fileSizeInKB = Math.round((filesize / 1024) * 10) / 10;
return (
<Drawer direction="left">
<DrawerTrigger asChild>
<div>
<CSVSummaryCard {...props} />
</div>
</DrawerTrigger>
<DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
<DrawerHeader className="flex justify-between">
<div className="space-y-2">
<DrawerTitle>Csv Raw Content</DrawerTitle>
<DrawerDescription>
{filename} ({inKB(filesize)} KB)
</DrawerDescription>
</div>
<DrawerClose asChild>
<Button variant="outline">Close</Button>
</DrawerClose>
</DrawerHeader>
<div className="m-4 max-h-[80%] overflow-auto">
<pre className="bg-secondary rounded-md p-4 block text-sm">
{content}
</pre>
</div>
</DrawerContent>
</Drawer>
);
}
function CSVSummaryCard(props: UploadCsvPreviewProps) {
const { onRemove, csv } = props;
return (
<div className="p-2 w-60 max-w-60 bg-secondary rounded-lg text-sm relative cursor-pointer">
<div className="p-2 w-80 bg-secondary rounded-lg text-sm relative">
<div className="flex flex-row items-center gap-2">
<div className="relative h-10 w-10 shrink-0 overflow-hidden rounded-md">
<Image
@@ -65,29 +26,21 @@ function CSVSummaryCard(props: UploadCsvPreviewProps) {
</div>
<div className="overflow-hidden">
<div className="truncate font-semibold">
{csv.filename} ({inKB(csv.filesize)} KB)
</div>
<div className="truncate text-token-text-tertiary flex items-center gap-2">
<span>Spreadsheet</span>
{filename} ({fileSizeInKB} KB)
</div>
<div className="truncate text-token-text-tertiary">Spreadsheet</div>
</div>
</div>
{onRemove && (
<div
className={cn(
"absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full",
)}
>
<XCircleIcon
className="w-6 h-6 bg-gray-500 text-white rounded-full"
onClick={onRemove}
/>
</div>
)}
<div
className={cn(
"absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full",
)}
>
<XCircleIcon
className="w-6 h-6 bg-gray-500 text-white rounded-full"
onClick={onRemove}
/>
</div>
</div>
);
}
function inKB(size: number) {
return Math.round((size / 1024) * 10) / 10;
}
@@ -74,11 +74,8 @@
* {
@apply border-border;
}
html {
@apply h-full;
}
body {
@apply bg-background text-foreground h-full;
@apply bg-background text-foreground;
font-feature-settings:
"rlig" 1,
"calt" 1;
@@ -3,13 +3,9 @@ import ChatSection from "./components/chat-section";
export default function Home() {
return (
<main className="h-full w-full flex justify-center items-center background-gradient">
<div className="space-y-2 lg:space-y-10 w-[90%] lg:w-[60rem]">
<Header />
<div className="h-[65vh] flex">
<ChatSection />
</div>
</div>
<main className="flex min-h-screen flex-col items-center gap-10 p-24 background-gradient">
<Header />
<ChatSection />
</main>
);
}
+7 -14
View File
@@ -10,9 +10,6 @@
"lint": "next lint"
},
"dependencies": {
"@apidevtools/swagger-parser": "^10.1.0",
"@e2b/code-interpreter": "^0.0.5",
"@llamaindex/pdf-viewer": "^1.1.1",
"@radix-ui/react-collapsible": "^1.0.3",
"@radix-ui/react-hover-card": "^1.0.7",
"@radix-ui/react-slot": "^1.0.2",
@@ -21,38 +18,34 @@
"class-variance-authority": "^0.7.0",
"clsx": "^2.1.1",
"dotenv": "^16.3.1",
"duck-duck-scrape": "^2.2.5",
"formdata-node": "^6.0.3",
"got": "^14.4.1",
"llamaindex": "0.4.6",
"llamaindex": "0.3.13",
"lucide-react": "^0.294.0",
"next": "^14.2.4",
"next": "^14.0.3",
"pdf2json": "3.0.5",
"react": "^18.2.0",
"react-dom": "^18.2.0",
"react-markdown": "^8.0.7",
"react-syntax-highlighter": "^15.5.0",
"rehype-katex": "^7.0.0",
"remark": "^14.0.3",
"remark-code-import": "^1.2.0",
"remark-gfm": "^3.0.1",
"remark-math": "^5.1.1",
"rehype-katex": "^7.0.0",
"supports-color": "^8.1.1",
"tailwind-merge": "^2.1.0",
"tiktoken": "^1.0.15",
"uuid": "^9.0.1",
"vaul": "^0.9.1"
"vaul": "^0.9.1",
"@llamaindex/pdf-viewer": "^1.1.1",
"@e2b/code-interpreter": "^0.0.5"
},
"devDependencies": {
"@types/node": "^20.10.3",
"@types/react": "^18.2.42",
"@types/react-dom": "^18.2.17",
"@types/react-syntax-highlighter": "^15.5.11",
"@types/uuid": "^9.0.8",
"autoprefixer": "^10.4.16",
"cross-env": "^7.0.3",
"eslint": "^8.55.0",
"eslint-config-next": "^14.2.4",
"eslint-config-next": "^14.0.3",
"eslint-config-prettier": "^8.10.0",
"postcss": "^8.4.32",
"prettier": "^3.2.5",