letta-nightly 0.7.0.dev20250423003112__py3-none-any.whl → 0.7.1.dev20250423104245__py3-none-any.whl

letta/server/server.py

@@ -28,7 +28,6 @@ from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerCo
 from letta.groups.helpers import load_multi_agent
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
-from letta.helpers.message_helper import prepare_input_message_create
 
 # TODO use custom interface
 from letta.interface import AgentInterface  # abstract
@@ -148,7 +147,7 @@ class Server(object):
         raise NotImplementedError
 
     @abstractmethod
-    def send_messages(self, user_id: str, agent_id: str, messages: Union[MessageCreate, List[Message]]) -> None:
+    def send_messages(self, user_id: str, agent_id: str, input_messages: List[MessageCreate]) -> None:
         """Send a list of messages to the agent"""
         raise NotImplementedError
 
@@ -372,19 +371,13 @@ class SyncServer(Server):
         self,
         actor: User,
         agent_id: str,
-        input_messages: Union[Message, List[Message]],
+        input_messages: List[MessageCreate],
         interface: Union[AgentInterface, None] = None,  # needed to getting responses
         put_inner_thoughts_first: bool = True,
         # timestamp: Optional[datetime],
     ) -> LettaUsageStatistics:
         """Send the input message through the agent"""
         # TODO: Thread actor directly through this function, since the top level caller most likely already retrieved the user
-        # Input validation
-        if isinstance(input_messages, Message):
-            input_messages = [input_messages]
-        if not all(isinstance(m, Message) for m in input_messages):
-            raise ValueError(f"messages should be a Message or a list of Message, got {type(input_messages)}")
-
         logger.debug(f"Got input messages: {input_messages}")
         letta_agent = None
         try:
@@ -400,8 +393,9 @@ class SyncServer(Server):
                 metadata = interface.metadata if hasattr(interface, "metadata") else None
             else:
                 metadata = None
+
             usage_stats = letta_agent.step(
-                messages=input_messages,
+                input_messages=input_messages,
                 chaining=self.chaining,
                 max_chaining_steps=self.max_chaining_steps,
                 stream=token_streaming,
@@ -572,23 +566,14 @@ class SyncServer(Server):
             )
 
             # NOTE: eventually deprecate and only allow passing Message types
-            # Convert to a Message object
-            if timestamp:
-                message = Message(
-                    agent_id=agent_id,
-                    role="user",
-                    content=[TextContent(text=packaged_user_message)],
-                    created_at=timestamp,
-                )
-            else:
-                message = Message(
-                    agent_id=agent_id,
-                    role="user",
-                    content=[TextContent(text=packaged_user_message)],
-                )
+            message = MessageCreate(
+                agent_id=agent_id,
+                role="user",
+                content=[TextContent(text=packaged_user_message)],
+            )
 
         # Run the agent state forward
-        usage = self._step(actor=actor, agent_id=agent_id, input_messages=message)
+        usage = self._step(actor=actor, agent_id=agent_id, input_messages=[message])
         return usage
 
     def system_message(
@@ -660,23 +645,14 @@ class SyncServer(Server):
         self,
         actor: User,
         agent_id: str,
-        messages: Union[List[MessageCreate], List[Message]],
+        input_messages: List[MessageCreate],
         wrap_user_message: bool = True,
         wrap_system_message: bool = True,
         interface: Union[AgentInterface, ChatCompletionsStreamingInterface, None] = None,  # needed for responses
         metadata: Optional[dict] = None,  # Pass through metadata to interface
         put_inner_thoughts_first: bool = True,
     ) -> LettaUsageStatistics:
-        """Send a list of messages to the agent.
-
-        If messages are of type MessageCreate, convert them to Message objects before sending.
-        """
-        if all(isinstance(m, MessageCreate) for m in messages):
-            message_objects = [prepare_input_message_create(m, agent_id, wrap_user_message, wrap_system_message) for m in messages]
-        elif all(isinstance(m, Message) for m in messages):
-            message_objects = messages
-        else:
-            raise ValueError(f"All messages must be of type Message or MessageCreate, got {[type(m) for m in messages]}")
+        """Send a list of messages to the agent."""
 
         # Store metadata in interface if provided
         if metadata and hasattr(interface, "metadata"):
