From d12da645a082a6249235d8db86d45a71611e5f43 Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Fri, 22 Nov 2024 12:10:36 -0800 Subject: [PATCH] feat: Add composio tools compatibility to sandboxes (#2097) --- .github/workflows/tests.yml | 1 + letta/functions/helpers.py | 4 -- letta/services/tool_execution_sandbox.py | 46 ++++++++++------------- letta/settings.py | 2 +- tests/test_tool_execution_sandbox.py | 47 ++++++++++++++++-------- 5 files changed, 52 insertions(+), 48 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 42ae9e6d..6f47b173 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -7,6 +7,7 @@ env: GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }} E2B_API_KEY: ${{ secrets.E2B_API_KEY }} + E2B_SANDBOX_TEMPLATE_ID: ${{ secrets.E2B_SANDBOX_TEMPLATE_ID }} on: push: diff --git a/letta/functions/helpers.py b/letta/functions/helpers.py index b73e4171..d94eb7da 100644 --- a/letta/functions/helpers.py +++ b/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} diff --git a/letta/services/tool_execution_sandbox.py b/letta/services/tool_execution_sandbox.py index a58d6dab..905e63e0 100644 --- a/letta/services/tool_execution_sandbox.py +++ b/letta/services/tool_execution_sandbox.py @@ -6,7 +6,6 @@ import pickle import runpy import sys import tempfile -import uuid from typing import Any, Optional from letta.log import get_logger @@ -24,10 +23,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 = "result_ZQqiequkcFwRwwGQMqkt" def __init__(self, tool_name: str, args: dict, user_id: str, force_recreate=False): self.tool_name = tool_name @@ -62,17 +60,17 @@ class ToolExecutionSandbox: """ 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,10 +106,11 @@ class ToolExecutionSandbox: 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 @@ -119,7 +118,7 @@ class ToolExecutionSandbox: 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 +138,6 @@ class ToolExecutionSandbox: 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 +162,7 @@ class ToolExecutionSandbox: 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 +222,13 @@ class ToolExecutionSandbox: 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 +264,15 @@ class ToolExecutionSandbox: # 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 diff --git a/letta/settings.py b/letta/settings.py index 1b443a4e..7271ff29 100644 --- a/letta/settings.py +++ b/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] = None # Updated manually class ModelSettings(BaseSettings): diff --git a/tests/test_tool_execution_sandbox.py b/tests/test_tool_execution_sandbox.py index f1d82b61..6bd10d8b 100644 --- a/tests/test_tool_execution_sandbox.py +++ b/tests/test_tool_execution_sandbox.py @@ -23,6 +23,7 @@ from letta.schemas.sandbox_config import ( 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"