letta-nightly 0.5.4.dev20241122104229__py3-none-any.whl → 0.5.4.dev20241124104049__py3-none-any.whl

letta/services/blocks_agents_manager.py

@@ -0,0 +1,84 @@
+import warnings
+from typing import List
+
+from letta.orm.blocks_agents import BlocksAgents as BlocksAgentsModel
+from letta.orm.errors import NoResultFound
+from letta.schemas.blocks_agents import BlocksAgents as PydanticBlocksAgents
+from letta.utils import enforce_types
+
+
+# TODO: DELETE THIS ASAP
+# TODO: So we have a patch where we manually specify CRUD operations
+# TODO: This is because Agent is NOT migrated to the ORM yet
+# TODO: Once we migrate Agent to the ORM, we should deprecate any agents relationship table managers
+class BlocksAgentsManager:
+    """Manager class to handle business logic related to Blocks and Agents."""
+
+    def __init__(self):
+        from letta.server.server import db_context
+
+        self.session_maker = db_context
+
+    @enforce_types
+    def add_block_to_agent(self, agent_id: str, block_id: str, block_label: str) -> PydanticBlocksAgents:
+        """Add a block to an agent. If the label already exists on that agent, this will error."""
+        with self.session_maker() as session:
+            try:
+                # Check if the block-label combination already exists for this agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                warnings.warn(f"Block label '{block_label}' already exists for agent '{agent_id}'.")
+            except NoResultFound:
+                blocks_agents_record = PydanticBlocksAgents(agent_id=agent_id, block_id=block_id, block_label=block_label)
+                blocks_agents_record = BlocksAgentsModel(**blocks_agents_record.model_dump(exclude_none=True))
+                blocks_agents_record.create(session)
+
+            return blocks_agents_record.to_pydantic()
+
+    @enforce_types
+    def remove_block_with_label_from_agent(self, agent_id: str, block_label: str) -> PydanticBlocksAgents:
+        """Remove a block with a label from an agent."""
+        with self.session_maker() as session:
+            try:
+                # Find and delete the block-label association for the agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                blocks_agents_record.hard_delete(session)
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block label '{block_label}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def remove_block_with_id_from_agent(self, agent_id: str, block_id: str) -> PydanticBlocksAgents:
+        """Remove a block with a label from an agent."""
+        with self.session_maker() as session:
+            try:
+                # Find and delete the block-label association for the agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_id=block_id)
+                blocks_agents_record.hard_delete(session)
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block id '{block_id}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def update_block_id_for_agent(self, agent_id: str, block_label: str, new_block_id: str) -> PydanticBlocksAgents:
+        """Update the block ID for a specific block label for an agent."""
+        with self.session_maker() as session:
+            try:
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                blocks_agents_record.block_id = new_block_id
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block label '{block_label}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def list_block_ids_for_agent(self, agent_id: str) -> List[str]:
+        """List all blocks associated with a specific agent."""
+        with self.session_maker() as session:
+            blocks_agents_record = BlocksAgentsModel.list(db_session=session, agent_id=agent_id)
+            return [record.block_id for record in blocks_agents_record]
+
+    @enforce_types
+    def list_agent_ids_with_block(self, block_id: str) -> List[str]:
+        """List all agents associated with a specific block."""
+        with self.session_maker() as session:
+            blocks_agents_record = BlocksAgentsModel.list(db_session=session, block_id=block_id)
+            return [record.agent_id for record in blocks_agents_record]

letta/services/sandbox_config_manager.py