@@ -686,7 +662,7 @@ class SyncServer(Server):
         return self._step(
             actor=actor,
             agent_id=agent_id,
-            input_messages=message_objects,
+            input_messages=input_messages,
             interface=interface,
             put_inner_thoughts_first=put_inner_thoughts_first,
         )
@@ -703,8 +679,6 @@ class SyncServer(Server):
     @trace_method
     def get_cached_llm_config(self, **kwargs):
         key = make_key(**kwargs)
-        print(self._llm_config_cache)
-        print("KEY", key)
         if key not in self._llm_config_cache:
             self._llm_config_cache[key] = self.get_llm_config_from_handle(**kwargs)
         return self._llm_config_cache[key]
@@ -1019,12 +993,8 @@ class SyncServer(Server):
         agent = self.load_agent(agent_id=sleeptime_agent.id, actor=actor)
         for passage in self.list_data_source_passages(source_id=source.id, user_id=actor.id):
             agent.step(
-                messages=[
-                    Message(
-                        role="user",
-                        content=[TextContent(text=passage.text)],
-                        agent_id=sleeptime_agent.id,
-                    ),
+                input_messages=[
+                    MessageCreate(role="user", content=passage.text),
                 ]
             )
         self.agent_manager.delete_agent(agent_id=sleeptime_agent.id, actor=actor)
@@ -1182,7 +1152,6 @@ class SyncServer(Server):
             provider = self.get_provider_from_name(provider_name)
 
             llm_configs = [config for config in provider.list_llm_models() if config.handle == handle]
-            print("LLM CONFIGS", llm_configs)
             if not llm_configs:
                 llm_configs = [config for config in provider.list_llm_models() if config.model == model_name]
             if not llm_configs:
@@ -1195,8 +1164,6 @@ class SyncServer(Server):
             if not llm_configs:
                 raise e
 
-        print("CONFIGS", llm_configs)
-
         if len(llm_configs) == 1:
             llm_config = llm_configs[0]
         elif len(llm_configs) > 1:
@@ -1343,17 +1310,17 @@ class SyncServer(Server):
 
         # Next, attempt to run the tool with the sandbox
         try:
-            sandbox_run_result = ToolExecutionSandbox(tool.name, tool_args, actor, tool_object=tool).run(
+            tool_execution_result = ToolExecutionSandbox(tool.name, tool_args, actor, tool_object=tool).run(
                 agent_state=agent_state, additional_env_vars=tool_env_vars
             )
             return ToolReturnMessage(
                 id="null",
                 tool_call_id="null",
                 date=get_utc_time(),
-                status=sandbox_run_result.status,
-                tool_return=str(sandbox_run_result.func_return),
-                stdout=sandbox_run_result.stdout,
-                stderr=sandbox_run_result.stderr,
+                status=tool_execution_result.status,
+                tool_return=str(tool_execution_result.func_return),
+                stdout=tool_execution_result.stdout,
+                stderr=tool_execution_result.stderr,
             )
 
         except Exception as e:
@@ -1567,7 +1534,7 @@ class SyncServer(Server):
         agent_id: str,
         actor: User,
         # role: MessageRole,
-        messages: Union[List[Message], List[MessageCreate]],
+        input_messages: List[MessageCreate],
         stream_steps: bool,
         stream_tokens: bool,
         # related to whether or not we return `LettaMessage`s or `Message`s
@@ -1647,7 +1614,7 @@ class SyncServer(Server):
                     self.send_messages,
                     actor=actor,
                     agent_id=agent_id,
-                    messages=messages,
+                    input_messages=input_messages,
                     interface=streaming_interface,
                     metadata=metadata,
                 )
@@ -1701,7 +1668,7 @@ class SyncServer(Server):
         self,
         group_id: str,
         actor: User,
-        messages: Union[List[Message], List[MessageCreate]],
+        input_messages: Union[List[Message], List[MessageCreate]],
         stream_steps: bool,
         stream_tokens: bool,
         chat_completion_mode: bool = False,
@@ -1751,7 +1718,7 @@ class SyncServer(Server):
         task = asyncio.create_task(
             asyncio.to_thread(
                 letta_multi_agent.step,
-                messages=messages,
+                input_messages=input_messages,
                 chaining=self.chaining,
                 max_chaining_steps=self.max_chaining_steps,
             )

letta/services/agent_manager.py

@@ -364,6 +364,7 @@ class AgentManager:
                 "base_template_id": agent_update.base_template_id,
                 "message_buffer_autoclear": agent_update.message_buffer_autoclear,
                 "enable_sleeptime": agent_update.enable_sleeptime,
+                "response_format": agent_update.response_format,
             }
             for col, val in scalar_updates.items():
                 if val is not None:

letta/services/llm_batch_manager.py

@@ -291,9 +291,7 @@ class LLMBatchManager:
             return [item.to_pydantic() for item in results]
 
     def bulk_update_llm_batch_items(
-        self,
-        llm_batch_id_agent_id_pairs: List[Tuple[str, str]],
-        field_updates: List[Dict[str, Any]],
+        self, llm_batch_id_agent_id_pairs: List[Tuple[str, str]], field_updates: List[Dict[str, Any]], strict: bool = True
     ) -> None:
         """
         Efficiently update multiple LLMBatchItem rows by (llm_batch_id, agent_id) pairs.
@@ -301,30 +299,43 @@ class LLMBatchManager:
         Args:
             llm_batch_id_agent_id_pairs: List of (llm_batch_id, agent_id) tuples identifying items to update
             field_updates: List of dictionaries containing the fields to update for each item
+            strict: Whether to error if any of the requested keys don't exist (default True).
+                    If False, missing pairs are skipped.
         """
         if not llm_batch_id_agent_id_pairs or not field_updates:
             return
 
         if len(llm_batch_id_agent_id_pairs) != len(field_updates):
-            raise ValueError("batch_id_agent_id_pairs and field_updates must have the same length")
+            raise ValueError("llm_batch_id_agent_id_pairs and field_updates must have the same length")
 
         with self.session_maker() as session:
-            # Lookup primary keys
+            # Lookup primary keys for all requested (batch_id, agent_id) pairs
             items = (
                 session.query(LLMBatchItem.id, LLMBatchItem.llm_batch_id, LLMBatchItem.agent_id)
                 .filter(tuple_(LLMBatchItem.llm_batch_id, LLMBatchItem.agent_id).in_(llm_batch_id_agent_id_pairs))
                 .all()
             )
-            pair_to_pk = {(b, a): id for id, b, a in items}
-
+            pair_to_pk = {(batch_id, agent_id): pk for pk, batch_id, agent_id in items}
+
+            if strict:
+                requested = set(llm_batch_id_agent_id_pairs)
+                found = set(pair_to_pk.keys())
+                missing = requested - found
+                if missing:
+                    raise ValueError(
+                        f"Cannot bulk-update batch items: no records for the following " f"(llm_batch_id, agent_id) pairs: {missing}"
+                    )
+
+            # Build mappings, skipping any missing when strict=False
             mappings = []
-            for (llm_batch_id, agent_id), fields in zip(llm_batch_id_agent_id_pairs, field_updates):
-                pk_id = pair_to_pk.get((llm_batch_id, agent_id))
-                if not pk_id:
+            for (batch_id, agent_id), fields in zip(llm_batch_id_agent_id_pairs, field_updates):
+                pk = pair_to_pk.get((batch_id, agent_id))
+                if pk is None:
+                    # skip missing in non-strict mode
                     continue
 
                 update_fields = fields.copy()
-                update_fields["id"] = pk_id
+                update_fields["id"] = pk
                 mappings.append(update_fields)
 
             if mappings:
@@ -332,10 +343,7 @@ class LLMBatchManager:
                 session.commit()
 
     @enforce_types
-    def bulk_update_batch_llm_items_results_by_agent(
-        self,
-        updates: List[ItemUpdateInfo],
-    ) -> None:
+    def bulk_update_batch_llm_items_results_by_agent(self, updates: List[ItemUpdateInfo], strict: bool = True) -> None:
         """Update request status and batch results for multiple batch items."""
         batch_id_agent_id_pairs = [(update.llm_batch_id, update.agent_id) for update in updates]
         field_updates = [
@@ -346,29 +354,23 @@ class LLMBatchManager:
             for update in updates
         ]
 
-        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates)
+        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
-    def bulk_update_llm_batch_items_step_status_by_agent(
-        self,
-        updates: List[StepStatusUpdateInfo],
-    ) -> None:
+    def bulk_update_llm_batch_items_step_status_by_agent(self, updates: List[StepStatusUpdateInfo], strict: bool = True) -> None:
         """Update step status for multiple batch items."""
         batch_id_agent_id_pairs = [(update.llm_batch_id, update.agent_id) for update in updates]
         field_updates = [{"step_status": update.step_status} for update in updates]
 
