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

letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client

letta/agent.py

@@ -3,7 +3,7 @@ import time
 import traceback
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from openai.types.beta.function_tool import FunctionTool as OpenAITool
 
@@ -17,6 +17,7 @@ from letta.constants import (
     LETTA_MULTI_AGENT_TOOL_MODULE_NAME,
     LLM_MAX_TOKENS,
     REQ_HEARTBEAT_MESSAGE,
+    SEND_MESSAGE_TOOL_NAME,
 )
 from letta.errors import ContextWindowExceededError
 from letta.functions.ast_parsers import coerce_dict_args_by_annotations, get_function_annotations_from_source
@@ -27,6 +28,7 @@ from letta.helpers import ToolRulesSolver
 from letta.helpers.composio_helpers import get_composio_api_key
 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
 from letta.interface import AgentInterface
 from letta.llm_api.helpers import calculate_summarizer_cutoff, get_token_counts_for_messages, is_context_overflow_error
 from letta.llm_api.llm_api_tools import create
@@ -42,12 +44,13 @@ from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import MessageRole
 from letta.schemas.letta_message_content import TextContent
 from letta.schemas.memory import ContextWindowOverview, Memory
-from letta.schemas.message import Message, ToolReturn
+from letta.schemas.message import Message, MessageCreate, ToolReturn
 from letta.schemas.openai.chat_completion_response import ChatCompletionResponse
 from letta.schemas.openai.chat_completion_response import Message as ChatCompletionMessage
 from letta.schemas.openai.chat_completion_response import UsageStatistics
-from letta.schemas.sandbox_config import SandboxRunResult
+from letta.schemas.response_format import ResponseFormatType
 from letta.schemas.tool import Tool
+from letta.schemas.tool_execution_result import ToolExecutionResult
 from letta.schemas.tool_rule import TerminalToolRule
 from letta.schemas.usage import LettaUsageStatistics
 from letta.services.agent_manager import AgentManager
