letta-nightly 0.6.48.dev20250407104216__py3-none-any.whl → 0.6.49.dev20250408104230__py3-none-any.whl

letta/services/block_manager.py

@@ -1,10 +1,13 @@
 import os
 from typing import List, Optional
 
+from sqlalchemy.orm import Session
+
 from letta.orm.block import Block as BlockModel
+from letta.orm.block_history import BlockHistory
+from letta.orm.enums import ActorType
 from letta.orm.errors import NoResultFound
 from letta.schemas.agent import AgentState as PydanticAgentState
-from letta.schemas.block import Block
 from letta.schemas.block import Block as PydanticBlock
 from letta.schemas.block import BlockUpdate, Human, Persona
 from letta.schemas.user import User as PydanticUser
@@ -21,7 +24,7 @@ class BlockManager:
         self.session_maker = db_context
 
     @enforce_types
-    def create_or_update_block(self, block: Block, actor: PydanticUser) -> PydanticBlock:
+    def create_or_update_block(self, block: PydanticBlock, actor: PydanticUser) -> PydanticBlock:
         """Create a new block based on the Block schema."""
         db_block = self.get_block_by_id(block.id, actor)
         if db_block:
@@ -140,3 +143,188 @@ class BlockManager:
             agents_pydantic = [agent.to_pydantic() for agent in agents_orm]
 
             return agents_pydantic