-        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates)
+        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
-    def bulk_update_llm_batch_items_request_status_by_agent(
-        self,
-        updates: List[RequestStatusUpdateInfo],
-    ) -> None:
+    def bulk_update_llm_batch_items_request_status_by_agent(self, updates: List[RequestStatusUpdateInfo], strict: bool = True) -> None:
         """Update request status for multiple batch items."""
         batch_id_agent_id_pairs = [(update.llm_batch_id, update.agent_id) for update in updates]
         field_updates = [{"request_status": update.request_status} for update in updates]
 
-        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates)
+        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
     def delete_llm_batch_item(self, item_id: str, actor: PydanticUser) -> None:

letta/services/tool_executor/tool_execution_manager.py

@@ -1,16 +1,17 @@
-from typing import Any, Dict, Optional, Tuple, Type
+import traceback
+from typing import Any, Dict, Optional, Type
 
 from letta.log import get_logger
 from letta.orm.enums import ToolType
 from letta.schemas.agent import AgentState
-from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult
+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.tool_executor.tool_executor import (
     ExternalComposioToolExecutor,
     ExternalMCPToolExecutor,
     LettaCoreToolExecutor,
-    LettaMemoryToolExecutor,
     LettaMultiAgentToolExecutor,
     SandboxToolExecutor,
     ToolExecutor,
@@ -24,8 +25,9 @@ class ToolExecutorFactory:
 
     _executor_map: Dict[ToolType, Type[ToolExecutor]] = {
         ToolType.LETTA_CORE: LettaCoreToolExecutor,
+        ToolType.LETTA_MEMORY_CORE: LettaCoreToolExecutor,
+        ToolType.LETTA_SLEEPTIME_CORE: LettaCoreToolExecutor,
         ToolType.LETTA_MULTI_AGENT_CORE: LettaMultiAgentToolExecutor,
-        ToolType.LETTA_MEMORY_CORE: LettaMemoryToolExecutor,
         ToolType.EXTERNAL_COMPOSIO: ExternalComposioToolExecutor,
         ToolType.EXTERNAL_MCP: ExternalMCPToolExecutor,
     }
@@ -33,13 +35,8 @@ class ToolExecutorFactory:
     @classmethod
     def get_executor(cls, tool_type: ToolType) -> ToolExecutor:
         """Get the appropriate executor for the given tool type."""
-        executor_class = cls._executor_map.get(tool_type)
-
-        if executor_class:
-            return executor_class()
-
-        # Default to sandbox executor for unknown types
-        return SandboxToolExecutor()
+        executor_class = cls._executor_map.get(tool_type, SandboxToolExecutor)
+        return executor_class()
 
 
 class ToolExecutionManager:
@@ -58,7 +55,7 @@ class ToolExecutionManager:
         self.sandbox_config = sandbox_config
         self.sandbox_env_vars = sandbox_env_vars
 
-    def execute_tool(self, function_name: str, function_args: dict, tool: Tool) -> Tuple[Any, Optional[SandboxRunResult]]:
+    def execute_tool(self, function_name: str, function_args: dict, tool: Tool) -> ToolExecutionResult:
         """
         Execute a tool and persist any state changes.
 
@@ -71,35 +68,43 @@ class ToolExecutionManager:
             Tuple containing the function response and sandbox run result (if applicable)
         """
         try:
-            # Get the appropriate executor for this tool type
             executor = ToolExecutorFactory.get_executor(tool.tool_type)
-
-            # Execute the tool
             return executor.execute(
-                function_name, function_args, self.agent_state, tool, self.actor, self.sandbox_config, self.sandbox_env_vars
+                function_name,
+                function_args,
+                self.agent_state,
+                tool,
+                self.actor,
+                self.sandbox_config,
+                self.sandbox_env_vars,
             )
 
         except Exception as e:
             self.logger.error(f"Error executing tool {function_name}: {str(e)}")
-            error_message = get_friendly_error_msg(function_name=function_name, exception_name=type(e).__name__, exception_message=str(e))
-            return error_message, SandboxRunResult(status="error")
+            error_message = get_friendly_error_msg(
+                function_name=function_name,
+                exception_name=type(e).__name__,
+                exception_message=str(e),
+            )
+            return ToolExecutionResult(
+                status="error",
+                func_return=error_message,
+                stderr=[traceback.format_exc()],
+            )
 
     @trace_method
-    async def execute_tool_async(self, function_name: str, function_args: dict, tool: Tool) -> Tuple[Any, Optional[SandboxRunResult]]:
+    async def execute_tool_async(self, function_name: str, function_args: dict, tool: Tool) -> ToolExecutionResult:
         """
         Execute a tool asynchronously and persist any state changes.
         """
         try:
-            # Get the appropriate executor for this tool type
+            executor = ToolExecutorFactory.get_executor(tool.tool_type)
             # TODO: Extend this async model to composio
-
-            if tool.tool_type == ToolType.CUSTOM:
-                executor = SandboxToolExecutor()
-                result_tuple = await executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
+            if isinstance(executor, SandboxToolExecutor):
+                result = await executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
             else:
-                executor = ToolExecutorFactory.get_executor(tool.tool_type)
-                result_tuple = executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
-            return result_tuple
+                result = executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
+            return result
 
         except Exception as e:
             self.logger.error(f"Error executing tool {function_name}: {str(e)}")
@@ -108,4 +113,8 @@ class ToolExecutionManager:
                 exception_name=type(e).__name__,
                 exception_message=str(e),
             )