@@ -78,7 +81,7 @@ class BaseAgent(ABC):
     @abstractmethod
     def step(
         self,
-        messages: Union[Message, List[Message]],
+        input_messages: List[MessageCreate],
     ) -> LettaUsageStatistics:
         """
         Top-level event message handler for the agent.
@@ -255,6 +258,28 @@ class Agent(BaseAgent):
         # Return updated messages
         return messages
 
+    def _runtime_override_tool_json_schema(
+        self,
+        functions_list: List[Dict | None],
+    ) -> List[Dict | None]:
+        """Override the tool JSON schema at runtime for a particular tool if conditions are met."""
+
+        # Currently just injects `send_message` with a `response_format` if provided to the agent.
+        if self.agent_state.response_format and self.agent_state.response_format.type != ResponseFormatType.text:
+            for func in functions_list:
+                if func["name"] == SEND_MESSAGE_TOOL_NAME:
+                    if self.agent_state.response_format.type == ResponseFormatType.json_schema:
+                        func["parameters"]["properties"]["message"] = self.agent_state.response_format.json_schema["schema"]
+                    if self.agent_state.response_format.type == ResponseFormatType.json_object:
+                        func["parameters"]["properties"]["message"] = {
+                            "type": "object",
+                            "description": "Message contents. All unicode (including emojis) are supported.",
+                            "additionalProperties": True,
+                            "properties": {},
+                        }
+                    break
+        return functions_list
+
     @trace_method
     def _get_ai_reply(
         self,
@@ -268,27 +293,26 @@ class Agent(BaseAgent):
         step_count: Optional[int] = None,
         last_function_failed: bool = False,
         put_inner_thoughts_first: bool = True,
-    ) -> ChatCompletionResponse:
+    ) -> ChatCompletionResponse | None:
         """Get response from LLM API with robust retry mechanism."""
         log_telemetry(self.logger, "_get_ai_reply start")
         available_tools = set([t.name for t in self.agent_state.tools])
-        allowed_tool_names = self.tool_rules_solver.get_allowed_tool_names(
-            available_tools=available_tools, last_function_response=self.last_function_response
-        )
         agent_state_tool_jsons = [t.json_schema for t in self.agent_state.tools]
 
-        allowed_functions = (
-            agent_state_tool_jsons
-            if not allowed_tool_names
-            else [func for func in agent_state_tool_jsons if func["name"] in allowed_tool_names]
-        )
+        # Get allowed tools or allow all if none are allowed
+        allowed_tool_names = self.tool_rules_solver.get_allowed_tool_names(
+            available_tools=available_tools, last_function_response=self.last_function_response
+        ) or list(available_tools)
 
         # Don't allow a tool to be called if it failed last time
         if last_function_failed and self.tool_rules_solver.tool_call_history:
-            allowed_functions = [f for f in allowed_functions if f["name"] != self.tool_rules_solver.tool_call_history[-1]]
-            if not allowed_functions:
+            allowed_tool_names = [f for f in allowed_tool_names if f != self.tool_rules_solver.tool_call_history[-1]]
+            if not allowed_tool_names:
                 return None
 
+        allowed_functions = [func for func in agent_state_tool_jsons if func["name"] in allowed_tool_names]
+        allowed_functions = self._runtime_override_tool_json_schema(allowed_functions)
+
         # For the first message, force the initial tool if one is specified
         force_tool_call = None
         if (
@@ -418,7 +442,7 @@ class Agent(BaseAgent):
                 tool_call_id = response_message.tool_calls[0].id
                 assert tool_call_id is not None  # should be defined
 
-            # only necessary to add the tool_cal_id to a function call (antipattern)
+            # only necessary to add the tool_call_id to a function call (antipattern)
             # response_message_dict = response_message.model_dump()
             # response_message_dict["tool_call_id"] = tool_call_id
 
@@ -513,6 +537,10 @@ class Agent(BaseAgent):
             # Failure case 3: function failed during execution
             # NOTE: the msg_obj associated with the "Running " message is the prior assistant message, not the function/tool role message
             #       this is because the function/tool role message is only created once the function/tool has executed/returned
+
+            # handle cases where we return a json message
+            if "message" in function_args:
+                function_args["message"] = str(function_args.get("message", ""))
             self.interface.function_message(f"Running {function_name}({function_args})", msg_obj=messages[-1], chunk_index=self.chunk_index)
             self.chunk_index += 1
             try:
@@ -529,22 +557,23 @@ class Agent(BaseAgent):
                     },
                 )
 
-                function_response, sandbox_run_result = self.execute_tool_and_persist_state(function_name, function_args, target_letta_tool)
+                tool_execution_result = self.execute_tool_and_persist_state(function_name, function_args, target_letta_tool)
+                function_response = tool_execution_result.func_return
 
                 log_event(
                     "tool_call_ended",
                     attributes={
                         "function_response": function_response,
-                        "sandbox_run_result": sandbox_run_result.model_dump() if sandbox_run_result else None,
+                        "tool_execution_result": tool_execution_result.model_dump(),
                     },
                 )
                 log_telemetry(
                     self.logger, "_handle_ai_response execute tool finish", function_name=function_name, function_args=function_args
                 )
 
-                if sandbox_run_result and sandbox_run_result.status == "error":
+                if tool_execution_result and tool_execution_result.status == "error":
                     tool_return = ToolReturn(
-                        status=sandbox_run_result.status, stdout=sandbox_run_result.stdout, stderr=sandbox_run_result.stderr
+                        status=tool_execution_result.status, stdout=tool_execution_result.stdout, stderr=tool_execution_result.stderr
                     )
                     messages = self._handle_function_error_response(
                         function_response,
@@ -598,14 +627,10 @@ class Agent(BaseAgent):
             # Step 4: check if function response is an error
             if function_response_string.startswith(ERROR_MESSAGE_PREFIX):
                 error_msg = function_response_string
-                tool_return = (
-                    ToolReturn(
-                        status=sandbox_run_result.status,
-                        stdout=sandbox_run_result.stdout,
-                        stderr=sandbox_run_result.stderr,
-                    )
-                    if sandbox_run_result
-                    else None
+                tool_return = ToolReturn(
+                    status=tool_execution_result.status,
+                    stdout=tool_execution_result.stdout,
+                    stderr=tool_execution_result.stderr,
                 )
                 messages = self._handle_function_error_response(
                     error_msg,
@@ -622,14 +647,10 @@ class Agent(BaseAgent):
 
             # If no failures happened along the way: ...
             # Step 5: send the info on the function call and function response to GPT
-            tool_return = (
-                ToolReturn(
-                    status=sandbox_run_result.status,
-                    stdout=sandbox_run_result.stdout,
-                    stderr=sandbox_run_result.stderr,
-                )
-                if sandbox_run_result
-                else None
+            tool_return = ToolReturn(
+                status=tool_execution_result.status,
+                stdout=tool_execution_result.stdout,
+                stderr=tool_execution_result.stderr,
             )
             messages.append(
                 Message(
@@ -641,7 +662,7 @@ class Agent(BaseAgent):
                     content=[TextContent(text=function_response)],
                     tool_call_id=tool_call_id,
                     # Letta extras
-                    tool_returns=[tool_return] if sandbox_run_result else None,
+                    tool_returns=[tool_return],
                     group_id=group_id,
                 )
             )  # extend conversation with function response
@@ -691,7 +712,7 @@ class Agent(BaseAgent):
     @trace_method
     def step(
         self,
-        messages: Union[Message, List[Message]],
+        input_messages: List[MessageCreate],
         # additional args
         chaining: bool = True,
         max_chaining_steps: Optional[int] = None,
@@ -704,7 +725,9 @@ class Agent(BaseAgent):
         # But just to be safe
         self.tool_rules_solver.clear_tool_history()
 
-        next_input_message = messages if isinstance(messages, list) else [messages]
+        # Convert MessageCreate objects to Message objects
+        message_objects = [prepare_input_message_create(m, self.agent_state.id, True, True) for m in input_messages]
+        next_input_messages = message_objects
         counter = 0
         total_usage = UsageStatistics()
         step_count = 0
@@ -715,7 +738,7 @@ class Agent(BaseAgent):
             kwargs["step_count"] = step_count
             kwargs["last_function_failed"] = function_failed
             step_response = self.inner_step(
-                messages=next_input_message,
+                messages=next_input_messages,
                 put_inner_thoughts_first=put_inner_thoughts_first,
                 **kwargs,
             )
@@ -745,36 +768,42 @@ class Agent(BaseAgent):
             # Chain handlers
             elif token_warning and summarizer_settings.send_memory_warning_message:
                 assert self.agent_state.created_by_id is not None
-                next_input_message = Message.dict_to_message(
-                    agent_id=self.agent_state.id,
-                    model=self.model,
-                    openai_message_dict={
-                        "role": "user",  # TODO: change to system?
-                        "content": get_token_limit_warning(),
-                    },
-                )
+                next_input_messages = [
+                    Message.dict_to_message(
+                        agent_id=self.agent_state.id,
+                        model=self.model,
+                        openai_message_dict={
+                            "role": "user",  # TODO: change to system?
+                            "content": get_token_limit_warning(),
+                        },
+                    ),
+                ]
                 continue  # always chain
             elif function_failed:
                 assert self.agent_state.created_by_id is not None
-                next_input_message = Message.dict_to_message(
-                    agent_id=self.agent_state.id,
-                    model=self.model,
-                    openai_message_dict={
-                        "role": "user",  # TODO: change to system?
-                        "content": get_heartbeat(FUNC_FAILED_HEARTBEAT_MESSAGE),
-                    },
-                )
+                next_input_messages = [
+                    Message.dict_to_message(
+                        agent_id=self.agent_state.id,
+                        model=self.model,
+                        openai_message_dict={
+                            "role": "user",  # TODO: change to system?
+                            "content": get_heartbeat(FUNC_FAILED_HEARTBEAT_MESSAGE),
+                        },
+                    )
+                ]
                 continue  # always chain
             elif heartbeat_request:
                 assert self.agent_state.created_by_id is not None
-                next_input_message = Message.dict_to_message(
-                    agent_id=self.agent_state.id,
-                    model=self.model,
-                    openai_message_dict={
-                        "role": "user",  # TODO: change to system?
-                        "content": get_heartbeat(REQ_HEARTBEAT_MESSAGE),
-                    },
-                )
+                next_input_messages = [
+                    Message.dict_to_message(
+                        agent_id=self.agent_state.id,
+                        model=self.model,
+                        openai_message_dict={
+                            "role": "user",  # TODO: change to system?
+                            "content": get_heartbeat(REQ_HEARTBEAT_MESSAGE),
+                        },
+                    )
+                ]
                 continue  # always chain
             # Letta no-op / yield
             else:
@@ -788,7 +817,7 @@ class Agent(BaseAgent):
 
     def inner_step(
         self,
-        messages: Union[Message, List[Message]],
+        messages: List[Message],
         first_message: bool = False,
         first_message_retry_limit: int = FIRST_MESSAGE_ATTEMPTS,
         skip_verify: bool = False,
@@ -814,11 +843,8 @@ class Agent(BaseAgent):
             self.update_memory_if_changed(current_persisted_memory)
 
             # Step 1: add user message
-            if isinstance(messages, Message):
-                messages = [messages]
-
             if not all(isinstance(m, Message) for m in messages):
-                raise ValueError(f"messages should be a Message or a list of Message, got {type(messages)}")
+                raise ValueError(f"messages should be a list of Message, got {[type(m) for m in messages]}")
 
             in_context_messages = self.agent_manager.get_in_context_messages(agent_id=self.agent_state.id, actor=self.user)
             input_message_sequence = in_context_messages + messages
@@ -1229,9 +1255,7 @@ class Agent(BaseAgent):
         return context_window_breakdown.context_window_size_current
 
     # TODO: Refactor into separate class v.s. large if/elses here
-    def execute_tool_and_persist_state(
-        self, function_name: str, function_args: dict, target_letta_tool: Tool
-    ) -> tuple[Any, Optional[SandboxRunResult]]:
+    def execute_tool_and_persist_state(self, function_name: str, function_args: dict, target_letta_tool: Tool) -> ToolExecutionResult:
         """
         Execute tool modifications and persist the state of the agent.
         Note: only some agent state modifications will be persisted, such as data in the AgentState ORM and block data
