From d60b3c5a96d1f51f8a42bf20d71fa3b46982956b Mon Sep 17 00:00:00 2001
From: leehuwuj <leehuwuj@gmail.com>
Date: Mon, 27 May 2024 13:09:59 +0700
Subject: [PATCH] refactor code and add changeset
---
.changeset/itchy-ads-travel.md | 5 ++++
helpers/index.ts | 5 ++++
.../engines/python/agent/tools/interpreter.py | 20 +++++++-------
.../streaming/fastapi/app/api/routers/chat.py | 8 +++++-
.../fastapi/app/api/routers/messaging.py | 26 ++++++++++++-------
templates/types/streaming/fastapi/main.py | 4 +--
6 files changed, 46 insertions(+), 22 deletions(-)
create mode 100644 .changeset/itchy-ads-travel.md
diff --git a/.changeset/itchy-ads-travel.md b/.changeset/itchy-ads-travel.md
new file mode 100644
index 00000000..ac3ee4cd
--- /dev/null
+++ b/.changeset/itchy-ads-travel.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Add support E2B code interpreter tool for FastAPI
diff --git a/helpers/index.ts b/helpers/index.ts
index ef602d28..e65998ba 100644
--- a/helpers/index.ts
+++ b/helpers/index.ts
@@ -171,6 +171,11 @@ export const installTemplate = async (
);
}
}
+
+ // Create tool-output directory
+ if (props.tools && props.tools.length > 0) {
+ await fsExtra.mkdir(path.join(props.root, "tool-output"));
+ }
} else {
// this is a frontend for a full-stack app, create .env file with model information
await createFrontendEnvFile(props.root, {
diff --git a/templates/components/engines/python/agent/tools/interpreter.py b/templates/components/engines/python/agent/tools/interpreter.py
index 55716a76..bc913935 100644
--- a/templates/components/engines/python/agent/tools/interpreter.py
+++ b/templates/components/engines/python/agent/tools/interpreter.py
@@ -32,11 +32,6 @@ class E2BCodeInterpreter:
self.api_key = api_key
self.filesever_url_prefix = filesever_url_prefix
- def code_interpret(
- self, code_interpreter: CodeInterpreter, code: str
- ) -> Tuple[List, List]:
- pass
-
def get_output_path(self, filename: str) -> str:
# if output directory doesn't exist, create it
if not os.path.exists(self.output_dir):
@@ -48,8 +43,12 @@ class E2BCodeInterpreter:
buffer = base64.b64decode(base64_data)
output_path = self.get_output_path(filename)
- with open(output_path, "wb") as file:
- file.write(buffer)
+ try:
+ with open(output_path, "wb") as file:
+ file.write(buffer)
+ except IOError as e:
+ logger.error(f"Failed to write to file {output_path}: {str(e)}")
+ raise e
logger.info(f"Saved file to {output_path}")
@@ -89,7 +88,7 @@ class E2BCodeInterpreter:
return output
- def interpret(self, code: str) -> Dict:
+ def interpret(self, code: str) -> E2BToolOutput:
with CodeInterpreter(api_key=self.api_key) as interpreter:
logger.info(
f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}"
@@ -106,7 +105,7 @@ class E2BCodeInterpreter:
output = E2BToolOutput(
is_error=False, logs=exec.logs, results=results
)
- return output.dict()
+ return output
def code_interpret(code: str) -> Dict:
@@ -127,7 +126,8 @@ def code_interpret(code: str) -> Dict:
interpreter = E2BCodeInterpreter(
api_key=api_key, filesever_url_prefix=filesever_url_prefix
)
- return interpreter.interpret(code)
+ output = interpreter.interpret(code)
+ return output.dict()
# Specify as functions tools to be loaded by the ToolFactory
diff --git a/templates/types/streaming/fastapi/app/api/routers/chat.py b/templates/types/streaming/fastapi/app/api/routers/chat.py
index c92ca3d4..a23cc440 100644
--- a/templates/types/streaming/fastapi/app/api/routers/chat.py
+++ b/templates/types/streaming/fastapi/app/api/routers/chat.py
@@ -93,7 +93,13 @@ async def chat(
event_handler = EventCallbackHandler()
chat_engine.callback_manager.handlers.append(event_handler) # type: ignore
- response = await chat_engine.astream_chat(last_message_content, messages)
+ try:
+ response = await chat_engine.astream_chat(last_message_content, messages)
+ except Exception as e:
+ raise HTTPException(
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+ detail=f"Error in chat engine: {e}",
+ )
async def content_generator():
# Yield the text response
diff --git a/templates/types/streaming/fastapi/app/api/routers/messaging.py b/templates/types/streaming/fastapi/app/api/routers/messaging.py
index a239657a..94cc5851 100644
--- a/templates/types/streaming/fastapi/app/api/routers/messaging.py
+++ b/templates/types/streaming/fastapi/app/api/routers/messaging.py
@@ -1,5 +1,6 @@
import json
import asyncio
+import logging
from typing import AsyncGenerator, Dict, Any, List, Optional
from llama_index.core.callbacks.base import BaseCallbackHandler
from llama_index.core.callbacks.schema import CBEventType
@@ -7,6 +8,9 @@ from llama_index.core.tools.types import ToolOutput
from pydantic import BaseModel
+logger = logging.getLogger(__name__)
+
+
class CallbackEvent(BaseModel):
event_type: CBEventType
payload: Optional[Dict[str, Any]] = None
@@ -72,15 +76,19 @@ class CallbackEvent(BaseModel):
}
def to_response(self):
- match self.event_type:
- case "retrieve":
- return self.get_retrieval_message()
- case "function_call":
- return self.get_tool_message()
- case "agent_step":
- return self.get_agent_tool_response()
- case _:
- return None
+ try:
+ match self.event_type:
+ case "retrieve":
+ return self.get_retrieval_message()
+ case "function_call":
+ return self.get_tool_message()
+ case "agent_step":
+ return self.get_agent_tool_response()
+ case _:
+ return None
+ except Exception as e:
+ logger.error(f"Error in converting event to response: {e}")
+ return None
class EventCallbackHandler(BaseCallbackHandler):
diff --git a/templates/types/streaming/fastapi/main.py b/templates/types/streaming/fastapi/main.py
index a7569a52..b095b616 100644
--- a/templates/types/streaming/fastapi/main.py
+++ b/templates/types/streaming/fastapi/main.py
@@ -41,8 +41,8 @@ if environment == "dev":
# Mount the data files to serve the file viewer
if os.path.exists("data"):
app.mount("/api/files/data", StaticFiles(directory="data"), name="data-static")
-# Mount the tool output files
-if os.path.exists("config/tools.yaml"):
+# Mount the output files from tools
+if os.path.exists("tool-output"):
app.mount(
"/api/files/tool-output",
StaticFiles(directory="tool-output"),
--
GitLab