-            return error_message, SandboxRunResult(status="error")
+            return ToolExecutionResult(
+                status="error",
+                func_return=error_message,
+                stderr=[traceback.format_exc()],
+            )

letta/services/tool_executor/tool_execution_sandbox.py

@@ -13,8 +13,9 @@ from typing import Any, Dict, Optional
 from letta.functions.helpers import generate_model_from_args_json_schema
 from letta.log import get_logger
 from letta.schemas.agent import AgentState
-from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult, SandboxType
+from letta.schemas.sandbox_config import SandboxConfig, SandboxType
 from letta.schemas.tool import Tool
+from letta.schemas.tool_execution_result import ToolExecutionResult
 from letta.schemas.user import User
 from letta.services.helpers.tool_execution_helper import (
     add_imports_and_pydantic_schemas_for_args,
@@ -72,7 +73,11 @@ class ToolExecutionSandbox:
         self.force_recreate = force_recreate
         self.force_recreate_venv = force_recreate_venv
 
-    def run(self, agent_state: Optional[AgentState] = None, additional_env_vars: Optional[Dict] = None) -> SandboxRunResult:
+    def run(
+        self,
+        agent_state: Optional[AgentState] = None,
+        additional_env_vars: Optional[Dict] = None,
+    ) -> ToolExecutionResult:
         """
         Run the tool in a sandbox environment.
 
@@ -81,7 +86,7 @@ class ToolExecutionSandbox:
             additional_env_vars (Optional[Dict]): Environment variables to inject into the sandbox
 
         Returns:
-            Tuple[Any, Optional[AgentState]]: Tuple containing (tool_result, agent_state)
+            ToolExecutionResult: Object containing tool execution outcome (e.g. status, response)
         """
         if tool_settings.e2b_api_key and not self.privileged_tools:
             logger.debug(f"Using e2b sandbox to execute {self.tool_name}")
@@ -115,7 +120,7 @@ class ToolExecutionSandbox:
     @trace_method
     def run_local_dir_sandbox(
         self, agent_state: Optional[AgentState] = None, additional_env_vars: Optional[Dict] = None
-    ) -> SandboxRunResult:
+    ) -> ToolExecutionResult:
         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()
 