@@ -1293,8 +1317,10 @@ class Agent(BaseAgent):
                     )
 
                 function_response, is_error = mcp_client.execute_tool(tool_name=function_name, tool_args=function_args)
-                sandbox_run_result = SandboxRunResult(status="error" if is_error else "success")
-                return function_response, sandbox_run_result
+                return ToolExecutionResult(
+                    status="error" if is_error else "success",
+                    func_return=function_response,
+                )
             else:
                 try:
                     # Parse the source code to extract function annotations
@@ -1311,23 +1337,29 @@ class Agent(BaseAgent):
                 agent_state_copy.tools = []
                 agent_state_copy.tool_rules = []
 
-                sandbox_run_result = ToolExecutionSandbox(function_name, function_args, self.user, tool_object=target_letta_tool).run(
+                tool_execution_result = ToolExecutionSandbox(function_name, function_args, self.user, tool_object=target_letta_tool).run(
                     agent_state=agent_state_copy
                 )
-                function_response, updated_agent_state = sandbox_run_result.func_return, sandbox_run_result.agent_state
                 assert orig_memory_str == self.agent_state.memory.compile(), "Memory should not be modified in a sandbox tool"
-                if updated_agent_state is not None:
-                    self.update_memory_if_changed(updated_agent_state.memory)
-                return function_response, sandbox_run_result
+                if tool_execution_result.agent_state is not None:
+                    self.update_memory_if_changed(tool_execution_result.agent_state.memory)
+                return tool_execution_result
         except Exception as e:
             # Need to catch error here, or else trunction wont happen
             # TODO: modify to function execution error
             function_response = get_friendly_error_msg(
                 function_name=function_name, exception_name=type(e).__name__, exception_message=str(e)
             )
-            return function_response, SandboxRunResult(status="error")
+            return ToolExecutionResult(
+                status="error",
+                func_return=function_response,
+                stderr=[traceback.format_exc()],
+            )
 
-        return function_response, None
+        return ToolExecutionResult(
+            status="success",
+            func_return=function_response,
+        )
 
 
 def save_agent(agent: Agent):

letta/agents/letta_agent.py

@@ -324,11 +324,11 @@ class LettaAgent(BaseAgent):
                 tool_execution_manager = ToolExecutionManager(agent_state=agent_state, actor=self.actor)
                 # TODO: Integrate sandbox result
                 log_event(name=f"start_{tool_name}_execution", attributes=tool_args)
-                function_response, _ = await tool_execution_manager.execute_tool_async(
+                tool_execution_result = await tool_execution_manager.execute_tool_async(
                     function_name=tool_name, function_args=tool_args, tool=target_tool
                 )
                 log_event(name=f"finish_{tool_name}_execution", attributes=tool_args)
-                return function_response, True
+                return tool_execution_result.func_return, True
         except Exception as e:
             return f"Failed to call tool. Error: {e}", False
 

letta/agents/letta_agent_batch.py

@@ -37,6 +37,7 @@ from letta.services.passage_manager import PassageManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
 from letta.services.tool_executor.tool_execution_manager import ToolExecutionManager
 from letta.settings import tool_settings
+from letta.tracing import log_event, trace_method
 from letta.utils import united_diff
 
 logger = get_logger(__name__)
@@ -82,12 +83,12 @@ async def execute_tool_wrapper(params: ToolExecutionParams):
             sandbox_config=params.sbx_config,
             sandbox_env_vars=params.sbx_env_vars,
         )
-        result, _ = await mgr.execute_tool_async(
+        tool_execution_result = await mgr.execute_tool_async(
             function_name=params.tool_call_name,
             function_args=params.tool_args,
             tool=target_tool,
         )
-        return params.agent_id, (result, True)
+        return params.agent_id, (tool_execution_result.func_return, True)
     except Exception as e:
         return params.agent_id, (f"Failed to call tool. Error: {e}", False)
 
@@ -120,55 +121,54 @@ class LettaAgentBatch:
         self.actor = actor
         self.max_steps = max_steps
 
+    @trace_method
     async def step_until_request(
         self,
         batch_requests: List[LettaBatchRequest],
         letta_batch_job_id: str,
         agent_step_state_mapping: Optional[Dict[str, AgentStepState]] = None,
     ) -> LettaBatchResponse:
-        # Basic checks
+        log_event(name="validate_inputs")
         if not batch_requests:
             raise ValueError("Empty list of batch_requests passed in!")
         if agent_step_state_mapping is None:
             agent_step_state_mapping = {}
 
+        log_event(name="load_and_prepare_agents")
         agent_messages_mapping: Dict[str, List[Message]] = {}
         agent_tools_mapping: Dict[str, List[dict]] = {}
         agent_states = []
-
         for batch_request in batch_requests:
             agent_id = batch_request.agent_id
             agent_state = self.agent_manager.get_agent_by_id(agent_id, actor=self.actor)
             agent_states.append(agent_state)
+
             agent_messages_mapping[agent_id] = self._get_in_context_messages_per_agent(
                 agent_state=agent_state, input_messages=batch_request.messages
             )
 
-            # TODO: Think about a cleaner way to do this?
             if agent_id not in agent_step_state_mapping:
                 agent_step_state_mapping[agent_id] = AgentStepState(
                     step_number=0, tool_rules_solver=ToolRulesSolver(tool_rules=agent_state.tool_rules)
                 )
 
-            agent_tools_mapping[agent_id] = self._prepare_tools_per_agent(
-                agent_state, agent_step_state_mapping.get(agent_id).tool_rules_solver
-            )
+            agent_tools_mapping[agent_id] = self._prepare_tools_per_agent(agent_state, agent_step_state_mapping[agent_id].tool_rules_solver)
 
-        # TODO: This is a hack, this is because LLM client expects a LLM config
-        # TODO: But that doesn't really work in batch land
-        # TODO: @caren will factor this out
+        log_event(name="init_llm_client")
         llm_client = LLMClient.create(
             llm_config=agent_states[0].llm_config,
             put_inner_thoughts_first=True,
         )
-        agent_llm_config_mapping = {agent_state.id: agent_state.llm_config for agent_state in agent_states}
+        agent_llm_config_mapping = {s.id: s.llm_config for s in agent_states}
+
+        log_event(name="send_llm_batch_request")
         batch_response = await llm_client.send_llm_batch_request_async(
             agent_messages_mapping=agent_messages_mapping,
             agent_tools_mapping=agent_tools_mapping,
             agent_llm_config_mapping=agent_llm_config_mapping,
         )
 
-        # Write the response into the jobs table, where it will get picked up by the next cron run
+        log_event(name="persist_llm_batch_job")
         llm_batch_job = self.batch_manager.create_llm_batch_job(
             llm_provider=ProviderType.anthropic,  # TODO: Expand to more providers
             create_batch_response=batch_response,
@@ -177,24 +177,26 @@ class LettaAgentBatch:
             letta_batch_job_id=letta_batch_job_id,
         )
 
-        # Create batch items in bulk for all agents
+        log_event(name="prepare_batch_items")
         batch_items = []
-        for agent_state in agent_states:
-            agent_step_state = agent_step_state_mapping.get(agent_state.id)
-            batch_item = LLMBatchItem(
-                llm_batch_id=llm_batch_job.id,
-                agent_id=agent_state.id,
-                llm_config=agent_state.llm_config,
-                request_status=JobStatus.created,
-                step_status=AgentStepStatus.paused,
-                step_state=agent_step_state,
+        for state in agent_states:
+            step_state = agent_step_state_mapping[state.id]
+            batch_items.append(
+                LLMBatchItem(
+                    llm_batch_id=llm_batch_job.id,
+                    agent_id=state.id,
+                    llm_config=state.llm_config,
+                    request_status=JobStatus.created,
+                    step_status=AgentStepStatus.paused,
+                    step_state=step_state,
+                )
             )
-            batch_items.append(batch_item)
 
-        # Create all batch items at once using the bulk operation
         if batch_items:
+            log_event(name="bulk_create_batch_items")
             self.batch_manager.create_llm_batch_items_bulk(batch_items, actor=self.actor)
 
+        log_event(name="return_batch_response")
         return LettaBatchResponse(
             letta_batch_id=llm_batch_job.letta_batch_job_id,
             last_llm_batch_id=llm_batch_job.id,
@@ -204,27 +206,27 @@ class LettaAgentBatch:
             created_at=llm_batch_job.created_at,
         )
 
+    @trace_method
     async def resume_step_after_request(self, letta_batch_id: str, llm_batch_id: str) -> LettaBatchResponse:
-        # 1. gather everything we need
+        log_event(name="load_context")
         llm_batch_job = self.batch_manager.get_llm_batch_job_by_id(llm_batch_id=llm_batch_id, actor=self.actor)
         ctx = await self._collect_resume_context(llm_batch_id)
 
-        # 2. persist request‑level status updates
+        log_event(name="update_statuses")
         self._update_request_statuses(ctx.request_status_updates)
 
-        # 3. run the tools in parallel
+        log_event(name="exec_tools")
         exec_results = await self._execute_tools(ctx)
 
-        # 4. create + save assistant/tool messages
+        log_event(name="persist_messages")
         msg_map = self._persist_tool_messages(exec_results, ctx)
 
-        # 5. mark steps complete
+        log_event(name="mark_steps_done")
         self._mark_steps_complete(llm_batch_id, ctx.agent_ids)
 
-        # 6. build next‑round requests / step‑state map
+        log_event(name="prepare_next")
         next_reqs, next_step_state = self._prepare_next_iteration(exec_results, ctx, msg_map)
         if len(next_reqs) == 0:
-            # mark batch job as completed
             self.job_manager.update_job_by_id(job_id=letta_batch_id, job_update=JobUpdate(status=JobStatus.completed), actor=self.actor)
             return LettaBatchResponse(
                 letta_batch_id=llm_batch_job.letta_batch_job_id,
@@ -235,15 +237,16 @@ class LettaAgentBatch:
                 created_at=llm_batch_job.created_at,
             )
 
-        # 7. recurse into the normal stepping pipeline
         return await self.step_until_request(
             batch_requests=next_reqs,
             letta_batch_job_id=letta_batch_id,
             agent_step_state_mapping=next_step_state,
         )
 
+    @trace_method
     async def _collect_resume_context(self, llm_batch_id: str) -> _ResumeContext:
-        batch_items = self.batch_manager.list_llm_batch_items(llm_batch_id=llm_batch_id)
+        # NOTE: We only continue for items with successful results
+        batch_items = self.batch_manager.list_llm_batch_items(llm_batch_id=llm_batch_id, request_status=JobStatus.completed)
 
         agent_ids, agent_state_map = [], {}
         provider_results, name_map, args_map, cont_map = {}, {}, {}, {}
@@ -300,6 +303,7 @@ class LettaAgentBatch:
         env = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(cfg.id, actor=self.actor, limit=100)
         return cfg, env
 
+    @trace_method
     async def _execute_tools(self, ctx: _ResumeContext) -> Sequence[Tuple[str, Tuple[str, bool]]]:
         sbx_cfg, sbx_env = self._build_sandbox()
         params = [