From 9c9decbb88bade76f7aa086a39dbac2dfdf595e0 Mon Sep 17 00:00:00 2001
From: Huu Le <39040748+leehuwuj@users.noreply.github.com>
Date: Fri, 14 Jun 2024 16:04:05 +0700
Subject: [PATCH] Reuse function tool instance and improve e2b interpreter tool
(#127)
---------
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>
---
.changeset/nervous-rats-lick.md | 5 ++
helpers/tools.ts | 13 ++-
.../engines/python/agent/tools/__init__.py | 11 +--
.../engines/python/agent/tools/interpreter.py | 82 +++++++++----------
.../python/agent/tools/openapi_action.py | 15 +++-
.../engines/python/agent/tools/weather.py | 3 +-
6 files changed, 62 insertions(+), 67 deletions(-)
create mode 100644 .changeset/nervous-rats-lick.md
diff --git a/.changeset/nervous-rats-lick.md b/.changeset/nervous-rats-lick.md
new file mode 100644
index 00000000..1e225ada
--- /dev/null
+++ b/.changeset/nervous-rats-lick.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Reuse function tool instances and improve e2b interpreter tool for Python
diff --git a/helpers/tools.ts b/helpers/tools.ts
index dcfa5d06..67d8f486 100644
--- a/helpers/tools.ts
+++ b/helpers/tools.ts
@@ -107,13 +107,12 @@ export const supportedTools: Tool[] = [
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for code interpreter tool.",
- 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.`,
+ 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.`,
},
],
},
diff --git a/templates/components/engines/python/agent/tools/__init__.py b/templates/components/engines/python/agent/tools/__init__.py
index f2c6c985..9e53efd7 100644
--- a/templates/components/engines/python/agent/tools/__init__.py
+++ b/templates/components/engines/python/agent/tools/__init__.py
@@ -19,15 +19,6 @@ class ToolFactory:
ToolType.LOCAL: "app.engine.tools",
}
- @staticmethod
- @cached(
- LRUCache(maxsize=100),
- key=lambda tool_type, tool_name, config: (
- tool_type,
- tool_name,
- json.dumps(config, sort_keys=True),
- ),
- )
def load_tools(tool_type: str, tool_name: str, config: dict) -> list[FunctionTool]:
source_package = ToolFactory.TOOL_SOURCE_PACKAGE_MAP[tool_type]
try:
@@ -40,7 +31,7 @@ class ToolFactory:
return tool_spec.to_tool_list()
else:
module = importlib.import_module(f"{source_package}.{tool_name}")
- tools = getattr(module, "tools")
+ tools = module.get_tools()
if not all(isinstance(tool, FunctionTool) for tool in tools):
raise ValueError(
f"The module {module} does not contain valid tools"
diff --git a/templates/components/engines/python/agent/tools/interpreter.py b/templates/components/engines/python/agent/tools/interpreter.py
index a89c6b9b..7184b7cb 100644
--- a/templates/components/engines/python/agent/tools/interpreter.py
+++ b/templates/components/engines/python/agent/tools/interpreter.py
@@ -29,9 +29,23 @@ class E2BCodeInterpreter:
output_dir = "tool-output"
- def __init__(self, api_key: str, filesever_url_prefix: str):
- self.api_key = api_key
+ def __init__(self):
+ 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"
+ )
+
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
@@ -101,50 +115,28 @@ class E2BCodeInterpreter:
return output
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}"
- )
- exec = interpreter.notebook.exec_cell(code)
+ """
+ Execute python code in a Jupyter notebook cell, the toll will return result, stdout, stderr, display_data, and error.
- 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=[])
- else:
- results = self.parse_result(exec.results[0])
- output = E2BToolOutput(
- is_error=False, logs=exec.logs, results=results
- )
- return output
-
-
-def code_interpret(code: str) -> Dict:
- """
- Execute python code in a Jupyter notebook cell and return any result, stdout, stderr, display_data, and error.
-
- Parameters:
- code (str): The python code to be executed in a single cell.
- """
- 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"
+ 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}"
)
-
- interpreter = E2BCodeInterpreter(
- api_key=api_key, filesever_url_prefix=filesever_url_prefix
- )
- output = interpreter.interpret(code)
- return output.dict()
+ 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=[])
+ else:
+ results = self.parse_result(exec.results[0])
+ output = E2BToolOutput(is_error=False, logs=exec.logs, results=results)
+ return output
-# Specify as functions tools to be loaded by the ToolFactory
-tools = [FunctionTool.from_defaults(code_interpret)]
+def get_tools():
+ return [FunctionTool.from_defaults(E2BCodeInterpreter().interpret)]
diff --git a/templates/components/engines/python/agent/tools/openapi_action.py b/templates/components/engines/python/agent/tools/openapi_action.py
index e9f19332..c19187d2 100644
--- a/templates/components/engines/python/agent/tools/openapi_action.py
+++ b/templates/components/engines/python/agent/tools/openapi_action.py
@@ -12,10 +12,17 @@ class OpenAPIActionToolSpec(OpenAPIToolSpec, RequestsToolSpec):
"""
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 = {}, **kwargs):
- # Load the OpenAPI spec
- openapi_spec, servers = self.load_openapi_spec(openapi_uri)
+ 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:
@@ -26,7 +33,7 @@ class OpenAPIActionToolSpec(OpenAPIToolSpec, RequestsToolSpec):
RequestsToolSpec.__init__(self, domain_headers)
@staticmethod
- def load_openapi_spec(uri: str) -> Tuple[Dict, List[str]]:
+ def _load_openapi_spec(uri: str) -> Tuple[Dict, List[str]]:
"""
Load an OpenAPI spec from a URI.
diff --git a/templates/components/engines/python/agent/tools/weather.py b/templates/components/engines/python/agent/tools/weather.py
index 3ea0fc03..08501310 100644
--- a/templates/components/engines/python/agent/tools/weather.py
+++ b/templates/components/engines/python/agent/tools/weather.py
@@ -69,4 +69,5 @@ class OpenMeteoWeather:
return response.json()
-tools = [FunctionTool.from_defaults(OpenMeteoWeather.get_weather_information)]
+def get_tools():
+ return [FunctionTool.from_defaults(OpenMeteoWeather.get_weather_information)]
--
GitLab