letta-nightly 0.6.53.dev20250418104238__tar.gz → 0.6.54.dev20250419104029__tar.gz

{letta_nightly-0.6.53.dev20250418104238 → letta_nightly-0.6.54.dev20250419104029}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: letta-nightly
-Version: 0.6.53.dev20250418104238
+Version: 0.6.54.dev20250419104029
 Summary: Create LLM agents with long-term memory and custom tools
 License: Apache License
 Author: Letta Team
@@ -22,6 +22,7 @@ Provides-Extra: postgres
 Provides-Extra: qdrant
 Provides-Extra: server
 Provides-Extra: tests
+Requires-Dist: aiomultiprocess (>=0.9.1,<0.10.0)
 Requires-Dist: alembic (>=1.13.3,<2.0.0)
 Requires-Dist: anthropic (>=0.49.0,<0.50.0)
 Requires-Dist: apscheduler (>=3.11.0,<4.0.0)

{letta_nightly-0.6.53.dev20250418104238 → letta_nightly-0.6.54.dev20250419104029}/letta/__init__.py

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

{letta_nightly-0.6.53.dev20250418104238 → letta_nightly-0.6.54.dev20250419104029}/letta/agent.py

@@ -36,7 +36,7 @@ from letta.log import get_logger
 from letta.memory import summarize_messages
 from letta.orm import User
 from letta.orm.enums import ToolType
-from letta.schemas.agent import AgentState, AgentStepResponse, UpdateAgent
+from letta.schemas.agent import AgentState, AgentStepResponse, UpdateAgent, get_prompt_template_for_agent_type
 from letta.schemas.block import BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import MessageRole
@@ -52,11 +52,7 @@ from letta.schemas.tool_rule import TerminalToolRule
 from letta.schemas.usage import LettaUsageStatistics
 from letta.services.agent_manager import AgentManager
 from letta.services.block_manager import BlockManager
-from letta.services.helpers.agent_manager_helper import (
-    check_supports_structured_output,
-    compile_memory_metadata_block,
-    compile_system_message,
-)
+from letta.services.helpers.agent_manager_helper import check_supports_structured_output, compile_memory_metadata_block
 from letta.services.job_manager import JobManager
 from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
@@ -204,7 +200,8 @@ class Agent(BaseAgent):
 
             # refresh memory from DB (using block ids)
             self.agent_state.memory = Memory(
-                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()]
+                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()],
+                prompt_template=get_prompt_template_for_agent_type(self.agent_state.agent_type),
             )
 
             # NOTE: don't do this since re-buildin the memory is handled at the start of the step
@@ -306,29 +303,6 @@ class Agent(BaseAgent):
         elif step_count is not None and step_count > 0 and len(allowed_tool_names) == 1:
             force_tool_call = allowed_tool_names[0]
 
-        if force_tool_call == "core_memory_insert":
-            current_system_message = message_sequence[0]
-            new_memory = Memory(
-                blocks=self.agent_state.memory.blocks,
-                prompt_template=(
-                    "{% for block in blocks %}"
-                    '<{{ block.label }} characters="{{ block.value|length }}/{{ block.limit }}">\n'
-                    "{% for line in block.value.splitlines() %}"
-                    "{{ loop.index0 }}: {{ line }}\n"
-                    "{% endfor %}"
-                    "</{{ block.label }}>"
-                    "{% if not loop.last %}\n{% endif %}"
-                    "{% endfor %}"
-                ),
-            )
-            new_system_message_str = compile_system_message(
-                system_prompt=self.agent_state.system,
-                in_context_memory=new_memory,
-                in_context_memory_last_edit=current_system_message.created_at,
-                previous_message_count=len(message_sequence),
-            )
-            message_sequence[0].content = [TextContent(text=new_system_message_str)]
-
         for attempt in range(1, empty_response_retry_limit + 1):
             try:
                 log_telemetry(self.logger, "_get_ai_reply create start")
@@ -834,7 +808,8 @@ class Agent(BaseAgent):
             # Step 0: update core memory
             # only pulling latest block data if shared memory is being used
             current_persisted_memory = Memory(
-                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()]
+                blocks=[self.block_manager.get_block_by_id(block.id, actor=self.user) for block in self.agent_state.memory.get_blocks()],
+                prompt_template=get_prompt_template_for_agent_type(self.agent_state.agent_type),
             )  # read blocks from DB
             self.update_memory_if_changed(current_persisted_memory)
 