@@ -162,7 +167,12 @@ class ToolExecutionSandbox:
             os.remove(temp_file_path)
 
     @trace_method
-    def run_local_dir_sandbox_venv(self, sbx_config: SandboxConfig, env: Dict[str, str], temp_file_path: str) -> SandboxRunResult:
+    def run_local_dir_sandbox_venv(
+        self,
+        sbx_config: SandboxConfig,
+        env: Dict[str, str],
+        temp_file_path: str,
+    ) -> ToolExecutionResult:
         local_configs = sbx_config.get_local_config()
         sandbox_dir = os.path.expanduser(local_configs.sandbox_dir)  # Expand tilde
         venv_path = os.path.join(sandbox_dir, local_configs.venv_name)
@@ -205,12 +215,12 @@ class ToolExecutionSandbox:
             func_result, stdout = self.parse_out_function_results_markers(result.stdout)
             func_return, agent_state = self.parse_best_effort(func_result)
 
-            return SandboxRunResult(
+            return ToolExecutionResult(
+                status="success",
                 func_return=func_return,
                 agent_state=agent_state,
                 stdout=[stdout] if stdout else [],
                 stderr=[result.stderr] if result.stderr else [],
-                status="success",
                 sandbox_config_fingerprint=sbx_config.fingerprint(),
             )
 
@@ -221,12 +231,12 @@ class ToolExecutionSandbox:
                 exception_name=type(e).__name__,
                 exception_message=str(e),
             )
