Skip to content
Snippets Groups Projects
Unverified Commit fae9bcd6 authored by Huu Le's avatar Huu Le Committed by GitHub
Browse files

add raw text e2b tool output response (#115)

parent 2091fea2
No related branches found
No related tags found
No related merge requests found
......@@ -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,8 +14,9 @@ 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):
......@@ -72,19 +73,30 @@ class E2BCodeInterpreter:
try:
formats = result.formats()
base64_data_arr = [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)
results = [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,
)
)
)
except Exception as error:
logger.error("Error when saving data to disk", error)
logger.exception(error, exc_info=True)
logger.error("Error when parsing output from E2b interpreter tool", error)
return output
......
......@@ -34,8 +34,9 @@ type InterpreterExtraType =
export type InterpreterExtraResult = {
type: InterpreterExtraType;
filename: string;
url: string;
content?: string;
filename?: string;
url?: string;
};
const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<InterpreterParameter>> = {
......@@ -118,23 +119,34 @@ 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
const results = 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 base64Data = base64DataArr[i];
if (ext && base64Data) {
const { filename } = this.saveToDisk(base64Data, ext);
output.push({
type: ext as InterpreterExtraType,
filename,
url: this.getFileUrl(filename),
});
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;
}
}
} catch (error) {
console.error("Error when saving data to disk", error);
console.error("Error when parsing e2b response", error);
}
return output;
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment