Skip to content

Commit

Permalink
finish
Browse files Browse the repository at this point in the history
  • Loading branch information
mattzh72 committed Nov 22, 2024
1 parent ae083fc commit 5f0aeaa
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 48 deletions.
4 changes: 0 additions & 4 deletions letta/functions/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}
Expand Down
48 changes: 21 additions & 27 deletions letta/services/tool_execution_sandbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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]:
Expand All @@ -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(),
)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion letta/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
47 changes: 31 additions & 16 deletions tests/test_tool_execution_sandbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
SandboxConfigCreate,
SandboxConfigUpdate,
SandboxEnvironmentVariableCreate,
SandboxType,
)
from letta.schemas.tool import Tool, ToolCreate
from letta.schemas.user import User
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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"

0 comments on commit 5f0aeaa

Please sign in to comment.