Compare commits

...

5 Commits

Author SHA1 Message Date
leehuwuj 4fba331e76 format code and fix missing default EMBEDDING_DIM 2024-10-08 14:01:35 +07:00
leehuwuj e3e277867c add e2b upload file for Nextjs 2024-06-06 11:38:37 +07:00
leehuwuj 2220a2e261 add sandbox file path to the meta 2024-06-06 10:37:55 +07:00
leehuwuj 84a5064020 add using tempfile for the csv 2024-06-06 10:06:24 +07:00
leehuwuj 229a11c576 update code interpreter to support upload file to e2b sandbox 2024-06-05 14:46:49 +07:00
9 changed files with 143 additions and 39 deletions
@@ -13,7 +13,6 @@ class ToolType:
class ToolFactory:
TOOL_SOURCE_PACKAGE_MAP = {
ToolType.LLAMAHUB: "llama_index.tools",
ToolType.LOCAL: "app.engine.tools",
@@ -3,7 +3,7 @@ import logging
import base64
import uuid
from pydantic import BaseModel
from typing import List, Tuple, Dict
from typing import List, Tuple, Dict, Optional
from llama_index.core.tools import FunctionTool
from e2b_code_interpreter import CodeInterpreter
from e2b_code_interpreter.models import Logs
@@ -14,18 +14,18 @@ logger = logging.getLogger(__name__)
class InterpreterExtraResult(BaseModel):
type: str
filename: str
url: str
content: Optional[str] = None
filename: Optional[str] = None
url: Optional[str] = None
class E2BToolOutput(BaseModel):
is_error: bool
logs: Logs
results: List[InterpreterExtraResult] = []
results: List[InterpreterExtraResult | str] = []
class E2BCodeInterpreter:
output_dir = "tool-output"
def __init__(self, api_key: str, filesever_url_prefix: str):
@@ -62,8 +62,9 @@ class E2BCodeInterpreter:
def parse_result(self, result) -> List[InterpreterExtraResult]:
"""
The result could include multiple formats (e.g. png, svg, etc.) but encoded in base64
We save each result to disk and return saved file metadata (extension, filename, url)
The result format could be either a base64 string (png, svg, etc.) or a raw text (text, html, markdown,...)
If it's base64, we save each result to disk and return saved file metadata (extension, filename, url),
otherwise just return the raw text content
"""
if not result:
return []
@@ -72,31 +73,46 @@ class E2BCodeInterpreter:
try:
formats = result.formats()
base64_data_arr = [result[format] for format in formats]
data_list = [result[format] for format in formats]
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)
for ext, data in zip(formats, data_list):
match ext:
case "png" | "jpeg" | "svg":
result = self.save_to_disk(data, ext)
filename = result["filename"]
output.append(
InterpreterExtraResult(
type=ext,
filename=filename,
url=self.get_file_url(filename),
)
)
)
break
case "text" | "html" | "markdown":
output.append(InterpreterExtraResult(type=ext, content=data))
except Exception as error:
logger.error("Error when saving data to disk", error)
return output
def interpret(self, code: str) -> E2BToolOutput:
def interpret(self, code: str, file_path: Optional[str] = None) -> E2BToolOutput:
with CodeInterpreter(api_key=self.api_key) as interpreter:
# Upload file to E2B sandbox
if file_path is not None:
with open(file_path, "rb") as f:
remote_path = interpreter.upload_file(f)
# Execute the code to analyze the file
logger.info(
f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}"
)
exec = interpreter.notebook.exec_cell(code)
if exec.error:
output = E2BToolOutput(is_error=True, logs=[exec.error])
logger.error(
f"Error when executing code in E2B sandbox: {exec.error} {exec.logs}"
)
output = E2BToolOutput(is_error=True, logs=exec.logs, results=[])
else:
if len(exec.results) == 0:
output = E2BToolOutput(is_error=False, logs=exec.logs, results=[])
@@ -108,9 +124,15 @@ class E2BCodeInterpreter:
return output
def code_interpret(code: str) -> Dict:
def code_interpret(code: str, local_file_path: str) -> Dict:
"""
Execute python code in a Jupyter notebook cell and return any result, stdout, stderr, display_data, and error.
Use this tool to analyze the provided data in a sandbox environment.
The tool will:
1. Upload the provided file from local to the sandbox. The uploaded file path will be /home/user/{filename}
2. Execute python code in a Jupyter notebook cell to analyze the uploaded file in the sandbox.
3. Get the result from the code in stdout, stderr, display_data, and error.
You must to provide the code and the provided file path to run this tool.
Your code should read the file from the sandbox path /home/user/{filename}.
"""
api_key = os.getenv("E2B_API_KEY")
filesever_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
@@ -126,7 +148,7 @@ def code_interpret(code: str) -> Dict:
interpreter = E2BCodeInterpreter(
api_key=api_key, filesever_url_prefix=filesever_url_prefix
)
output = interpreter.interpret(code)
output = interpreter.interpret(code, local_file_path)
return output.dict()
@@ -7,6 +7,7 @@ import path from "node:path";
export type InterpreterParameter = {
code: string;
localFilePath: string;
};
export type InterpreterToolParams = {
@@ -34,14 +35,21 @@ type InterpreterExtraType =
export type InterpreterExtraResult = {
type: InterpreterExtraType;
content?: string;
filename: string;
url: string;
};
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<InterpreterParameter>> = {
name: "interpreter",
description:
"Execute python code in a Jupyter notebook cell and return any result, stdout, stderr, display_data, and error.",
description: `Use this tool to analyze the provided data in a sandbox environment.
The tool will:
1. Upload the provided file from local to the sandbox. The uploaded file path will be /home/user/{filename}
2. Execute python code in a Jupyter notebook cell to analyze the uploaded file in the sandbox.
3. Get the result from the code in stdout, stderr, display_data, and error.
You must to provide the code and the provided file path to run this tool.
Your code should read the file from the sandbox path /home/user/{filename}.
`,
parameters: {
type: "object",
properties: {
@@ -49,6 +57,10 @@ const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<InterpreterParameter>> = {
type: "string",
description: "The python code to execute in a single cell.",
},
localFilePath: {
type: "string",
description: "The local file path to upload to the sandbox.",
},
},
required: ["code"],
},
@@ -88,11 +100,22 @@ export class InterpreterTool implements BaseTool<InterpreterParameter> {
return this.codeInterpreter;
}
public async codeInterpret(code: string): Promise<InterpreterToolOutput> {
public async codeInterpret(
code: string,
localFilePath: string,
): Promise<InterpreterToolOutput> {
const interpreter = await this.initInterpreter();
// Upload file to sandbox
console.log(`Uploading file ${localFilePath} to sandbox`);
const fileBuffer = fs.readFileSync(localFilePath);
const fileName = path.basename(localFilePath);
await interpreter.uploadFile(fileBuffer, fileName);
console.log(`Uploaded file ${fileName} to sandbox`);
// Execute code in sandbox
console.log(
`\n${"=".repeat(50)}\n> Running following AI-generated code:\n${code}\n${"=".repeat(50)}`,
);
const interpreter = await this.initInterpreter();
const exec = await interpreter.notebook.execCell(code);
if (exec.error) console.error("[Code Interpreter error]", exec.error);
const extraResult = await this.getExtraResult(exec.results[0]);
@@ -105,7 +128,7 @@ export class InterpreterTool implements BaseTool<InterpreterParameter> {
}
async call(input: InterpreterParameter): Promise<InterpreterToolOutput> {
const result = await this.codeInterpret(input.code);
const result = await this.codeInterpret(input.code, input.localFilePath);
await this.codeInterpreter?.close();
return result;
}
@@ -119,18 +142,26 @@ export class InterpreterTool implements BaseTool<InterpreterParameter> {
try {
const formats = res.formats(); // formats available for the result. Eg: ['png', ...]
const base64DataArr = formats.map((f) => res[f as keyof Result]); // get base64 data for each format
console.log("data", base64DataArr);
// save base64 data to file and return the url
for (let i = 0; i < formats.length; i++) {
const ext = formats[i];
const base64Data = base64DataArr[i];
if (ext && base64Data) {
if (ext === "png" && base64Data) {
const { filename } = this.saveToDisk(base64Data, ext);
output.push({
type: ext as InterpreterExtraType,
filename,
url: this.getFileUrl(filename),
});
} else {
output.push({
type: ext as InterpreterExtraType,
content: base64Data,
filename: `output.${ext}`,
url: "",
});
}
}
} catch (error) {
@@ -1,4 +1,5 @@
import os
from llama_index.vector_stores.milvus import MilvusVectorStore
@@ -15,6 +16,6 @@ def get_vector_store():
user=os.getenv("MILVUS_USERNAME"),
password=os.getenv("MILVUS_PASSWORD"),
collection_name=collection,
dim=int(os.getenv("EMBEDDING_DIM")),
dim=int(os.getenv("EMBEDDING_DIM", 768)),
)
return store
@@ -1,5 +1,6 @@
import os
import logging
import tempfile
from pydantic import BaseModel, Field, validator
from pydantic.alias_generators import to_camel
from typing import List, Any, Optional, Dict
@@ -21,6 +22,29 @@ class CsvFile(BaseModel):
filesize: int
id: str
type: str
local_file_path: Optional[str] = None
def __init__(self, **data):
super().__init__(**data)
# Write the content to a temporary file
saved_path = self.write_to_temp_file(self.content)
self.local_file_path = saved_path
@staticmethod
def write_to_temp_file(file_content: str) -> str:
"""
Write the content to a temporary file and return the file path
"""
csv_file = tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv")
csv_file.write(file_content)
file_path = csv_file.name
return file_path
def __del__(self):
# Remove the temporary file once the object is deleted
if self.local_file_path:
os.remove(self.local_file_path)
class DataParserOptions(BaseModel):
@@ -47,9 +71,16 @@ class DataParserOptions(BaseModel):
def to_raw_content(self) -> str:
if self.csv_files is not None and len(self.csv_files) > 0:
return "Use data from following CSV raw contents" + "\n".join(
[f"```csv\n{csv_file.content}\n```" for csv_file in self.csv_files]
)
saved_path = self.csv_files[0].local_file_path
saved_file_name = os.path.basename(saved_path)
content = self.csv_files[0].content
csv_meta = {
"local_file_path": saved_path,
"example_data": content[: min(200, len(content))],
"sandbox_file_path": f"/home/user/{saved_file_name}",
}
return f"Provided CSV file metadata:\n{csv_meta}"
def to_response_data(self) -> list[dict] | None:
output = []
@@ -21,7 +21,6 @@ STORAGE_DIR = os.getenv("STORAGE_DIR", "storage")
def get_doc_store():
# If the storage directory is there, load the document store from it.
# If not, set up an in-memory document store since we can't load from a directory that doesn't exist.
if os.path.exists(STORAGE_DIR):
@@ -14,11 +14,13 @@ import {
} from "llamaindex";
import { AgentStreamChatResponse } from "llamaindex/agent/base";
import path from "path";
import {
CsvFile,
appendCsvData,
appendImageData,
appendSourceData,
writeTempCsvFiles,
} from "./stream-helper";
type LlamaIndexResponse =
@@ -51,14 +53,22 @@ export const convertMessageContent = (
}
if (additionalData?.csvFiles?.length) {
const rawContents = additionalData.csvFiles.map((csv) => {
return "```csv\n" + csv.content + "\n```";
});
const tmpFile = writeTempCsvFiles(additionalData.csvFiles);
// Get a few lines of the CSV file as sample content
const sampleContent = additionalData.csvFiles
.map((csv) => csv.content.split("\n").slice(1, 4).join("\n"))
.join("\n\n");
const metadata = {
localFilePath: tmpFile.name,
sampleContent: sampleContent,
sandboxFilePath: `/home/user/${path.basename(tmpFile.name)}`,
};
const csvContent =
"Use data from following CSV raw contents:\n" + rawContents.join("\n\n");
"Provided CSV file metadata:\n" + JSON.stringify(metadata, null, 2);
console.log(csvContent);
content.push({
type: "text",
text: `${csvContent}\n\n${textMessage}`,
text: `${textMessage}\n\n${csvContent}`,
});
}
@@ -1,4 +1,5 @@
import { StreamData } from "ai";
import fs from "fs";
import {
CallbackManager,
Metadata,
@@ -6,6 +7,7 @@ import {
ToolCall,
ToolOutput,
} from "llamaindex";
import tmp from "tmp";
export function appendImageData(data: StreamData, imageUrl?: string) {
if (!imageUrl) return;
@@ -127,6 +129,7 @@ export type CsvFile = {
filename: string;
filesize: number;
id: string;
localFilePath: string;
};
export function appendCsvData(data: StreamData, csvFiles?: CsvFile[]) {
@@ -138,3 +141,10 @@ export function appendCsvData(data: StreamData, csvFiles?: CsvFile[]) {
},
});
}
export function writeTempCsvFiles(csvFiles: CsvFile[]) {
const csvFile = csvFiles[0];
const tmpFile = tmp.fileSync({ postfix: ".csv" });
fs.writeFileSync(tmpFile.name, csvFile.content);
return tmpFile;
}
@@ -36,7 +36,8 @@
"vaul": "^0.9.1",
"@llamaindex/pdf-viewer": "^1.1.1",
"@e2b/code-interpreter": "^0.0.5",
"uuid": "^9.0.1"
"uuid": "^9.0.1",
"tmp": "^0.2.3"
},
"devDependencies": {
"@types/node": "^20.10.3",