@@ -0,0 +1,256 @@
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from letta.log import get_logger
+from letta.orm.errors import NoResultFound
+from letta.orm.sandbox_config import SandboxConfig as SandboxConfigModel
+from letta.orm.sandbox_config import SandboxEnvironmentVariable as SandboxEnvVarModel
+from letta.schemas.sandbox_config import E2BSandboxConfig, LocalSandboxConfig
+from letta.schemas.sandbox_config import SandboxConfig as PydanticSandboxConfig
+from letta.schemas.sandbox_config import SandboxConfigCreate, SandboxConfigUpdate
+from letta.schemas.sandbox_config import SandboxEnvironmentVariable as PydanticEnvVar
+from letta.schemas.sandbox_config import (
+    SandboxEnvironmentVariableCreate,
+    SandboxEnvironmentVariableUpdate,
+    SandboxType,
+)
+from letta.schemas.user import User as PydanticUser
+from letta.utils import enforce_types, printd
+
+logger = get_logger(__name__)
+
+
+class SandboxConfigManager:
+    """Manager class to handle business logic related to SandboxConfig and SandboxEnvironmentVariable."""
+
+    def __init__(self, settings):
+        from letta.server.server import db_context
+
+        self.session_maker = db_context
+        self.e2b_template_id = settings.e2b_sandbox_template_id
+
+    @enforce_types
+    def get_or_create_default_sandbox_config(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
+        sandbox_config = self.get_sandbox_config_by_type(sandbox_type, actor=actor)
+        if not sandbox_config:
+            logger.info(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.")
+
+            # TODO: Add more sandbox types later
+            if sandbox_type == SandboxType.E2B:
+                default_config = E2BSandboxConfig(template=self.e2b_template_id).model_dump(exclude_none=True)
+            else:
+                default_local_sandbox_path = str(Path(__file__).parent / "tool_sandbox_env")
+                default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path).model_dump(exclude_none=True)
+
+            sandbox_config = self.create_or_update_sandbox_config(SandboxConfigCreate(config=default_config), actor=actor)
+        return sandbox_config
+
+    @enforce_types
+    def create_or_update_sandbox_config(self, sandbox_config_create: SandboxConfigCreate, actor: PydanticUser) -> PydanticSandboxConfig:
+        """Create or update a sandbox configuration based on the PydanticSandboxConfig schema."""
+        config = sandbox_config_create.config
+        sandbox_type = config.type
+        sandbox_config = PydanticSandboxConfig(
+            type=sandbox_type, config=config.model_dump(exclude_none=True), organization_id=actor.organization_id
+        )
+
+        # Attempt to retrieve the existing sandbox configuration by type within the organization
+        db_sandbox = self.get_sandbox_config_by_type(sandbox_config.type, actor=actor)
+        if db_sandbox:
+            # Prepare the update data, excluding fields that should not be reset
+            update_data = sandbox_config.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(db_sandbox, key) != value}
+
+            # If there are changes, update the sandbox configuration
+            if update_data:
+                db_sandbox = self.update_sandbox_config(db_sandbox.id, SandboxConfigUpdate(**update_data), actor)
+            else:
+                printd(
+                    f"`create_or_update_sandbox_config` was called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"type={sandbox_config.type}, but found existing configuration with nothing to update."
+                )
+
+            return db_sandbox
+        else:
+            # If the sandbox configuration doesn't exist, create a new one
+            with self.session_maker() as session:
+                db_sandbox = SandboxConfigModel(**sandbox_config.model_dump(exclude_none=True))
+                db_sandbox.create(session, actor=actor)
+                return db_sandbox.to_pydantic()
+
+    @enforce_types
+    def update_sandbox_config(
+        self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
+    ) -> PydanticSandboxConfig:
+        """Update an existing sandbox configuration."""
+        with self.session_maker() as session:
+            sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
+            # We need to check that the sandbox_update provided is the same type as the original sandbox
+            if sandbox.type != sandbox_update.config.type:
+                raise ValueError(
+                    f"Mismatched type for sandbox config update: tried to update sandbox_config of type {sandbox.type} with config of type {sandbox_update.config.type}"
+                )
+
+            update_data = sandbox_update.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(sandbox, key) != value}
+
+            if update_data:
+                for key, value in update_data.items():
+                    setattr(sandbox, key, value)
+                sandbox.update(db_session=session, actor=actor)
+            else:
+                printd(
+                    f"`update_sandbox_config` called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"name={sandbox.type}, but nothing to update."
+                )
+            return sandbox.to_pydantic()
+
+    @enforce_types
+    def delete_sandbox_config(self, sandbox_config_id: str, actor: PydanticUser) -> PydanticSandboxConfig:
+        """Delete a sandbox configuration by its ID."""
+        with self.session_maker() as session:
+            sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
+            sandbox.hard_delete(db_session=session, actor=actor)
+            return sandbox.to_pydantic()
+
+    @enforce_types
+    def list_sandbox_configs(
+        self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50
+    ) -> List[PydanticSandboxConfig]:
+        """List all sandbox configurations with optional pagination."""
+        with self.session_maker() as session:
+            sandboxes = SandboxConfigModel.list(
+                db_session=session,
+                cursor=cursor,
+                limit=limit,
+                organization_id=actor.organization_id,
+            )
+            return [sandbox.to_pydantic() for sandbox in sandboxes]
+
+    @enforce_types
+    def get_sandbox_config_by_id(self, sandbox_config_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
+        """Retrieve a sandbox configuration by its ID."""
+        with self.session_maker() as session:
+            try:
+                sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
+                return sandbox.to_pydantic()
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    def get_sandbox_config_by_type(self, type: SandboxType, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
+        """Retrieve a sandbox config by its type."""
+        with self.session_maker() as session:
+            try:
+                sandboxes = SandboxConfigModel.list(
+                    db_session=session,
+                    type=type,
+                    organization_id=actor.organization_id,
+                    limit=1,
+                )
+                if sandboxes:
+                    return sandboxes[0].to_pydantic()
+                return None
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    def create_sandbox_env_var(
+        self, env_var_create: SandboxEnvironmentVariableCreate, sandbox_config_id: str, actor: PydanticUser
+    ) -> PydanticEnvVar:
+        """Create a new sandbox environment variable."""
+        env_var = PydanticEnvVar(**env_var_create.model_dump(), sandbox_config_id=sandbox_config_id, organization_id=actor.organization_id)
+
+        db_env_var = self.get_sandbox_env_var_by_key_and_sandbox_config_id(env_var.key, env_var.sandbox_config_id, actor=actor)
+        if db_env_var:
+            update_data = env_var.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(db_env_var, key) != value}
+            # If there are changes, update the environment variable
+            if update_data:
+                db_env_var = self.update_sandbox_env_var(db_env_var.id, SandboxEnvironmentVariableUpdate(**update_data), actor)
+            else:
+                printd(
+                    f"`create_or_update_sandbox_env_var` was called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"key={env_var.key}, but found existing variable with nothing to update."
+                )
+
+            return db_env_var
+        else:
+            with self.session_maker() as session:
+                env_var = SandboxEnvVarModel(**env_var.model_dump(exclude_none=True))
+                env_var.create(session, actor=actor)
+            return env_var.to_pydantic()
+
+    @enforce_types
+    def update_sandbox_env_var(
+        self, env_var_id: str, env_var_update: SandboxEnvironmentVariableUpdate, actor: PydanticUser
+    ) -> PydanticEnvVar:
+        """Update an existing sandbox environment variable."""
+        with self.session_maker() as session:
+            env_var = SandboxEnvVarModel.read(db_session=session, identifier=env_var_id, actor=actor)
+            update_data = env_var_update.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(env_var, key) != value}
+
+            if update_data:
+                for key, value in update_data.items():
+                    setattr(env_var, key, value)
+                env_var.update(db_session=session, actor=actor)
+            else:
+                printd(
+                    f"`update_sandbox_env_var` called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"key={env_var.key}, but nothing to update."
+                )
+            return env_var.to_pydantic()
+
+    @enforce_types
+    def delete_sandbox_env_var(self, env_var_id: str, actor: PydanticUser) -> PydanticEnvVar:
+        """Delete a sandbox environment variable by its ID."""
+        with self.session_maker() as session:
+            env_var = SandboxEnvVarModel.read(db_session=session, identifier=env_var_id, actor=actor)
+            env_var.hard_delete(db_session=session, actor=actor)
+            return env_var.to_pydantic()
+
+    @enforce_types
+    def list_sandbox_env_vars(
+        self, sandbox_config_id: str, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50
+    ) -> List[PydanticEnvVar]:
+        """List all sandbox environment variables with optional pagination."""
+        with self.session_maker() as session:
+            env_vars = SandboxEnvVarModel.list(
+                db_session=session,
+                cursor=cursor,
+                limit=limit,
+                organization_id=actor.organization_id,
+                sandbox_config_id=sandbox_config_id,
+            )
+            return [env_var.to_pydantic() for env_var in env_vars]
+
+    @enforce_types
+    def get_sandbox_env_vars_as_dict(
+        self, sandbox_config_id: str, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50
+    ) -> Dict[str, str]:
+        env_vars = self.list_sandbox_env_vars(sandbox_config_id, actor, cursor, limit)
+        result = {}
+        for env_var in env_vars:
+            result[env_var.key] = env_var.value
+        return result
+
+    @enforce_types
+    def get_sandbox_env_var_by_key_and_sandbox_config_id(
+        self, key: str, sandbox_config_id: str, actor: Optional[PydanticUser] = None
+    ) -> Optional[PydanticEnvVar]:
+        """Retrieve a sandbox environment variable by its key and sandbox_config_id."""
+        with self.session_maker() as session:
+            try:
+                env_var = SandboxEnvVarModel.list(
+                    db_session=session,
+                    key=key,
+                    sandbox_config_id=sandbox_config_id,
+                    organization_id=actor.organization_id,
+                    limit=1,
+                )
+                if env_var:
+                    return env_var[0].to_pydantic()
+                return None
+            except NoResultFound:
+                return None