+
+    # Block History Functions
+
+    @enforce_types
+    def checkpoint_block(
+        self,
+        block_id: str,
+        actor: PydanticUser,
+        agent_id: Optional[str] = None,
+        use_preloaded_block: Optional[BlockModel] = None,  # For concurrency tests
+    ) -> PydanticBlock:
+        """
+        Create a new checkpoint for the given Block by copying its
+        current state into BlockHistory, using SQLAlchemy's built-in
+        version_id_col for concurrency checks.
+
+        - If the block was undone to an earlier checkpoint, we remove
+          any "future" checkpoints beyond the current state to keep a
+          strictly linear history.
+        - A single commit at the end ensures atomicity.
+        """
+        with self.session_maker() as session:
+            # 1) Load the Block
+            if use_preloaded_block is not None:
+                block = session.merge(use_preloaded_block)
+            else:
+                block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
+
+            # 2) Identify the block's current checkpoint (if any)
+            current_entry = None
+            if block.current_history_entry_id:
+                current_entry = session.get(BlockHistory, block.current_history_entry_id)
+
+            # The current sequence, or 0 if no checkpoints exist
+            current_seq = current_entry.sequence_number if current_entry else 0
+
+            # 3) Truncate any future checkpoints
+            #    If we are at seq=2, but there's a seq=3 or higher from a prior "redo chain",
+            #    remove those, so we maintain a strictly linear undo/redo stack.
+            session.query(BlockHistory).filter(BlockHistory.block_id == block.id, BlockHistory.sequence_number > current_seq).delete()
+
+            # 4) Determine the next sequence number
+            next_seq = current_seq + 1
+
+            # 5) Create a new BlockHistory row reflecting the block's current state
+            history_entry = BlockHistory(
+                organization_id=actor.organization_id,
+                block_id=block.id,
+                sequence_number=next_seq,
+                description=block.description,
+                label=block.label,
+                value=block.value,
+                limit=block.limit,
+                metadata_=block.metadata_,
+                actor_type=ActorType.LETTA_AGENT if agent_id else ActorType.LETTA_USER,
+                actor_id=agent_id if agent_id else actor.id,
+            )
+            history_entry.create(session, actor=actor, no_commit=True)
+
+            # 6) Update the block’s pointer to the new checkpoint
+            block.current_history_entry_id = history_entry.id
+
+            # 7) Flush changes, then commit once
+            block = block.update(db_session=session, actor=actor, no_commit=True)
+            session.commit()
+
+            return block.to_pydantic()
+
+    @enforce_types
+    def _move_block_to_sequence(self, session: Session, block: BlockModel, target_seq: int, actor: PydanticUser) -> BlockModel:
+        """
+        Internal helper that moves the 'block' to the specified 'target_seq' within BlockHistory.
+        1) Find the BlockHistory row at sequence_number=target_seq
+        2) Copy fields into the block
+        3) Update and flush (no_commit=True) - the caller is responsible for final commit
+
+        Raises:
+            NoResultFound: if no BlockHistory row for (block_id, target_seq)
+        """
+        if not block.id:
+            raise ValueError("Block is missing an ID. Cannot move sequence.")
+
+        target_entry = (
+            session.query(BlockHistory)
+            .filter(
+                BlockHistory.block_id == block.id,
+                BlockHistory.sequence_number == target_seq,
+            )
+            .one_or_none()
+        )
+        if not target_entry:
+            raise NoResultFound(f"No BlockHistory row found for block_id={block.id} at sequence={target_seq}")
+
+        # Copy fields from target_entry to block
+        block.description = target_entry.description  # type: ignore
+        block.label = target_entry.label  # type: ignore
+        block.value = target_entry.value  # type: ignore
+        block.limit = target_entry.limit  # type: ignore
+        block.metadata_ = target_entry.metadata_  # type: ignore
+        block.current_history_entry_id = target_entry.id  # type: ignore
+
+        # Update in DB (optimistic locking).
+        # We'll do a flush now; the caller does final commit.
+        updated_block = block.update(db_session=session, actor=actor, no_commit=True)
+        return updated_block
+
+    @enforce_types
+    def undo_checkpoint_block(self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None) -> PydanticBlock:
+        """
+        Move the block to the immediately previous checkpoint in BlockHistory.
+        If older sequences have been pruned, we jump to the largest sequence
+        number that is still < current_seq.
+        """
+        with self.session_maker() as session:
+            # 1) Load the current block
+            block = (
+                session.merge(use_preloaded_block)
+                if use_preloaded_block
+                else BlockModel.read(db_session=session, identifier=block_id, actor=actor)
+            )
+
+            if not block.current_history_entry_id:
+                raise ValueError(f"Block {block_id} has no history entry - cannot undo.")
+
+            current_entry = session.get(BlockHistory, block.current_history_entry_id)
+            if not current_entry:
+                raise NoResultFound(f"BlockHistory row not found for id={block.current_history_entry_id}")
+
+            current_seq = current_entry.sequence_number
+
+            # 2) Find the largest sequence < current_seq
+            previous_entry = (
+                session.query(BlockHistory)
+                .filter(BlockHistory.block_id == block.id, BlockHistory.sequence_number < current_seq)
+                .order_by(BlockHistory.sequence_number.desc())
+                .first()
+            )
+            if not previous_entry:
+                # No earlier checkpoint available
+                raise ValueError(f"Block {block_id} is already at the earliest checkpoint (seq={current_seq}). Cannot undo further.")
+
+            # 3) Move to that sequence
+            block = self._move_block_to_sequence(session, block, previous_entry.sequence_number, actor)
+
+            # 4) Commit
+            session.commit()
+            return block.to_pydantic()
+
+    @enforce_types
+    def redo_checkpoint_block(self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None) -> PydanticBlock:
+        """
+        Move the block to the next checkpoint if it exists.
+        If some middle checkpoints have been pruned, we jump to the smallest
+        sequence > current_seq that remains.
+        """
+        with self.session_maker() as session:
+            block = (
+                session.merge(use_preloaded_block)
+                if use_preloaded_block
+                else BlockModel.read(db_session=session, identifier=block_id, actor=actor)
+            )
+
+            if not block.current_history_entry_id:
+                raise ValueError(f"Block {block_id} has no history entry - cannot redo.")
+
+            current_entry = session.get(BlockHistory, block.current_history_entry_id)
+            if not current_entry:
+                raise NoResultFound(f"BlockHistory row not found for id={block.current_history_entry_id}")
+
+            current_seq = current_entry.sequence_number
+
+            # Find the smallest sequence that is > current_seq
+            next_entry = (
+                session.query(BlockHistory)
+                .filter(BlockHistory.block_id == block.id, BlockHistory.sequence_number > current_seq)
+                .order_by(BlockHistory.sequence_number.asc())
+                .first()
+            )
+            if not next_entry:
+                raise ValueError(f"Block {block_id} is at the highest checkpoint (seq={current_seq}). Cannot redo further.")
+
+            block = self._move_block_to_sequence(session, block, next_entry.sequence_number, actor)
+
+            session.commit()
+            return block.to_pydantic()

