feat: Add composio tools compatibility to sandboxes

letta/functions/helpers.py

@@ -13,8 +13,6 @@ def generate_composio_tool_wrapper(action: "ActionType") -> tuple[str, str]:
 
     wrapper_function_str = f"""
 def {func_name}(**kwargs):
-    if 'self' in kwargs:
-        del kwargs['self']
     from composio import Action, App, Tag
     from composio_langchain import ComposioToolSet
 
@@ -46,8 +44,6 @@ def generate_langchain_tool_wrapper(
     # Combine all parts into the wrapper function
     wrapper_function_str = f"""
 def {func_name}(**kwargs):
-    if 'self' in kwargs:
-        del kwargs['self']
     import importlib
     {import_statement}
     {extra_module_imports}

letta/services/tool_execution_sandbox.py

@@ -3,10 +3,11 @@
 import io
 import os
 import pickle
+import random
 import runpy
+import string
 import sys
 import tempfile
-import uuid
 from typing import Any, Optional
 
 from letta.log import get_logger
@@ -24,10 +25,9 @@ class ToolExecutionSandbox:
     METADATA_CONFIG_STATE_KEY = "config_state"
     REQUIREMENT_TXT_NAME = "requirements.txt"
 
-    # For generating long, random marker hashes
-    NAMESPACE = uuid.NAMESPACE_DNS
-    LOCAL_SANDBOX_RESULT_START_MARKER = str(uuid.uuid5(NAMESPACE, "local-sandbox-result-start-marker"))
-    LOCAL_SANDBOX_RESULT_END_MARKER = str(uuid.uuid5(NAMESPACE, "local-sandbox-result-end-marker"))
+    # This is the variable name in the auto-generated code that contains the function results
+    # We make this a long random string to avoid collisions with any variables in the user's code
+    LOCAL_SANDBOX_RESULT_VAR_NAME = "".join(random.choices(string.ascii_letters + "_", k=20))
 
     def __init__(self, tool_name: str, args: dict, user_id: str, force_recreate=False):
         self.tool_name = tool_name
@@ -62,17 +62,17 @@ def run(self, agent_state: Optional[AgentState] = None) -> Optional[SandboxRunRe
         """
         if tool_settings.e2b_api_key:
             logger.info(f"Using e2b sandbox to execute {self.tool_name}")
-            code = self.generate_execution_script(wrap_print_with_markers=False, agent_state=agent_state)
+            code = self.generate_execution_script(agent_state=agent_state)
             result = self.run_e2b_sandbox(code=code)
         else:
             logger.info(f"Using local sandbox to execute {self.tool_name}")
-            code = self.generate_execution_script(wrap_print_with_markers=True, agent_state=agent_state)
+            code = self.generate_execution_script(agent_state=agent_state)
             result = self.run_local_dir_sandbox(code=code)
 
         # Log out any stdout from the tool run
         logger.info(f"Executed tool '{self.tool_name}', logging stdout from tool run: \n")
         for log_line in result.stdout:
-            logger.info(f"{log_line}\n")
+            logger.info(f"{log_line}")
         logger.info(f"Ending stdout log from tool run.")
 
         # Return result
@@ -108,18 +108,19 @@ def run_local_dir_sandbox(self, code: str) -> Optional[SandboxRunResult]:
             temp_file.flush()
             temp_file_path = temp_file.name
 
+        # Save the old stdout
+        old_stdout = sys.stdout
         try:
             # Redirect stdout to capture script output
             captured_stdout = io.StringIO()
-            old_stdout = sys.stdout
             sys.stdout = captured_stdout
 
             # Execute the temp file
             with self.temporary_env_vars(env_vars):
                 result = runpy.run_path(temp_file_path, init_globals=env_vars)
 
             # Fetch the result
-            func_result = result.get("result")
+            func_result = result.get(self.LOCAL_SANDBOX_RESULT_VAR_NAME)
             func_return, agent_state = self.parse_best_effort(func_result)
 
             # Restore stdout and collect captured output
@@ -139,12 +140,6 @@ def run_local_dir_sandbox(self, code: str) -> Optional[SandboxRunResult]:
             sys.stdout = old_stdout
             os.remove(temp_file_path)
 
-    def parse_out_function_results_markers(self, text: str):
-        marker_len = len(self.LOCAL_SANDBOX_RESULT_START_MARKER)
-        start_index = text.index(self.LOCAL_SANDBOX_RESULT_START_MARKER) + marker_len
-        end_index = text.index(self.LOCAL_SANDBOX_RESULT_END_MARKER)
-        return text[start_index:end_index], text[: start_index - marker_len] + text[end_index + +marker_len :]
-
     # e2b sandbox specific functions
 
     def run_e2b_sandbox(self, code: str) -> Optional[SandboxRunResult]:
@@ -169,7 +164,7 @@ def run_e2b_sandbox(self, code: str) -> Optional[SandboxRunResult]:
             return SandboxRunResult(
                 func_return=func_return,
                 agent_state=agent_state,
-                stdout=execution.logs.stdout,
+                stdout=execution.logs.stdout + execution.logs.stderr,
                 sandbox_config_fingerprint=sbx_config.fingerprint(),
             )
 
@@ -229,14 +224,13 @@ def parse_function_arguments(self, source_code: str, tool_name: str):
                     args.append(arg.arg)
         return args
 
-    def generate_execution_script(self, agent_state: AgentState, wrap_print_with_markers: bool = False) -> str:
+    def generate_execution_script(self, agent_state: AgentState) -> str:
         """
         Generate code to run inside of execution sandbox.
         Passes into a serialized agent state into the code, to be accessed by the tool.
 
         Args:
             agent_state (AgentState): The agent state
-            wrap_print_with_markers (bool): Whether to wrap print statements (?)
 
         Returns:
             code (str): The generated code strong
@@ -272,15 +266,15 @@ def generate_execution_script(self, agent_state: AgentState, wrap_print_with_mar
         # TODO: handle wrapped print
 
         code += (
-            'result = {"results": ' + self.invoke_function_call(inject_agent_state=inject_agent_state) + ', "agent_state": agent_state}\n'
+            self.LOCAL_SANDBOX_RESULT_VAR_NAME
+            + ' = {"results": '
+            + self.invoke_function_call(inject_agent_state=inject_agent_state)
+            + ', "agent_state": agent_state}\n'
         )
-        code += "result = base64.b64encode(pickle.dumps(result)).decode('utf-8')\n"
-        if wrap_print_with_markers:
-            code += f"sys.stdout.write('{self.LOCAL_SANDBOX_RESULT_START_MARKER}')\n"
-            code += f"sys.stdout.write(str(result))\n"
-            code += f"sys.stdout.write('{self.LOCAL_SANDBOX_RESULT_END_MARKER}')\n"
-        else:
-            code += "result\n"
+        code += (
+            f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME} = base64.b64encode(pickle.dumps({self.LOCAL_SANDBOX_RESULT_VAR_NAME})).decode('utf-8')\n"
+        )
+        code += f"{self.LOCAL_SANDBOX_RESULT_VAR_NAME}\n"
 
         return code
 

letta/settings.py

@@ -12,7 +12,7 @@ class ToolSettings(BaseSettings):
 
     # Sandbox configurations
     e2b_api_key: Optional[str] = None
-    e2b_sandbox_template_id: Optional[str] = "ngtrcfmr9wyzs9yjd8l2"  # Updated manually
+    e2b_sandbox_template_id: Optional[str] = "m2lu1nfx7ztyiuzlbl89"  # Updated manually
 
 
 class ModelSettings(BaseSettings):

tests/test_tool_execution_sandbox.py

@@ -23,6 +23,7 @@
     SandboxConfigCreate,
     SandboxConfigUpdate,
     SandboxEnvironmentVariableCreate,
+    SandboxType,
 )
 from letta.schemas.tool import Tool, ToolCreate
 from letta.schemas.user import User