-            return SandboxRunResult(
+            return ToolExecutionResult(
+                status="error",
                 func_return=func_return,
                 agent_state=None,
                 stdout=[e.stdout] if e.stdout else [],
                 stderr=[e.stderr] if e.stderr else [],
-                status="error",
                 sandbox_config_fingerprint=sbx_config.fingerprint(),
             )
 
@@ -238,7 +248,12 @@ class ToolExecutionSandbox:
             raise e
 
     @trace_method
-    def run_local_dir_sandbox_directly(self, sbx_config: SandboxConfig, env: Dict[str, str], temp_file_path: str) -> SandboxRunResult:
+    def run_local_dir_sandbox_directly(
+        self,
+        sbx_config: SandboxConfig,
+        env: Dict[str, str],
+        temp_file_path: str,
+    ) -> ToolExecutionResult:
         status = "success"
         func_return, agent_state, stderr = None, None, None
 
@@ -288,12 +303,12 @@ class ToolExecutionSandbox:
         stdout_output = [captured_stdout.getvalue()] if captured_stdout.getvalue() else []
         stderr_output = [captured_stderr.getvalue()] if captured_stderr.getvalue() else []
 
-        return SandboxRunResult(
+        return ToolExecutionResult(
+            status=status,
             func_return=func_return,
             agent_state=agent_state,
             stdout=stdout_output,
             stderr=stderr_output,
-            status=status,
             sandbox_config_fingerprint=sbx_config.fingerprint(),
         )
 
@@ -307,7 +322,11 @@ class ToolExecutionSandbox:
 
     # e2b sandbox specific functions
 
-    def run_e2b_sandbox(self, agent_state: Optional[AgentState] = None, additional_env_vars: Optional[Dict] = None) -> SandboxRunResult:
+    def run_e2b_sandbox(
+        self,
+        agent_state: Optional[AgentState] = None,
+        additional_env_vars: Optional[Dict] = None,
+    ) -> ToolExecutionResult:
         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:
@@ -348,12 +367,12 @@ class ToolExecutionSandbox:
         else:
             raise ValueError(f"Tool {self.tool_name} returned execution with None")
 
-        return SandboxRunResult(
+        return ToolExecutionResult(
+            status="error" if execution.error else "success",
             func_return=func_return,
             agent_state=agent_state,
             stdout=execution.logs.stdout,
             stderr=execution.logs.stderr,
-            status="error" if execution.error else "success",
             sandbox_config_fingerprint=sbx_config.fingerprint(),
         )
 
@@ -535,7 +554,7 @@ class ToolExecutionSandbox:
         Generate the code string to call the function.
 
         Args:
-            inject_agent_state (bool): Whether to inject the axgent's state as an input into the tool
+            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