letta/services/group_manager.py

@@ -71,11 +71,20 @@ class GroupManager:
                 case ManagerType.supervisor:
                     new_group.manager_type = ManagerType.supervisor
                     new_group.manager_agent_id = group.manager_config.manager_agent_id
+                case ManagerType.sleeptime:
+                    new_group.manager_type = ManagerType.sleeptime
+                    new_group.manager_agent_id = group.manager_config.manager_agent_id
+                    new_group.sleeptime_agent_frequency = group.manager_config.sleeptime_agent_frequency
+                    if new_group.sleeptime_agent_frequency:
+                        new_group.turns_counter = 0
                 case _:
                     raise ValueError(f"Unsupported manager type: {group.manager_config.manager_type}")
 
             self._process_agent_relationship(session=session, group=new_group, agent_ids=group.agent_ids, allow_partial=False)
 
+            if group.shared_block_ids:
+                self._process_shared_block_relationship(session=session, group=new_group, block_ids=group.shared_block_ids)
+
             new_group.create(session, actor=actor)
             return new_group.to_pydantic()
 
@@ -84,6 +93,7 @@ class GroupManager:
         with self.session_maker() as session:
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
 
+            sleeptime_agent_frequency = None
             max_turns = None
             termination_token = None
             manager_agent_id = None
@@ -99,9 +109,16 @@ class GroupManager:
                         termination_token = group_update.manager_config.termination_token
                     case ManagerType.supervisor:
                         manager_agent_id = group_update.manager_config.manager_agent_id
+                    case ManagerType.sleeptime:
+                        manager_agent_id = group_update.manager_config.manager_agent_id
+                        sleeptime_agent_frequency = group_update.manager_config.sleeptime_agent_frequency
+                        if sleeptime_agent_frequency and group.turns_counter is None:
+                            group.turns_counter = 0
                     case _:
                         raise ValueError(f"Unsupported manager type: {group_update.manager_config.manager_type}")
 
+            if sleeptime_agent_frequency:
+                group.sleeptime_agent_frequency = sleeptime_agent_frequency
             if max_turns:
                 group.max_turns = max_turns
             if termination_token:
@@ -174,6 +191,30 @@ class GroupManager:
 
             session.commit()
 
+    @enforce_types
+    def bump_turns_counter(self, group_id: str, actor: PydanticUser) -> int:
+        with self.session_maker() as session:
+            # Ensure group is loadable by user
+            group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
+
+            # Update turns counter
+            group.turns_counter = (group.turns_counter + 1) % group.sleeptime_agent_frequency
+            group.update(session, actor=actor)
+            return group.turns_counter
+
+    @enforce_types
+    def get_last_processed_message_id_and_update(self, group_id: str, last_processed_message_id: str, actor: PydanticUser) -> str:
+        with self.session_maker() as session:
+            # Ensure group is loadable by user
+            group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
+
+            # Update last processed message id
+            prev_last_processed_message_id = group.last_processed_message_id
+            group.last_processed_message_id = last_processed_message_id
+            group.update(session, actor=actor)
+
+            return prev_last_processed_message_id
+
     def _process_agent_relationship(self, session: Session, group: GroupModel, agent_ids: List[str], allow_partial=False, replace=True):
         if not agent_ids:
             if replace:
@@ -203,3 +244,30 @@ class GroupManager:
             setattr(group, "agent_ids", agent_ids)
         else:
             raise ValueError("Extend relationship is not supported for groups.")