{letta_nightly-0.6.53.dev20250418104238 → letta_nightly-0.6.54.dev20250419104029}/letta/agents/letta_agent.py

@@ -76,6 +76,7 @@ class LettaAgent(BaseAgent):
                 agent_state=agent_state,
                 tool_rules_solver=tool_rules_solver,
                 stream=False,
+                # TODO: also pass in reasoning content
             )
 
             tool_call = response.choices[0].message.tool_calls[0]

letta_nightly-0.6.54.dev20250419104029/letta/agents/letta_agent_batch.py

@@ -0,0 +1,515 @@
+import json
+import uuid
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from aiomultiprocess import Pool
+from anthropic.types.beta.messages import BetaMessageBatchCanceledResult, BetaMessageBatchErroredResult, BetaMessageBatchSucceededResult
+
+from letta.agents.helpers import _prepare_in_context_messages
+from letta.helpers import ToolRulesSolver
+from letta.helpers.datetime_helpers import get_utc_time
+from letta.helpers.tool_execution_helper import enable_strict_mode
+from letta.jobs.types import RequestStatusUpdateInfo, StepStatusUpdateInfo
+from letta.llm_api.llm_client import LLMClient
+from letta.local_llm.constants import INNER_THOUGHTS_KWARG
+from letta.log import get_logger
+from letta.orm.enums import ToolType
+from letta.schemas.agent import AgentState, AgentStepState
+from letta.schemas.enums import AgentStepStatus, JobStatus, ProviderType
+from letta.schemas.job import JobUpdate
+from letta.schemas.letta_message_content import OmittedReasoningContent, ReasoningContent, RedactedReasoningContent, TextContent
+from letta.schemas.letta_request import LettaBatchRequest
+from letta.schemas.letta_response import LettaBatchResponse
+from letta.schemas.llm_batch_job import LLMBatchItem
+from letta.schemas.message import Message, MessageCreate, MessageUpdate
+from letta.schemas.openai.chat_completion_response import ToolCall as OpenAIToolCall
+from letta.schemas.sandbox_config import SandboxConfig, SandboxType
+from letta.schemas.user import User
+from letta.server.rest_api.utils import create_heartbeat_system_message, create_letta_messages_from_llm_response
+from letta.services.agent_manager import AgentManager
+from letta.services.block_manager import BlockManager
+from letta.services.helpers.agent_manager_helper import compile_system_message
+from letta.services.job_manager import JobManager
+from letta.services.llm_batch_manager import LLMBatchManager
+from letta.services.message_manager import MessageManager
+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.utils import united_diff
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class ToolExecutionParams:
+    agent_id: str
+    tool_call_name: str
+    tool_args: Dict[str, Any]
+    agent_state: AgentState
+    actor: User
+    sbx_config: SandboxConfig
+    sbx_env_vars: Dict[str, Any]
+
+
+@dataclass
+class _ResumeContext:
+    batch_items: List[LLMBatchItem]
+    agent_ids: List[str]
+    agent_state_map: Dict[str, AgentState]
+    provider_results: Dict[str, Any]
+    tool_call_name_map: Dict[str, str]
+    tool_call_args_map: Dict[str, Dict[str, Any]]
+    should_continue_map: Dict[str, bool]
+    request_status_updates: List[RequestStatusUpdateInfo]
+
+
+async def execute_tool_wrapper(params: ToolExecutionParams):
+    """
+    Executes the tool in an out‑of‑process worker and returns:
+        (agent_id, (tool_result:str, success_flag:bool))
+    """
+    # locate the tool on the agent
+    target_tool = next((t for t in params.agent_state.tools if t.name == params.tool_call_name), None)
+    if not target_tool:
+        return params.agent_id, (f"Tool not found: {params.tool_call_name}", False)
+
+    try:
+        mgr = ToolExecutionManager(
+            agent_state=params.agent_state,
+            actor=params.actor,
+            sandbox_config=params.sbx_config,
+            sandbox_env_vars=params.sbx_env_vars,
+        )
+        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)
+    except Exception as e:
+        return params.agent_id, (f"Failed to call tool. Error: {e}", False)
+
+
+# TODO: Limitations ->
+# TODO: Only works with anthropic for now
+class LettaAgentBatch:
+
+    def __init__(
+        self,
+        message_manager: MessageManager,
+        agent_manager: AgentManager,
+        block_manager: BlockManager,
+        passage_manager: PassageManager,
+        batch_manager: LLMBatchManager,
+        sandbox_config_manager: SandboxConfigManager,
+        job_manager: JobManager,
+        actor: User,
+        use_assistant_message: bool = True,
+        max_steps: int = 10,
+    ):
+        self.message_manager = message_manager
+        self.agent_manager = agent_manager
+        self.block_manager = block_manager
+        self.passage_manager = passage_manager
+        self.batch_manager = batch_manager
+        self.sandbox_config_manager = sandbox_config_manager
+        self.job_manager = job_manager
+        self.use_assistant_message = use_assistant_message
+        self.actor = actor
+        self.max_steps = max_steps
+
+    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
+        if not batch_requests:
+            raise ValueError("Empty list of batch_requests passed in!")
+        if agent_step_state_mapping is None:
+            agent_step_state_mapping = {}
+
+        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
+            )
+
+        # 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
+        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}
+        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
+        llm_batch_job = self.batch_manager.create_llm_batch_job(
+            llm_provider=ProviderType.anthropic,  # TODO: Expand to more providers
+            create_batch_response=batch_response,
+            actor=self.actor,
+            status=JobStatus.running,
+            letta_batch_job_id=letta_batch_job_id,
+        )
+
+        # Create batch items in bulk for all agents
+        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,
+            )
+            batch_items.append(batch_item)
+
+        # Create all batch items at once using the bulk operation
+        if batch_items:
+            self.batch_manager.create_llm_batch_items_bulk(batch_items, actor=self.actor)
+
+        return LettaBatchResponse(
+            letta_batch_id=llm_batch_job.letta_batch_job_id,
+            last_llm_batch_id=llm_batch_job.id,
+            status=llm_batch_job.status,
+            agent_count=len(agent_states),
+            last_polled_at=get_utc_time(),
+            created_at=llm_batch_job.created_at,
+        )
+
+    async def resume_step_after_request(self, letta_batch_id: str, llm_batch_id: str) -> LettaBatchResponse:
+        # 1. gather everything we need
+        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
+        self._update_request_statuses(ctx.request_status_updates)
+
+        # 3. run the tools in parallel
+        exec_results = await self._execute_tools(ctx)
+
+        # 4. create + save assistant/tool messages
+        msg_map = self._persist_tool_messages(exec_results, ctx)
+
+        # 5. mark steps complete
+        self._mark_steps_complete(llm_batch_id, ctx.agent_ids)
+
+        # 6. build next‑round requests / step‑state map
+        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,
+                last_llm_batch_id=llm_batch_job.id,
+                status=JobStatus.completed,
+                agent_count=len(ctx.agent_ids),
+                last_polled_at=get_utc_time(),
+                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,
+        )
+
+    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)
+
+        agent_ids, agent_state_map = [], {}
+        provider_results, name_map, args_map, cont_map = {}, {}, {}, {}
+        request_status_updates: List[RequestStatusUpdateInfo] = []
+
+        for item in batch_items:
+            aid = item.agent_id
+            agent_ids.append(aid)
+            agent_state_map[aid] = self.agent_manager.get_agent_by_id(aid, actor=self.actor)
+            provider_results[aid] = item.batch_request_result.result
+
+            # status bookkeeping
+            pr = provider_results[aid]
+            status = (
+                JobStatus.completed
+                if isinstance(pr, BetaMessageBatchSucceededResult)
+                else (
+                    JobStatus.failed
+                    if isinstance(pr, BetaMessageBatchErroredResult)
+                    else JobStatus.cancelled if isinstance(pr, BetaMessageBatchCanceledResult) else JobStatus.expired
+                )
+            )
+            request_status_updates.append(RequestStatusUpdateInfo(llm_batch_id=llm_batch_id, agent_id=aid, request_status=status))
+
+            # translate provider‑specific response → OpenAI‑style tool call (unchanged)
+            llm_client = LLMClient.create(llm_config=item.llm_config, put_inner_thoughts_first=True)
+            tool_call = (
+                llm_client.convert_response_to_chat_completion(response_data=pr.message.model_dump(), input_messages=[])
+                .choices[0]
+                .message.tool_calls[0]
+            )
+
+            name, args, cont = self._extract_tool_call_and_decide_continue(tool_call, item.step_state)
+            name_map[aid], args_map[aid], cont_map[aid] = name, args, cont
+
+        return _ResumeContext(
+            batch_items=batch_items,
+            agent_ids=agent_ids,
+            agent_state_map=agent_state_map,
+            provider_results=provider_results,
+            tool_call_name_map=name_map,
+            tool_call_args_map=args_map,
+            should_continue_map=cont_map,
+            request_status_updates=request_status_updates,
+        )
+
+    def _update_request_statuses(self, updates: List[RequestStatusUpdateInfo]) -> None:
+        if updates:
+            self.batch_manager.bulk_update_llm_batch_items_request_status_by_agent(updates=updates)
+
+    def _build_sandbox(self) -> Tuple[SandboxConfig, Dict[str, Any]]:
+        sbx_type = SandboxType.E2B if tool_settings.e2b_api_key else SandboxType.LOCAL
+        cfg = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=sbx_type, actor=self.actor)
+        env = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(cfg.id, actor=self.actor, limit=100)
+        return cfg, env
+
+    async def _execute_tools(self, ctx: _ResumeContext) -> Sequence[Tuple[str, Tuple[str, bool]]]:
+        sbx_cfg, sbx_env = self._build_sandbox()
+        params = [
+            ToolExecutionParams(
+                agent_id=aid,
+                tool_call_name=ctx.tool_call_name_map[aid],
+                tool_args=ctx.tool_call_args_map[aid],
+                agent_state=ctx.agent_state_map[aid],
+                actor=self.actor,
+                sbx_config=sbx_cfg,
+                sbx_env_vars=sbx_env,
+            )
+            for aid in ctx.agent_ids
+        ]
+        async with Pool() as pool:
+            return await pool.map(execute_tool_wrapper, params)
+
+    def _persist_tool_messages(
+        self,
+        exec_results: Sequence[Tuple[str, Tuple[str, bool]]],
+        ctx: _ResumeContext,
+    ) -> Dict[str, List[Message]]:
+        msg_map: Dict[str, List[Message]] = {}
+        for aid, (tool_res, success) in exec_results:
+            msgs = self._create_tool_call_messages(
+                agent_state=ctx.agent_state_map[aid],
+                tool_call_name=ctx.tool_call_name_map[aid],
+                tool_call_args=ctx.tool_call_args_map[aid],
+                tool_exec_result=tool_res,
+                success_flag=success,
+                reasoning_content=None,
+            )
+            msg_map[aid] = msgs
+        # flatten & persist
+        self.message_manager.create_many_messages([m for msgs in msg_map.values() for m in msgs], actor=self.actor)
+        return msg_map
+
+    def _mark_steps_complete(self, llm_batch_id: str, agent_ids: List[str]) -> None:
+        updates = [
+            StepStatusUpdateInfo(llm_batch_id=llm_batch_id, agent_id=aid, step_status=AgentStepStatus.completed) for aid in agent_ids
+        ]
+        self.batch_manager.bulk_update_llm_batch_items_step_status_by_agent(updates)
+
+    def _prepare_next_iteration(
+        self,
+        exec_results: Sequence[Tuple[str, Tuple[str, bool]]],
+        ctx: _ResumeContext,
+        msg_map: Dict[str, List[Message]],
+    ) -> Tuple[List[LettaBatchRequest], Dict[str, AgentStepState]]:
+        # who continues?
+        continues = [aid for aid, cont in ctx.should_continue_map.items() if cont]
+
+        success_flag_map = {aid: flag for aid, (_res, flag) in exec_results}
+
+        batch_reqs: List[LettaBatchRequest] = []
+        for aid in continues:
+            heartbeat = create_heartbeat_system_message(
+                agent_id=aid,
+                model=ctx.agent_state_map[aid].llm_config.model,
+                function_call_success=success_flag_map[aid],
+                actor=self.actor,
+            )
+            batch_reqs.append(
+                LettaBatchRequest(
+                    agent_id=aid, messages=[MessageCreate.model_validate(heartbeat.model_dump(include={"role", "content", "name", "otid"}))]
+                )
+            )
+
+        # extend in‑context ids when necessary
+        for aid, new_msgs in msg_map.items():
+            ast = ctx.agent_state_map[aid]
+            if not ast.message_buffer_autoclear:
+                self.agent_manager.set_in_context_messages(
+                    agent_id=aid,
+                    message_ids=ast.message_ids + [m.id for m in new_msgs],
+                    actor=self.actor,
+                )
+
+        # bump step number
+        step_map = {
+            item.agent_id: item.step_state.model_copy(update={"step_number": item.step_state.step_number + 1}) for item in ctx.batch_items
+        }
+        return batch_reqs, step_map
+
+    def _create_tool_call_messages(
+        self,
+        agent_state: AgentState,
+        tool_call_name: str,
+        tool_call_args: Dict[str, Any],
+        tool_exec_result: str,
+        success_flag: bool,
+        reasoning_content: Optional[List[Union[TextContent, ReasoningContent, RedactedReasoningContent, OmittedReasoningContent]]] = None,
+    ) -> List[Message]:
+        tool_call_id = f"call_{uuid.uuid4().hex[:8]}"
+
+        tool_call_messages = create_letta_messages_from_llm_response(
+            agent_id=agent_state.id,
+            model=agent_state.llm_config.model,
+            function_name=tool_call_name,
+            function_arguments=tool_call_args,
+            tool_call_id=tool_call_id,
+            function_call_success=success_flag,
+            function_response=tool_exec_result,
+            actor=self.actor,
+            add_heartbeat_request_system_message=False,
+            reasoning_content=reasoning_content,
+            pre_computed_assistant_message_id=None,
+            pre_computed_tool_message_id=None,
+        )
+
+        return tool_call_messages
+
+    # TODO: This is doing a lot of dict passing
+    # TODO: Make the passing here typed
+    def _extract_tool_call_and_decide_continue(
+        self, tool_call: OpenAIToolCall, agent_step_state: AgentStepState
+    ) -> Tuple[str, Dict[str, Any], bool]:
+        """
+        Now that streaming is done, handle the final AI response.
+        This might yield additional SSE tokens if we do stalling.
+        At the end, set self._continue_execution accordingly.
+        """
+        tool_call_name = tool_call.function.name
+        tool_call_args_str = tool_call.function.arguments
+
+        try:
+            tool_args = json.loads(tool_call_args_str)
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to JSON decode tool call argument string: {tool_call_args_str}")
+            tool_args = {}
+
+        # Get request heartbeats and coerce to bool
+        request_heartbeat = tool_args.pop("request_heartbeat", False)
+        # Pre-emptively pop out inner_thoughts
+        tool_args.pop(INNER_THOUGHTS_KWARG, "")
+
+        # So this is necessary, because sometimes non-structured outputs makes mistakes
+        if isinstance(request_heartbeat, str):
+            request_heartbeat = request_heartbeat.lower() == "true"
+        else:
+            request_heartbeat = bool(request_heartbeat)
+
+        continue_stepping = request_heartbeat
+        tool_rules_solver = agent_step_state.tool_rules_solver
+        tool_rules_solver.register_tool_call(tool_name=tool_call_name)
+        if tool_rules_solver.is_terminal_tool(tool_name=tool_call_name):
+            continue_stepping = False
+        elif tool_rules_solver.has_children_tools(tool_name=tool_call_name):
+            continue_stepping = True
+        elif tool_rules_solver.is_continue_tool(tool_name=tool_call_name):
+            continue_stepping = True
+
+        step_count = agent_step_state.step_number
+        if step_count >= self.max_steps:
+            logger.warning("Hit max steps, stopping agent loop prematurely.")
+            continue_stepping = False
+
+        return tool_call_name, tool_args, continue_stepping
+
+    def _prepare_tools_per_agent(self, agent_state: AgentState, tool_rules_solver: ToolRulesSolver) -> List[dict]:
+        tools = [t for t in agent_state.tools if t.tool_type in {ToolType.CUSTOM, ToolType.LETTA_CORE, ToolType.LETTA_MEMORY_CORE}]
+        valid_tool_names = tool_rules_solver.get_allowed_tool_names(available_tools=set([t.name for t in tools]))
+        return [enable_strict_mode(t.json_schema) for t in tools if t.name in set(valid_tool_names)]
+
+    def _get_in_context_messages_per_agent(self, agent_state: AgentState, input_messages: List[MessageCreate]) -> List[Message]:
+        current_in_context_messages, new_in_context_messages = _prepare_in_context_messages(
+            input_messages, agent_state, self.message_manager, self.actor
+        )
+
+        in_context_messages = self._rebuild_memory(current_in_context_messages + new_in_context_messages, agent_state)
+        return in_context_messages
+
+    # TODO: Make this a bullk function
+    def _rebuild_memory(self, in_context_messages: List[Message], agent_state: AgentState) -> List[Message]:
+        agent_state = self.agent_manager.refresh_memory(agent_state=agent_state, actor=self.actor)
+
+        # TODO: This is a pretty brittle pattern established all over our code, need to get rid of this
+        curr_system_message = in_context_messages[0]
+        curr_memory_str = agent_state.memory.compile()
+        curr_system_message_text = curr_system_message.content[0].text
+        if curr_memory_str in curr_system_message_text:
+            # NOTE: could this cause issues if a block is removed? (substring match would still work)
+            logger.debug(
+                f"Memory hasn't changed for agent id={agent_state.id} and actor=({self.actor.id}, {self.actor.name}), skipping system prompt rebuild"
+            )
+            return in_context_messages
+
+        memory_edit_timestamp = get_utc_time()
+
+        num_messages = self.message_manager.size(actor=self.actor, agent_id=agent_state.id)
+        num_archival_memories = self.passage_manager.size(actor=self.actor, agent_id=agent_state.id)
+
+        new_system_message_str = compile_system_message(
+            system_prompt=agent_state.system,
+            in_context_memory=agent_state.memory,
+            in_context_memory_last_edit=memory_edit_timestamp,
+            previous_message_count=num_messages,
+            archival_memory_size=num_archival_memories,
+        )
+
+        diff = united_diff(curr_system_message_text, new_system_message_str)
+        if len(diff) > 0:
+            logger.debug(f"Rebuilding system with new memory...\nDiff:\n{diff}")
+
+            new_system_message = self.message_manager.update_message_by_id(
+                curr_system_message.id, message_update=MessageUpdate(content=new_system_message_str), actor=self.actor
+            )
+
+            # Skip pulling down the agent's memory again to save on a db call
+            return [new_system_message] + in_context_messages[1:]
+
+        else:
+            return in_context_messages

