142 lines
6.1 KiB
Python
142 lines
6.1 KiB
Python
import traceback
|
|
from typing import Any, Dict, Optional
|
|
|
|
from letta.functions.ast_parsers import coerce_dict_args_by_annotations, get_function_annotations_from_source
|
|
from letta.log import get_logger
|
|
from letta.otel.tracing import trace_method
|
|
from letta.schemas.agent import AgentState
|
|
from letta.schemas.enums import SandboxType, ToolSourceType
|
|
from letta.schemas.sandbox_config import SandboxConfig
|
|
from letta.schemas.tool import Tool
|
|
from letta.schemas.tool_execution_result import ToolExecutionResult
|
|
from letta.schemas.user import User
|
|
from letta.services.agent_manager import AgentManager
|
|
from letta.services.tool_executor.tool_executor_base import ToolExecutor
|
|
from letta.services.tool_sandbox.local_sandbox import AsyncToolSandboxLocal
|
|
from letta.settings import tool_settings
|
|
from letta.types import JsonDict
|
|
from letta.utils import get_friendly_error_msg
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class SandboxToolExecutor(ToolExecutor):
|
|
"""Executor for sandboxed tools."""
|
|
|
|
@trace_method
|
|
async def execute(
|
|
self,
|
|
function_name: str,
|
|
function_args: JsonDict,
|
|
tool: Tool,
|
|
actor: User,
|
|
agent_state: Optional[AgentState] = None,
|
|
sandbox_config: Optional[SandboxConfig] = None,
|
|
sandbox_env_vars: Optional[Dict[str, Any]] = None,
|
|
) -> ToolExecutionResult:
|
|
|
|
# Store original memory state
|
|
if agent_state:
|
|
orig_memory_str = await agent_state.memory.compile_in_thread_async()
|
|
else:
|
|
orig_memory_str = None
|
|
|
|
try:
|
|
# Prepare function arguments
|
|
function_args = self._prepare_function_args(function_args, tool, function_name)
|
|
|
|
agent_state_copy = self._create_agent_state_copy(agent_state) if agent_state else None
|
|
|
|
# Execute in sandbox depending on API key
|
|
if tool_settings.sandbox_type == SandboxType.E2B:
|
|
from letta.services.tool_sandbox.e2b_sandbox import AsyncToolSandboxE2B
|
|
|
|
sandbox = AsyncToolSandboxE2B(
|
|
function_name, function_args, actor, tool_object=tool, sandbox_config=sandbox_config, sandbox_env_vars=sandbox_env_vars
|
|
)
|
|
# TODO (cliandy): this is just for testing right now, separate this out into it's own subclass and handling logic
|
|
elif tool_settings.sandbox_type == SandboxType.MODAL:
|
|
from letta.services.tool_sandbox.modal_sandbox import AsyncToolSandboxModal, TypescriptToolSandboxModal
|
|
|
|
if tool.source_type == ToolSourceType.typescript:
|
|
sandbox = TypescriptToolSandboxModal(
|
|
function_name,
|
|
function_args,
|
|
actor,
|
|
tool_object=tool,
|
|
sandbox_config=sandbox_config,
|
|
sandbox_env_vars=sandbox_env_vars,
|
|
)
|
|
elif tool.source_type == ToolSourceType.python:
|
|
sandbox = AsyncToolSandboxModal(
|
|
function_name,
|
|
function_args,
|
|
actor,
|
|
tool_object=tool,
|
|
sandbox_config=sandbox_config,
|
|
sandbox_env_vars=sandbox_env_vars,
|
|
)
|
|
else:
|
|
raise ValueError(f"Tool source type was {tool.source_type} but is required to be python or typescript to run in Modal.")
|
|
else:
|
|
sandbox = AsyncToolSandboxLocal(
|
|
function_name, function_args, actor, tool_object=tool, sandbox_config=sandbox_config, sandbox_env_vars=sandbox_env_vars
|
|
)
|
|
|
|
tool_execution_result = await sandbox.run(agent_state=agent_state_copy)
|
|
|
|
log_lines = (tool_execution_result.stdout or []) + (tool_execution_result.stderr or [])
|
|
logger.debug("Tool execution log: %s", "\n".join(log_lines))
|
|
|
|
# Verify memory integrity
|
|
if agent_state:
|
|
new_memory_str = await agent_state.memory.compile_in_thread_async()
|
|
assert orig_memory_str == new_memory_str, "Memory should not be modified in a sandbox tool"
|
|
|
|
# Update agent memory if needed
|
|
if tool_execution_result.agent_state is not None:
|
|
await AgentManager().update_memory_if_changed_async(agent_state.id, tool_execution_result.agent_state.memory, actor)
|
|
|
|
return tool_execution_result
|
|
|
|
except Exception as e:
|
|
return self._handle_execution_error(e, function_name, traceback.format_exc())
|
|
|
|
@staticmethod
|
|
def _prepare_function_args(function_args: JsonDict, tool: Tool, function_name: str) -> dict:
|
|
"""Prepare function arguments with proper type coercion."""
|
|
try:
|
|
# Parse the source code to extract function annotations
|
|
annotations = get_function_annotations_from_source(tool.source_code, function_name)
|
|
# Coerce the function arguments to the correct types based on the annotations
|
|
return coerce_dict_args_by_annotations(function_args, annotations)
|
|
except ValueError:
|
|
# Just log the error and continue with original args
|
|
# This is defensive programming - we try to coerce but fall back if it fails
|
|
return function_args
|
|
|
|
@staticmethod
|
|
def _create_agent_state_copy(agent_state: AgentState):
|
|
"""Create a copy of agent state for sandbox execution."""
|
|
agent_state_copy = agent_state.__deepcopy__()
|
|
# Remove tools from copy to prevent nested tool execution
|
|
agent_state_copy.tools = []
|
|
agent_state_copy.tool_rules = []
|
|
return agent_state_copy
|
|
|
|
@staticmethod
|
|
def _handle_execution_error(
|
|
exception: Exception,
|
|
function_name: str,
|
|
stderr: str,
|
|
) -> ToolExecutionResult:
|
|
"""Handle tool execution errors."""
|
|
error_message = get_friendly_error_msg(
|
|
function_name=function_name, exception_name=type(exception).__name__, exception_message=str(exception)
|
|
)
|
|
return ToolExecutionResult(
|
|
status="error",
|
|
func_return=error_message,
|
|
stderr=[stderr],
|
|
)
|