+
+    def _process_shared_block_relationship(
+        self,
+        session: Session,
+        group: GroupModel,
+        block_ids: List[str],
+    ):
+        """Process shared block relationships for a group and its agents."""
+        from letta.orm import Agent, Block, BlocksAgents
+
+        # Add blocks to group
+        blocks = session.query(Block).filter(Block.id.in_(block_ids)).all()
+        group.shared_blocks = blocks
+
+        # Add blocks to all agents
+        if group.agent_ids:
+            agents = session.query(Agent).filter(Agent.id.in_(group.agent_ids)).all()
+            for agent in agents:
+                for block in blocks:
+                    session.add(BlocksAgents(agent_id=agent.id, block_id=block.id, block_label=block.label))
+
+        # Add blocks to manager agent if exists
+        if group.manager_agent_id:
+            manager_agent = session.query(Agent).filter(Agent.id == group.manager_agent_id).first()
+            if manager_agent:
+                for block in blocks:
+                    session.add(BlocksAgents(agent_id=manager_agent.id, block_id=block.id, block_label=block.label))

letta/services/helpers/agent_manager_helper.py

@@ -89,13 +89,15 @@ def _process_tags(agent: AgentModel, tags: List[str], replace=True):
         agent.tags.extend([tag for tag in new_tags if tag.tag not in existing_tags])
 
 
-def derive_system_message(agent_type: AgentType, system: Optional[str] = None):
+def derive_system_message(agent_type: AgentType, enable_sleeptime: Optional[bool] = None, system: Optional[str] = None):
     if system is None:
         # TODO: don't hardcode
-        if agent_type == AgentType.memgpt_agent:
+        if agent_type == AgentType.memgpt_agent and not enable_sleeptime:
             system = gpt_system.get_system_text("memgpt_chat")
-        elif agent_type == AgentType.offline_memory_agent:
-            system = gpt_system.get_system_text("memgpt_offline_memory")
+        elif agent_type == AgentType.memgpt_agent and enable_sleeptime:
+            system = gpt_system.get_system_text("memgpt_sleeptime_chat")
+        elif agent_type == AgentType.sleeptime_agent:
+            system = gpt_system.get_system_text("sleeptime")
         else:
             raise ValueError(f"Invalid agent type: {agent_type}")
 

letta/services/llm_batch_manager.py