{letta_nightly-0.6.53.dev20250418104238 → letta_nightly-0.6.54.dev20250419104029}/letta/constants.py

@@ -56,10 +56,10 @@ BASE_MEMORY_TOOLS = ["core_memory_append", "core_memory_replace"]
 BASE_SLEEPTIME_CHAT_TOOLS = ["send_message", "conversation_search", "archival_memory_search"]
 # Base memory tools for sleeptime agent
 BASE_SLEEPTIME_TOOLS = [
-    "rethink_memory",
-    "finish_rethinking_memory",
-    "view_core_memory_with_line_numbers",
-    "core_memory_insert",
+    "memory_replace",
+    "memory_insert",
+    "memory_rethink",
+    "memory_finish_edits",
     "archival_memory_insert",
     "archival_memory_search",
     "conversation_search",
@@ -103,6 +103,11 @@ ERROR_MESSAGE_PREFIX = "Error"
 
 NON_USER_MSG_PREFIX = "[This is an automated system message hidden from the user] "
 
+CORE_MEMORY_LINE_NUMBER_WARNING = (
+    "# NOTE: Line numbers shown below are to help during editing. Do NOT include line number prefixes in your memory edit tool calls."
+)
+
+
 # Constants to do with summarization / conversation length window
 # The max amount of tokens supported by the underlying model (eg 8k for gpt-4 and Mistral 7B)
 LLM_MAX_TOKENS = {
@@ -110,6 +115,12 @@ LLM_MAX_TOKENS = {
     "deepseek-chat": 64000,
     "deepseek-reasoner": 64000,
     ## OpenAI models: https://platform.openai.com/docs/models/overview
+    "gpt-4.1": 1047576,
+    "gpt-4.1-2025-04-14": 1047576,
+    "gpt-4.1-mini": 1047576,
+    "gpt-4.1-mini-2025-04-14": 1047576,
+    "gpt-4.1-nano": 1047576,
+    "gpt-4.1-nano-2025-04-14": 1047576,
     # gpt-4.5-preview
     "gpt-4.5-preview": 128000,
     "gpt-4.5-preview-2025-02-27": 128000,