Reuse function tool instance and improve e2b interpreter tool (#127)

---------
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/nervous-rats-lick.md

+---
+"create-llama": patch
+---
+
+Reuse function tool instances and improve e2b interpreter tool for Python

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.`,
       },
     ],
   },

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"

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)]

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.
 

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)]