@@ -0,0 +1,139 @@
+import datetime
+from typing import Optional
+
+from anthropic.types.beta.messages import BetaMessageBatch, BetaMessageBatchIndividualResponse
+
+from letta.log import get_logger
+from letta.orm.llm_batch_items import LLMBatchItem
+from letta.orm.llm_batch_job import LLMBatchJob
+from letta.schemas.agent import AgentStepState
+from letta.schemas.enums import AgentStepStatus, JobStatus
+from letta.schemas.llm_batch_job import LLMBatchItem as PydanticLLMBatchItem
+from letta.schemas.llm_batch_job import LLMBatchJob as PydanticLLMBatchJob
+from letta.schemas.llm_config import LLMConfig
+from letta.schemas.user import User as PydanticUser
+from letta.utils import enforce_types
+
+logger = get_logger(__name__)
+
+
+class LLMBatchManager:
+    """Manager for handling both LLMBatchJob and LLMBatchItem operations."""
+
+    def __init__(self):
+        from letta.server.db import db_context
+
+        self.session_maker = db_context
+
+    @enforce_types
+    def create_batch_request(
+        self,
+        llm_provider: str,
+        create_batch_response: BetaMessageBatch,
+        actor: PydanticUser,
+        status: JobStatus = JobStatus.created,
+    ) -> PydanticLLMBatchJob:
+        """Create a new LLM batch job."""
+        with self.session_maker() as session:
+            batch = LLMBatchJob(
+                status=status,
+                llm_provider=llm_provider,
+                create_batch_response=create_batch_response,
+                organization_id=actor.organization_id,
+            )
+            batch.create(session, actor=actor)
+            return batch.to_pydantic()
+
+    @enforce_types
+    def get_batch_request_by_id(self, batch_id: str, actor: PydanticUser) -> PydanticLLMBatchJob:
+        """Retrieve a single batch job by ID."""
+        with self.session_maker() as session:
+            batch = LLMBatchJob.read(db_session=session, identifier=batch_id, actor=actor)
+            return batch.to_pydantic()
+
+    @enforce_types
+    def update_batch_status(
+        self,
+        batch_id: str,
+        status: JobStatus,
+        actor: PydanticUser,
+        latest_polling_response: Optional[BetaMessageBatch] = None,
+    ) -> PydanticLLMBatchJob:
+        """Update a batch job’s status and optionally its polling response."""
+        with self.session_maker() as session:
+            batch = LLMBatchJob.read(db_session=session, identifier=batch_id, actor=actor)
+            batch.status = status
+            batch.latest_polling_response = latest_polling_response
+            batch.last_polled_at = datetime.datetime.now(datetime.timezone.utc)
+            return batch.update(db_session=session, actor=actor).to_pydantic()
+
+    @enforce_types
+    def delete_batch_request(self, batch_id: str, actor: PydanticUser) -> None:
+        """Hard delete a batch job by ID."""
+        with self.session_maker() as session:
+            batch = LLMBatchJob.read(db_session=session, identifier=batch_id, actor=actor)
+            batch.hard_delete(db_session=session, actor=actor)
+
+    @enforce_types
+    def create_batch_item(
+        self,
+        batch_id: str,
+        agent_id: str,
+        llm_config: LLMConfig,
+        actor: PydanticUser,
+        request_status: JobStatus = JobStatus.created,
+        step_status: AgentStepStatus = AgentStepStatus.paused,
+        step_state: Optional[AgentStepState] = None,
+    ) -> PydanticLLMBatchItem:
+        """Create a new batch item."""
+        with self.session_maker() as session:
+            item = LLMBatchItem(
+                batch_id=batch_id,
+                agent_id=agent_id,
+                llm_config=llm_config,
+                request_status=request_status,
+                step_status=step_status,
+                step_state=step_state,
+                organization_id=actor.organization_id,
+            )
+            item.create(session, actor=actor)
+            return item.to_pydantic()
+
+    @enforce_types
+    def get_batch_item_by_id(self, item_id: str, actor: PydanticUser) -> PydanticLLMBatchItem:
+        """Retrieve a single batch item by ID."""
+        with self.session_maker() as session:
+            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
+            return item.to_pydantic()
+
+    @enforce_types
+    def update_batch_item(
+        self,
+        item_id: str,
+        actor: PydanticUser,
+        request_status: Optional[JobStatus] = None,
+        step_status: Optional[AgentStepStatus] = None,
+        llm_request_response: Optional[BetaMessageBatchIndividualResponse] = None,
+        step_state: Optional[AgentStepState] = None,
+    ) -> PydanticLLMBatchItem:
+        """Update fields on a batch item."""
+        with self.session_maker() as session:
+            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
+
+            if request_status:
+                item.request_status = request_status
+            if step_status:
+                item.step_status = step_status
+            if llm_request_response:
+                item.batch_request_result = llm_request_response
+            if step_state:
+                item.step_state = step_state
+
+            return item.update(db_session=session, actor=actor).to_pydantic()
+
+    @enforce_types
+    def delete_batch_item(self, item_id: str, actor: PydanticUser) -> None:
+        """Hard delete a batch item by ID."""
+        with self.session_maker() as session:
+            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
+            item.hard_delete(db_session=session, actor=actor)

letta/services/message_manager.py

@@ -1,7 +1,7 @@
 import json
 from typing import List, Optional, Sequence
 
-from sqlalchemy import and_, exists, func, or_, select, text
+from sqlalchemy import exists, func, select, text
 
 from letta.log import get_logger
 from letta.orm.agent import Agent as AgentModel
