diff --git a/.changeset/nervous-rats-lick.md b/.changeset/nervous-rats-lick.md
new file mode 100644
index 0000000000000000000000000000000000000000..1e225ada14b3d1422e54240575eb624fc1905320
--- /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 dcfa5d06ed20277ef82be6aaca33a15ab48877ad..67d8f486f36e45f37f69dbb8f34668b3e1283b9d 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 f2c6c98536206d7f30d818e562635be05c201900..9e53efd7d9b58b82573de3ea8092e4175cd64ac7 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 a89c6b9bbcc9e9b0f1d80de98324f5a80a34de00..7184b7cb7ab66c9314f5ad836129b2b06f798eb9 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 e9f1933224a713f9dea3ca3461c8894d02ea63a5..c19187d2f78d3431beb2ca53d63a0dba002722e6 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 3ea0fc03fcdfd48443366833dd506d5025c74a19..08501310a0fa7db1c9a14d49c9fc1c3b738f091c 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)]