letta/services/tool_execution_sandbox.py

@@ -0,0 +1,326 @@
+import ast
+import base64
+import io
+import os
+import pickle
+import runpy
+import sys
+import tempfile
+from typing import Any, Optional
+
+from letta.log import get_logger
+from letta.schemas.agent import AgentState
+from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult, SandboxType
+from letta.services.sandbox_config_manager import SandboxConfigManager
+from letta.services.tool_manager import ToolManager
+from letta.services.user_manager import UserManager
+from letta.settings import tool_settings
+
+logger = get_logger(__name__)
+
+
+class ToolExecutionSandbox:
+    METADATA_CONFIG_STATE_KEY = "config_state"
+    REQUIREMENT_TXT_NAME = "requirements.txt"
+
+    # 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
+        self.args = args
+
+        # Get the user
+        # This user corresponds to the agent_state's user_id field
+        # agent_state is the state of the agent that invoked this run
+        self.user = UserManager().get_user_by_id(user_id=user_id)
+
+        # Get the tool
+        # TODO: So in theory, it's possible this retrieves a tool not provisioned to the agent
+        # TODO: That would probably imply that agent_state is incorrectly configured
+        self.tool = ToolManager().get_tool_by_name(tool_name=tool_name, actor=self.user)
+        if not self.tool:
+            raise ValueError(
+                f"Agent attempted to invoke tool {self.tool_name} that does not exist for organization {self.user.organization_id}"
+            )
+
+        self.sandbox_config_manager = SandboxConfigManager(tool_settings)
+        self.force_recreate = force_recreate
+
+    def run(self, agent_state: Optional[AgentState] = None) -> Optional[SandboxRunResult]:
+        """
+        Run the tool in a sandbox environment.
+
+        Args:
+            agent_state (Optional[AgentState]): The state of the agent invoking the tool
+
+        Returns:
+            Tuple[Any, Optional[AgentState]]: Tuple containing (tool_result, agent_state)
+        """
+        if tool_settings.e2b_api_key:
+            logger.info(f"Using e2b sandbox to execute {self.tool_name}")
+            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(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}")
+        logger.info(f"Ending stdout log from tool run.")
+
+        # Return result
+        return result
+
+    # local sandbox specific functions
+    from contextlib import contextmanager
+
+    @contextmanager
+    def temporary_env_vars(self, env_vars: dict):
+        original_env = os.environ.copy()  # Backup original environment variables
+        os.environ.update(env_vars)  # Update with the new variables
+        try:
+            yield
+        finally:
+            os.environ.clear()
+            os.environ.update(original_env)  # Restore original environment variables
+
+    def run_local_dir_sandbox(self, code: str) -> Optional[SandboxRunResult]:
+        sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.LOCAL, actor=self.user)
+        local_configs = sbx_config.get_local_config()
+
+        # Get environment variables for the sandbox
+        env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100)
+
+        # Safety checks
+        if not os.path.isdir(local_configs.sandbox_dir):
+            raise FileNotFoundError(f"Sandbox directory does not exist: {local_configs.sandbox_dir}")
+
+        # Write the code to a temp file in the sandbox_dir
+        with tempfile.NamedTemporaryFile(mode="w", dir=local_configs.sandbox_dir, suffix=".py", delete=False) as temp_file:
+            temp_file.write(code)
+            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()
+            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(self.LOCAL_SANDBOX_RESULT_VAR_NAME)
+            func_return, agent_state = self.parse_best_effort(func_result)
+
+            # Restore stdout and collect captured output
+            sys.stdout = old_stdout
+            stdout_output = captured_stdout.getvalue()
+
+            return SandboxRunResult(
+                func_return=func_return,
+                agent_state=agent_state,
+                stdout=[stdout_output],
+                sandbox_config_fingerprint=sbx_config.fingerprint(),
+            )
+        except Exception as e:
+            raise RuntimeError(f"Executing tool {self.tool_name} has an unexpected error: {e}")
+        finally:
+            # Clean up the temp file and restore stdout
+            sys.stdout = old_stdout
+            os.remove(temp_file_path)
+
+    # e2b sandbox specific functions
+
+    def run_e2b_sandbox(self, code: str) -> Optional[SandboxRunResult]:
+        sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=self.user)
+        sbx = self.get_running_e2b_sandbox_with_same_state(sbx_config)
+        if not sbx or self.force_recreate:
+            sbx = self.create_e2b_sandbox_with_metadata_hash(sandbox_config=sbx_config)
+
+        # Since this sandbox was used, we extend its lifecycle by the timeout
+        sbx.set_timeout(sbx_config.get_e2b_config().timeout)
+
+        # Get environment variables for the sandbox
+        # TODO: We set limit to 100 here, but maybe we want it uncapped? Realistically this should be fine.
+        env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100)
+        execution = sbx.run_code(code, envs=env_vars)
+        if execution.error is not None:
+            raise Exception(f"Executing tool {self.tool_name} failed with {execution.error}. Generated code: \n\n{code}")
+        elif len(execution.results) == 0:
+            return None
+        else:
+            func_return, agent_state = self.parse_best_effort(execution.results[0].text)
+            return SandboxRunResult(
+                func_return=func_return,
+                agent_state=agent_state,
+                stdout=execution.logs.stdout + execution.logs.stderr,
+                sandbox_config_fingerprint=sbx_config.fingerprint(),
+            )
+
+    def get_running_e2b_sandbox_with_same_state(self, sandbox_config: SandboxConfig) -> Optional["Sandbox"]:
+        from e2b_code_interpreter import Sandbox
+
+        # List running sandboxes and access metadata.
+        running_sandboxes = self.list_running_e2b_sandboxes()
+
+        # Hash the config to check the state
+        state_hash = sandbox_config.fingerprint()
+        for sandbox in running_sandboxes:
+            if self.METADATA_CONFIG_STATE_KEY in sandbox.metadata and sandbox.metadata[self.METADATA_CONFIG_STATE_KEY] == state_hash:
+                return Sandbox.connect(sandbox.sandbox_id)
+
+        return None
+
+    def create_e2b_sandbox_with_metadata_hash(self, sandbox_config: SandboxConfig) -> "Sandbox":
+        from e2b_code_interpreter import Sandbox
+
+        state_hash = sandbox_config.fingerprint()
+        e2b_config = sandbox_config.get_e2b_config()
+        if e2b_config.template:
+            sbx = Sandbox(sandbox_config.get_e2b_config().template, metadata={self.METADATA_CONFIG_STATE_KEY: state_hash})
+        else:
+            # no template
+            sbx = Sandbox(metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(exclude={"pip_requirements"}))
+
+        # install pip requirements
+        if e2b_config.pip_requirements:
+            for package in e2b_config.pip_requirements:
+                sbx.commands.run(f"pip install {package}")
+        return sbx
+
+    def list_running_e2b_sandboxes(self):
+        from e2b_code_interpreter import Sandbox
+
+        # List running sandboxes and access metadata.
+        return Sandbox.list()
+
+    # general utility functions
+
+    def parse_best_effort(self, text: str) -> Any:
+        result = pickle.loads(base64.b64decode(text))
+        agent_state = None
+        if not result["agent_state"] is None:
+            agent_state = result["agent_state"]
+        return result["results"], agent_state
+
+    def parse_function_arguments(self, source_code: str, tool_name: str):
+        """Get arguments of a function from its source code"""
+        tree = ast.parse(source_code)
+        args = []
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef) and node.name == tool_name:
+                for arg in node.args.args:
+                    args.append(arg.arg)
+        return args
+
+    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
+
+        Returns:
+            code (str): The generated code strong
+        """
+        # dump JSON representation of agent state to re-load
+        code = "from typing import *\n"
+        code += "import pickle\n"
+        code += "import sys\n"
+        code += "import base64\n"
+
+        # Load the agent state data into the program
+        if agent_state:
+            code += "import letta\n"
+            code += "from letta import * \n"
+            import pickle
+
+            agent_state_pickle = pickle.dumps(agent_state)
+            code += f"agent_state = pickle.loads({agent_state_pickle})\n"
+        else:
+            # agent state is None
+            code += "agent_state = None\n"
+
+        for param in self.args:
+            code += self.initialize_param(param, self.args[param])
+
+        if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
+            inject_agent_state = True
+        else:
+            inject_agent_state = False
+
+        code += "\n" + self.tool.source_code + "\n"
+
+        # TODO: handle wrapped print
+
+        code += (
+            self.LOCAL_SANDBOX_RESULT_VAR_NAME
+            + ' = {"results": '
+            + self.invoke_function_call(inject_agent_state=inject_agent_state)
+            + ', "agent_state": agent_state}\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
+
+    def initialize_param(self, name: str, raw_value: str) -> str:
+        params = self.tool.json_schema["parameters"]["properties"]
+        spec = params.get(name)
+        if spec is None:
+            # ignore extra params (like 'self') for now
+            return ""
+
+        param_type = spec.get("type")
+        if param_type is None and spec.get("parameters"):
+            param_type = spec["parameters"].get("type")
+
+        if param_type == "string":
+            value = '"' + raw_value + '"'
+        elif param_type == "integer" or param_type == "boolean":
+            value = raw_value
+        else:
+            raise TypeError(f"unsupported type: {param_type}")
+
+        return name + " = " + str(value) + "\n"
+
+    def invoke_function_call(self, inject_agent_state: bool) -> str:
+        """
+        Generate the code string to call the function.
+
+        Args:
+            inject_agent_state (bool): Whether to inject the agent's state as an input into the tool
+
+        Returns:
+            str: Generated code string for calling the tool
+        """
+        kwargs = []
+        for name in self.args:
+            if name in self.tool.json_schema["parameters"]["properties"]:
+                kwargs.append(name)
+
+        param_list = [f"{arg}={arg}" for arg in kwargs]
+        if inject_agent_state:
+            param_list.append("agent_state=agent_state")
+        params = ", ".join(param_list)
+        # if "agent_state" in kwargs:
+        #    params += ", agent_state=agent_state"
+        # TODO: fix to figure out when to insert agent state or not
+        # params += "agent_state=agent_state"
+
+        func_call_str = self.tool.name + "(" + params + ")"
+        return func_call_str
+
+    #