@@ -270,19 +270,20 @@ class MessageManager:
         Most performant query to list messages for an agent by directly querying the Message table.
 
         This function filters by the agent_id (leveraging the index on messages.agent_id)
-        and applies efficient pagination using (created_at, id) as the cursor.
+        and applies pagination using sequence_id as the cursor.
         If query_text is provided, it will filter messages whose text content partially matches the query.
         If role is provided, it will filter messages by the specified role.
 
         Args:
             agent_id: The ID of the agent whose messages are queried.
             actor: The user performing the action (used for permission checks).
-            after: A message ID; if provided, only messages *after* this message (per sort order) are returned.
-            before: A message ID; if provided, only messages *before* this message are returned.
+            after: A message ID; if provided, only messages *after* this message (by sequence_id) are returned.
+            before: A message ID; if provided, only messages *before* this message (by sequence_id) are returned.
             query_text: Optional string to partially match the message text content.
             roles: Optional MessageRole to filter messages by role.
             limit: Maximum number of messages to return.
-            ascending: If True, sort by (created_at, id) ascending; if False, sort descending.
+            ascending: If True, sort by sequence_id ascending; if False, sort descending.
+            group_id: Optional group ID to filter messages by group_id.
 
         Returns:
             List[PydanticMessage]: A list of messages (converted via .to_pydantic()).
