letta-nightly 0.6.53.dev20250418104238__py3-none-any.whl → 0.6.54.dev20250419194309__py3-none-any.whl

letta/services/llm_batch_manager.py

@@ -1,10 +1,10 @@
 import datetime
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from anthropic.types.beta.messages import BetaMessageBatch, BetaMessageBatchIndividualResponse
-from sqlalchemy import tuple_
+from sqlalchemy import func, tuple_
 
-from letta.jobs.types import BatchPollingResult, ItemUpdateInfo
+from letta.jobs.types import BatchPollingResult, ItemUpdateInfo, RequestStatusUpdateInfo, StepStatusUpdateInfo
 from letta.log import get_logger
 from letta.orm.llm_batch_items import LLMBatchItem
 from letta.orm.llm_batch_job import LLMBatchJob
@@ -28,11 +28,12 @@ class LLMBatchManager:
         self.session_maker = db_context
 
     @enforce_types
-    def create_batch_job(
+    def create_llm_batch_job(
         self,
         llm_provider: ProviderType,
         create_batch_response: BetaMessageBatch,
         actor: PydanticUser,
+        letta_batch_job_id: str,
         status: JobStatus = JobStatus.created,
     ) -> PydanticLLMBatchJob:
         """Create a new LLM batch job."""
@@ -42,51 +43,52 @@ class LLMBatchManager:
                 llm_provider=llm_provider,
                 create_batch_response=create_batch_response,
                 organization_id=actor.organization_id,
+                letta_batch_job_id=letta_batch_job_id,
             )
             batch.create(session, actor=actor)
             return batch.to_pydantic()
 
     @enforce_types
-    def get_batch_job_by_id(self, batch_id: str, actor: Optional[PydanticUser] = None) -> PydanticLLMBatchJob:
+    def get_llm_batch_job_by_id(self, llm_batch_id: str, actor: Optional[PydanticUser] = None) -> 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)
+            batch = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
             return batch.to_pydantic()
 
     @enforce_types