@@ -275,6 +276,22 @@ def test_local_sandbox_env(mock_e2b_api_key_none, get_env_tool, test_user):
     assert long_random_string in result.func_return
 
 
+@pytest.mark.local_sandbox
+def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_composio_key_set, composio_github_star_tool, test_user):
+    # Add the composio key
+    manager = SandboxConfigManager(tool_settings)
+    config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.LOCAL, actor=test_user)
+
+    manager.create_sandbox_env_var(
+        SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=tool_settings.composio_api_key),
+        sandbox_config_id=config.id,
+        actor=test_user,
+    )
+
+    result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user_id=test_user.id).run()
+    assert result.func_return["details"] == "Action executed successfully"
+
+
 # E2B sandbox tests
 
 
@@ -407,19 +424,17 @@ def test_e2b_sandbox_with_list_rv(check_e2b_key_is_set, list_tool, test_user):
     assert len(result.func_return) == 5
 
 
-# TODO: Add tests for composio
-# def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_set, composio_github_star_tool, test_user):
-#     # Add the composio key
-#     manager = SandboxConfigManager(tool_settings)
-#     config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=test_user)
-#
-#     manager.create_sandbox_env_var(
-#         SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=tool_settings.composio_api_key),
-#         sandbox_config_id=config.id,
-#         actor=test_user,
-#     )
-#
-#     result = ToolExecutionSandbox(composio_github_star_tool.name, {}, user_id=test_user.id).run()
-#     import ipdb
-#
-#     ipdb.set_trace()
+@pytest.mark.e2b_sandboxfunc
+def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_set, composio_github_star_tool, test_user):
+    # Add the composio key
+    manager = SandboxConfigManager(tool_settings)
+    config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=test_user)
+
+    manager.create_sandbox_env_var(
+        SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=tool_settings.composio_api_key),
+        sandbox_config_id=config.id,
+        actor=test_user,
+    )
+
+    result = ToolExecutionSandbox(composio_github_star_tool.name, {"owner": "letta-ai", "repo": "letta"}, user_id=test_user.id).run()
+    assert result.func_return["details"] == "Action executed successfully"