@@ -290,6 +291,7 @@ class MessageManager:
         Raises:
             NoResultFound: If the provided after/before message IDs do not exist.
         """
+
         with self.session_maker() as session:
             # Permission check: raise if the agent doesn't exist or actor is not allowed.
             AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
@@ -301,7 +303,7 @@ class MessageManager:
             if group_id:
                 query = query.filter(MessageModel.group_id == group_id)
 
-            # If query_text is provided, filter messages using subquery.
+            # If query_text is provided, filter messages using subquery + json_array_elements.
             if query_text:
                 content_element = func.json_array_elements(MessageModel.content).alias("content_element")
                 query = query.filter(
@@ -313,48 +315,32 @@ class MessageManager:
                     )
                 )
 
-            # If role is provided, filter messages by role.
+            # If role(s) are provided, filter messages by those roles.
             if roles:
                 role_values = [r.value for r in roles]
                 query = query.filter(MessageModel.role.in_(role_values))
 
             # Apply 'after' pagination if specified.
             if after:
-                after_ref = session.query(MessageModel.created_at, MessageModel.id).filter(MessageModel.id == after).limit(1).one_or_none()
+                after_ref = session.query(MessageModel.sequence_id).filter(MessageModel.id == after).one_or_none()
                 if not after_ref:
                     raise NoResultFound(f"No message found with id '{after}' for agent '{agent_id}'.")
-                query = query.filter(
-                    or_(
-                        MessageModel.created_at > after_ref.created_at,
-                        and_(
-                            MessageModel.created_at == after_ref.created_at,
-                            MessageModel.id > after_ref.id,
-                        ),
-                    )
-                )
+                # Filter out any messages with a sequence_id <= after_ref.sequence_id
+                query = query.filter(MessageModel.sequence_id > after_ref.sequence_id)
 
             # Apply 'before' pagination if specified.
             if before:
-                before_ref = (
-                    session.query(MessageModel.created_at, MessageModel.id).filter(MessageModel.id == before).limit(1).one_or_none()
-                )
+                before_ref = session.query(MessageModel.sequence_id).filter(MessageModel.id == before).one_or_none()
                 if not before_ref:
                     raise NoResultFound(f"No message found with id '{before}' for agent '{agent_id}'.")
-                query = query.filter(
-                    or_(
-                        MessageModel.created_at < before_ref.created_at,
-                        and_(
-                            MessageModel.created_at == before_ref.created_at,
-                            MessageModel.id < before_ref.id,
-                        ),
-                    )
-                )
+                # Filter out any messages with a sequence_id >= before_ref.sequence_id
+                query = query.filter(MessageModel.sequence_id < before_ref.sequence_id)
 
             # Apply ordering based on the ascending flag.
             if ascending:
-                query = query.order_by(MessageModel.created_at.asc(), MessageModel.id.asc())
+                query = query.order_by(MessageModel.sequence_id.asc())
             else:
-                query = query.order_by(MessageModel.created_at.desc(), MessageModel.id.desc())
+                query = query.order_by(MessageModel.sequence_id.desc())
 
             # Limit the number of results.
             query = query.limit(limit)

letta/services/tool_executor/tool_execution_sandbox.py

@@ -387,9 +387,7 @@ class ToolExecutionSandbox:
             sbx = Sandbox(sandbox_config.get_e2b_config().template, metadata={self.METADATA_CONFIG_STATE_KEY: state_hash})
         else:
             # no template
-            sbx = Sandbox(
-                metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(to_orm=True, exclude={"pip_requirements"})
-            )
+            sbx = Sandbox(metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(exclude={"pip_requirements"}))
 
         # install pip requirements
         if e2b_config.pip_requirements:

letta/services/tool_executor/tool_executor.py

@@ -1,4 +1,3 @@
-import ast
 import math
 from abc import ABC, abstractmethod
 from typing import Any, Optional, Tuple
@@ -15,7 +14,9 @@ from letta.schemas.user import User
 from letta.services.agent_manager import AgentManager
 from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
-from letta.services.tool_executor.async_tool_execution_sandbox import AsyncToolExecutionSandbox
+from letta.services.tool_sandbox.e2b_sandbox import AsyncToolSandboxE2B
+from letta.services.tool_sandbox.local_sandbox import AsyncToolSandboxLocal
+from letta.settings import tool_settings
 from letta.utils import get_friendly_error_msg
 
 
@@ -40,6 +41,7 @@ class LettaCoreToolExecutor(ToolExecutor):
             "send_message": self.send_message,
             "conversation_search": self.conversation_search,
             "archival_memory_search": self.archival_memory_search,
+            "archival_memory_insert": self.archival_memory_insert,
         }
 
         if function_name not in function_map:
@@ -334,16 +336,13 @@ class SandboxToolExecutor(ToolExecutor):
 
             agent_state_copy = self._create_agent_state_copy(agent_state)
 
-            # TODO: This is brittle, think about better way to do this?
-            if "agent_state" in self.parse_function_arguments(tool.source_code, tool.name):
-                inject_agent_state = True
+            # Execute in sandbox depending on API key
+            if tool_settings.e2b_api_key:
+                sandbox = AsyncToolSandboxE2B(function_name, function_args, actor, tool_object=tool)
             else:
-                inject_agent_state = False
+                sandbox = AsyncToolSandboxLocal(function_name, function_args, actor, tool_object=tool)
 
-            # Execute in sandbox
-            sandbox_run_result = await AsyncToolExecutionSandbox(function_name, function_args, actor, tool_object=tool).run(
-                agent_state=agent_state_copy, inject_agent_state=inject_agent_state
-            )
+            sandbox_run_result = await sandbox.run(agent_state=agent_state_copy)
 
             function_response, updated_agent_state = sandbox_run_result.func_return, sandbox_run_result.agent_state
 
@@ -371,16 +370,6 @@ class SandboxToolExecutor(ToolExecutor):
             # This is defensive programming - we try to coerce but fall back if it fails
             return function_args
 
-    def parse_function_arguments(self, source_code: str, tool_name: str):
-        """Get arguments of a function from its source code"""
-        tree = ast.parse(source_code)
-        args = []
-        for node in ast.walk(tree):
-            if isinstance(node, ast.FunctionDef) and node.name == tool_name:
-                for arg in node.args.args:
-                    args.append(arg.arg)
-        return args
-
     def _create_agent_state_copy(self, agent_state: AgentState):
         """Create a copy of agent state for sandbox execution."""
         agent_state_copy = agent_state.__deepcopy__()