-    def update_batch_status(
+    def update_llm_batch_status(
         self,
-        batch_id: str,
+        llm_batch_id: str,
         status: JobStatus,
         actor: Optional[PydanticUser] = None,
         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 = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
             batch.status = status
             batch.latest_polling_response = latest_polling_response
             batch.last_polled_at = datetime.datetime.now(datetime.timezone.utc)
             batch = batch.update(db_session=session, actor=actor)
             return batch.to_pydantic()
 
-    def bulk_update_batch_statuses(
+    def bulk_update_llm_batch_statuses(
         self,
         updates: List[BatchPollingResult],
     ) -> None:
         """
         Efficiently update many LLMBatchJob rows. This is used by the cron jobs.
 
-        `updates` = [(batch_id, new_status, polling_response_or_None), …]
+        `updates` = [(llm_batch_id, new_status, polling_response_or_None), …]
         """
         now = datetime.datetime.now(datetime.timezone.utc)
 
         with self.session_maker() as session:
             mappings = []
-            for batch_id, status, response in updates:
+            for llm_batch_id, status, response in updates:
                 mappings.append(
                     {
-                        "id": batch_id,
+                        "id": llm_batch_id,
                         "status": status,
                         "latest_polling_response": response,
                         "last_polled_at": now,
@@ -97,14 +99,51 @@ class LLMBatchManager:
             session.commit()
 
     @enforce_types
-    def delete_batch_request(self, batch_id: str, actor: PydanticUser) -> None:
+    def list_llm_batch_jobs(
+        self,
+        letta_batch_id: str,
+        limit: Optional[int] = None,
+        actor: Optional[PydanticUser] = None,
+        after: Optional[str] = None,
+    ) -> List[PydanticLLMBatchItem]:
+        """
+        List all batch items for a given llm_batch_id, optionally filtered by additional criteria and limited in count.
+
+        Optional filters:
+            - after: A cursor string. Only items with an `id` greater than this value are returned.
+            - agent_id: Restrict the result set to a specific agent.
+            - request_status: Filter items based on their request status (e.g., created, completed, expired).
+            - step_status: Filter items based on their step execution status.
+
+        The results are ordered by their id in ascending order.
+        """
+        with self.session_maker() as session:
+            query = session.query(LLMBatchJob).filter(LLMBatchJob.letta_batch_job_id == letta_batch_id)
+
+            if actor is not None:
+                query = query.filter(LLMBatchJob.organization_id == actor.organization_id)
+
+            # Additional optional filters
+            if after is not None:
+                query = query.filter(LLMBatchJob.id > after)
+
+            query = query.order_by(LLMBatchJob.id.asc())
+
+            if limit is not None:
+                query = query.limit(limit)
+
+            results = query.all()
+            return [item.to_pydantic() for item in results]
+
+    @enforce_types
+    def delete_llm_batch_request(self, llm_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 = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
             batch.hard_delete(db_session=session, actor=actor)
 
     @enforce_types
-    def list_running_batches(self, actor: Optional[PydanticUser] = None) -> List[PydanticLLMBatchJob]:
+    def list_running_llm_batches(self, actor: Optional[PydanticUser] = None) -> List[PydanticLLMBatchJob]:
         """Return all running LLM batch jobs, optionally filtered by actor's organization."""
         with self.session_maker() as session:
             query = session.query(LLMBatchJob).filter(LLMBatchJob.status == JobStatus.running)
@@ -116,9 +155,9 @@ class LLMBatchManager:
             return [batch.to_pydantic() for batch in results]
 
     @enforce_types
-    def create_batch_item(
+    def create_llm_batch_item(
         self,
-        batch_id: str,
+        llm_batch_id: str,
         agent_id: str,
         llm_config: LLMConfig,
         actor: PydanticUser,
@@ -129,7 +168,7 @@ class LLMBatchManager:
         """Create a new batch item."""
         with self.session_maker() as session:
             item = LLMBatchItem(
-                batch_id=batch_id,
+                llm_batch_id=llm_batch_id,
                 agent_id=agent_id,
                 llm_config=llm_config,
                 request_status=request_status,
@@ -141,14 +180,47 @@ class LLMBatchManager:
             return item.to_pydantic()
 
     @enforce_types
-    def get_batch_item_by_id(self, item_id: str, actor: PydanticUser) -> PydanticLLMBatchItem:
+    def create_llm_batch_items_bulk(self, llm_batch_items: List[PydanticLLMBatchItem], actor: PydanticUser) -> List[PydanticLLMBatchItem]:
+        """
+        Create multiple batch items in bulk for better performance.
+
+        Args:
+            llm_batch_items: List of batch items to create
+            actor: User performing the action
+
+        Returns:
+            List of created batch items as Pydantic models
+        """
+        with self.session_maker() as session:
+            # Convert Pydantic models to ORM objects
+            orm_items = []
+            for item in llm_batch_items:
+                orm_item = LLMBatchItem(
+                    llm_batch_id=item.llm_batch_id,
+                    agent_id=item.agent_id,
+                    llm_config=item.llm_config,
+                    request_status=item.request_status,
+                    step_status=item.step_status,
+                    step_state=item.step_state,
+                    organization_id=actor.organization_id,
+                )
+                orm_items.append(orm_item)
+
+            # Use the batch_create method to create all items at once
+            created_items = LLMBatchItem.batch_create(orm_items, session, actor=actor)
+
+            # Convert back to Pydantic models
+            return [item.to_pydantic() for item in created_items]
+
+    @enforce_types
+    def get_llm_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(
+    def update_llm_batch_item(
         self,
         item_id: str,
         actor: PydanticUser,
@@ -173,78 +245,149 @@ class LLMBatchManager:
             return item.update(db_session=session, actor=actor).to_pydantic()
 
     @enforce_types
-    def list_batch_items(
+    def list_llm_batch_items(
         self,
-        batch_id: str,
+        llm_batch_id: str,
         limit: Optional[int] = None,
         actor: Optional[PydanticUser] = None,
+        after: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        request_status: Optional[JobStatus] = None,
+        step_status: Optional[AgentStepStatus] = None,
     ) -> List[PydanticLLMBatchItem]:
-        """List all batch items for a given batch_id, optionally filtered by organization and limited in count."""
+        """
+        List all batch items for a given llm_batch_id, optionally filtered by additional criteria and limited in count.
+
+        Optional filters:
+            - after: A cursor string. Only items with an `id` greater than this value are returned.
+            - agent_id: Restrict the result set to a specific agent.
+            - request_status: Filter items based on their request status (e.g., created, completed, expired).
+            - step_status: Filter items based on their step execution status.
+
+        The results are ordered by their id in ascending order.
+        """
         with self.session_maker() as session:
-            query = session.query(LLMBatchItem).filter(LLMBatchItem.batch_id == batch_id)
+            query = session.query(LLMBatchItem).filter(LLMBatchItem.llm_batch_id == llm_batch_id)
 
             if actor is not None:
                 query = query.filter(LLMBatchItem.organization_id == actor.organization_id)
 
-            if limit:
+            # Additional optional filters
+            if agent_id is not None:
+                query = query.filter(LLMBatchItem.agent_id == agent_id)
+            if request_status is not None:
+                query = query.filter(LLMBatchItem.request_status == request_status)
+            if step_status is not None:
+                query = query.filter(LLMBatchItem.step_status == step_status)
+            if after is not None:
+                query = query.filter(LLMBatchItem.id > after)
+
+            query = query.order_by(LLMBatchItem.id.asc())
+
+            if limit is not None:
                 query = query.limit(limit)
 
             results = query.all()
             return [item.to_pydantic() for item in results]
 
-    def bulk_update_batch_items_by_agent(
+    def bulk_update_llm_batch_items(
         self,
-        updates: List[ItemUpdateInfo],
+        llm_batch_id_agent_id_pairs: List[Tuple[str, str]],
+        field_updates: List[Dict[str, Any]],
     ) -> None:
         """
-        Efficiently update LLMBatchItem rows by (batch_id, agent_id).
+        Efficiently update multiple LLMBatchItem rows by (llm_batch_id, agent_id) pairs.
 
         Args:
-            updates: List of tuples:
-              (batch_id, agent_id, new_request_status, batch_request_result)
+            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
         """
-        with self.session_maker() as session:
-            # For bulk_update_mappings, we need the primary key of each row
-            # So we must map (batch_id, agent_id) → actual PK (id)
-            # We'll do it in one DB query using the (batch_id, agent_id) sets
+        if not llm_batch_id_agent_id_pairs or not field_updates:
+            return
 
-            # 1. Gather the pairs
-            key_pairs = [(b_id, a_id) for (b_id, a_id, *_rest) in updates]
+        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")
 
-            # 2. Query items in a single step
+        with self.session_maker() as session:
+            # Lookup primary keys
             items = (
-                session.query(LLMBatchItem.id, LLMBatchItem.batch_id, LLMBatchItem.agent_id)
-                .filter(tuple_(LLMBatchItem.batch_id, LLMBatchItem.agent_id).in_(key_pairs))
+                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}
 
-            # Build a map from (batch_id, agent_id) → PK id
-            pair_to_pk = {}
-            for row_id, row_batch_id, row_agent_id in items:
-                pair_to_pk[(row_batch_id, row_agent_id)] = row_id
-
-            # 3. Construct mappings for the PK-based bulk update
             mappings = []
-            for batch_id, agent_id, new_status, new_result in updates:
-                pk_id = pair_to_pk.get((batch_id, agent_id))
+            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:
-                    # Nonexistent or mismatch → skip
                     continue
-                mappings.append(
-                    {
-                        "id": pk_id,
-                        "request_status": new_status,
-                        "batch_request_result": new_result,
-                    }
-                )
+
+                update_fields = fields.copy()
+                update_fields["id"] = pk_id
+                mappings.append(update_fields)
 
             if mappings:
                 session.bulk_update_mappings(LLMBatchItem, mappings)
                 session.commit()
 
     @enforce_types
-    def delete_batch_item(self, item_id: str, actor: PydanticUser) -> None:
+    def bulk_update_batch_llm_items_results_by_agent(
+        self,
+        updates: List[ItemUpdateInfo],
+    ) -> 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 = [
+            {
+                "request_status": update.request_status,
+                "batch_request_result": update.batch_request_result,
+            }
+            for update in updates
+        ]
+
+        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates)
+
+    @enforce_types
+    def bulk_update_llm_batch_items_step_status_by_agent(
+        self,
+        updates: List[StepStatusUpdateInfo],
+    ) -> 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)
+
+    @enforce_types
+    def bulk_update_llm_batch_items_request_status_by_agent(
+        self,
+        updates: List[RequestStatusUpdateInfo],
+    ) -> 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)
+
+    @enforce_types
+    def delete_llm_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)
+
+    @enforce_types
+    def count_llm_batch_items(self, llm_batch_id: str) -> int:
+        """
+        Efficiently count the number of batch items for a given llm_batch_id.
+
+        Args:
+            llm_batch_id (str): The batch identifier to count items for.
+
+        Returns:
+            int: The total number of batch items associated with the given llm_batch_id.
+        """
+        with self.session_maker() as session:
+            count = session.query(func.count(LLMBatchItem.id)).filter(LLMBatchItem.llm_batch_id == llm_batch_id).scalar()
+            return count or 0

letta/services/message_manager.py

@@ -1,7 +1,7 @@
 import json
 from typing import List, Optional, Sequence
 
-from sqlalchemy import exists, func, select, text
+from sqlalchemy import delete, exists, func, select, text
 
 from letta.log import get_logger
 from letta.orm.agent import Agent as AgentModel
@@ -348,3 +348,25 @@ class MessageManager:
             # Execute and convert each Message to its Pydantic representation.
             results = query.all()
             return [msg.to_pydantic() for msg in results]
+
+    @enforce_types
+    def delete_all_messages_for_agent(self, agent_id: str, actor: PydanticUser) -> int:
+        """
+        Efficiently deletes all messages associated with a given agent_id,
+        while enforcing permission checks and avoiding any ORM‑level loads.
+        """
+        with self.session_maker() as session:
+            # 1) verify the agent exists and the actor has access
+            AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+
+            # 2) issue a CORE DELETE against the mapped class
+            stmt = (
+                delete(MessageModel).where(MessageModel.agent_id == agent_id).where(MessageModel.organization_id == actor.organization_id)
+            )
+            result = session.execute(stmt)
+
+            # 3) commit once
+            session.commit()
+
+            # 4) return the number of rows deleted
+            return result.rowcount

letta/services/passage_manager.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import List, Optional
 
 from openai import OpenAI
@@ -49,7 +49,7 @@ class PassageManager:
             "organization_id": data["organization_id"],
             "metadata_": data.get("metadata", {}),
             "is_deleted": data.get("is_deleted", False),
-            "created_at": data.get("created_at", datetime.utcnow()),
+            "created_at": data.get("created_at", datetime.now(timezone.utc)),
         }
 
         if "agent_id" in data and data["agent_id"]:

letta/services/tool_executor/tool_execution_manager.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, Optional, Tuple, 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 SandboxRunResult
+from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult
 from letta.schemas.tool import Tool
 from letta.schemas.user import User
 from letta.services.tool_executor.tool_executor import (
@@ -45,10 +45,18 @@ class ToolExecutorFactory:
 class ToolExecutionManager:
     """Manager class for tool execution operations."""
 
-    def __init__(self, agent_state: AgentState, actor: User):
+    def __init__(
+        self,
+        agent_state: AgentState,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
+    ):
         self.agent_state = agent_state
         self.logger = get_logger(__name__)
         self.actor = actor
+        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]]:
         """
@@ -67,7 +75,9 @@ class ToolExecutionManager:
             executor = ToolExecutorFactory.get_executor(tool.tool_type)
 
             # Execute the tool
-            return executor.execute(function_name, function_args, self.agent_state, tool, self.actor)
+            return executor.execute(
+                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)}")

letta/services/tool_executor/tool_execution_sandbox.py

@@ -148,7 +148,6 @@ class ToolExecutionSandbox:
             temp_file.write(code)
             temp_file.flush()
             temp_file_path = temp_file.name
-
         try:
             if local_configs.use_venv:
                 return self.run_local_dir_sandbox_venv(sbx_config, env, temp_file_path)

letta/services/tool_executor/tool_executor.py

@@ -1,6 +1,6 @@
 import math
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 from letta.constants import COMPOSIO_ENTITY_ENV_VAR_KEY, RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
 from letta.functions.ast_parsers import coerce_dict_args_by_annotations, get_function_annotations_from_source
@@ -8,7 +8,7 @@ from letta.functions.helpers import execute_composio_action, generate_composio_a
 from letta.helpers.composio_helpers import get_composio_api_key
 from letta.helpers.json_helpers import json_dumps
 from letta.schemas.agent import AgentState
-from letta.schemas.sandbox_config import SandboxRunResult
+from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult
 from letta.schemas.tool import Tool
 from letta.schemas.user import User
 from letta.services.agent_manager import AgentManager
@@ -25,7 +25,14 @@ class ToolExecutor(ABC):
 
     @abstractmethod
     def execute(
-        self, function_name: str, function_args: dict, agent_state: AgentState, tool: Tool, actor: User
+        self,
+        function_name: str,
+        function_args: dict,
+        agent_state: AgentState,
+        tool: Tool,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Any, Optional[SandboxRunResult]]:
         """Execute the tool and return the result."""
 
@@ -34,7 +41,14 @@ class LettaCoreToolExecutor(ToolExecutor):
     """Executor for LETTA core tools with direct implementation of functions."""
 
     def execute(
-        self, function_name: str, function_args: dict, agent_state: AgentState, tool: Tool, actor: User
+        self,
+        function_name: str,
+        function_args: dict,
+        agent_state: AgentState,
+        tool: Tool,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Any, Optional[SandboxRunResult]]:
         # Map function names to method calls
         function_map = {
@@ -184,7 +198,14 @@ class LettaMemoryToolExecutor(ToolExecutor):
     """Executor for LETTA memory core tools with direct implementation."""
 
     def execute(
-        self, function_name: str, function_args: dict, agent_state: AgentState, tool: Tool, actor: User
+        self,
+        function_name: str,
+        function_args: dict,
+        agent_state: AgentState,
+        tool: Tool,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Any, Optional[SandboxRunResult]]:
         # Map function names to method calls
         function_map = {
@@ -244,7 +265,14 @@ class ExternalComposioToolExecutor(ToolExecutor):
     """Executor for external Composio tools."""
 
     def execute(
-        self, function_name: str, function_args: dict, agent_state: AgentState, tool: Tool, actor: User
+        self,
+        function_name: str,
+        function_args: dict,
+        agent_state: AgentState,
+        tool: Tool,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Any, Optional[SandboxRunResult]]:
         action_name = generate_composio_action_from_func_name(tool.name)
 
@@ -324,7 +352,14 @@ class SandboxToolExecutor(ToolExecutor):
     """Executor for sandboxed tools."""
 
     async def execute(
-        self, function_name: str, function_args: dict, agent_state: AgentState, tool: Tool, actor: User
+        self,
+        function_name: str,
+        function_args: dict,
+        agent_state: AgentState,
+        tool: Tool,
+        actor: User,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
     ) -> Tuple[Any, Optional[SandboxRunResult]]:
 
         # Store original memory state
@@ -338,9 +373,13 @@ class SandboxToolExecutor(ToolExecutor):
 
             # Execute in sandbox depending on API key
             if tool_settings.e2b_api_key:
-                sandbox = AsyncToolSandboxE2B(function_name, function_args, actor, tool_object=tool)
+                sandbox = AsyncToolSandboxE2B(
+                    function_name, function_args, actor, tool_object=tool, sandbox_config=sandbox_config, sandbox_env_vars=sandbox_env_vars
+                )
             else:
-                sandbox = AsyncToolSandboxLocal(function_name, function_args, actor, tool_object=tool)
+                sandbox = AsyncToolSandboxLocal(
+                    function_name, function_args, actor, tool_object=tool, sandbox_config=sandbox_config, sandbox_env_vars=sandbox_env_vars
+                )
 
             sandbox_run_result = await sandbox.run(agent_state=agent_state_copy)
 

letta/services/tool_sandbox/base.py

@@ -7,9 +7,9 @@ from typing import Any, Dict, Optional, Tuple
 
 from letta.functions.helpers import generate_model_from_args_json_schema
 from letta.schemas.agent import AgentState
-from letta.schemas.sandbox_config import SandboxRunResult
+from letta.schemas.sandbox_config import SandboxConfig, SandboxRunResult
+from letta.schemas.tool import Tool
 from letta.services.helpers.tool_execution_helper import add_imports_and_pydantic_schemas_for_args
-from letta.services.organization_manager import OrganizationManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
 from letta.services.tool_manager import ToolManager
 
@@ -20,12 +20,18 @@ class AsyncToolSandboxBase(ABC):
     LOCAL_SANDBOX_RESULT_END_MARKER = str(uuid.uuid5(NAMESPACE, "local-sandbox-result-end-marker"))
     LOCAL_SANDBOX_RESULT_VAR_NAME = "result_ZQqiequkcFwRwwGQMqkt"
 
-    def __init__(self, tool_name: str, args: dict, user, tool_object=None):
+    def __init__(
+        self,
+        tool_name: str,
+        args: dict,
+        user,
+        tool_object: Optional[Tool] = None,
+        sandbox_config: Optional[SandboxConfig] = None,
+        sandbox_env_vars: Optional[Dict[str, Any]] = None,
+    ):
         self.tool_name = tool_name
         self.args = args
         self.user = user
-        self.organization = OrganizationManager().get_organization_by_id(self.user.organization_id)
-        self.privileged_tools = self.organization.privileged_tools
 
         self.tool = tool_object or ToolManager().get_tool_by_name(tool_name=tool_name, actor=self.user)
         if self.tool is None:
@@ -33,7 +39,12 @@ class AsyncToolSandboxBase(ABC):
                 f"Agent attempted to invoke tool {self.tool_name} that does not exist for organization {self.user.organization_id}"
             )
 
-        self.sandbox_config_manager = SandboxConfigManager()
+        # Store provided values or create manager to fetch them later
+        self.provided_sandbox_config = sandbox_config
+        self.provided_sandbox_env_vars = sandbox_env_vars
+
+        # Only create the manager if we need to (lazy initialization)
+        self._sandbox_config_manager = None
 
         # See if we should inject agent_state or not based on the presence of the "agent_state" arg
         if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
@@ -41,6 +52,13 @@ class AsyncToolSandboxBase(ABC):
         else:
             self.inject_agent_state = False
 
+    # Lazily initialize the manager only when needed
+    @property
+    def sandbox_config_manager(self):
+        if self._sandbox_config_manager is None:
+            self._sandbox_config_manager = SandboxConfigManager()
+        return self._sandbox_config_manager
+
     @abstractmethod
     async def run(
         self,