letta-nightly 0.10.0.dev20250806104523__py3-none-any.whl → 0.11.0.dev20250807104511__py3-none-any.whl

letta/services/passage_manager.py

@@ -9,14 +9,16 @@ from sqlalchemy import select
 from letta.constants import MAX_EMBEDDING_DIM
 from letta.embeddings import embedding_model, parse_and_chunk_text
 from letta.helpers.decorators import async_redis_cache
+from letta.orm import ArchivesAgents
 from letta.orm.errors import NoResultFound
-from letta.orm.passage import AgentPassage, SourcePassage
+from letta.orm.passage import ArchivalPassage, SourcePassage
 from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState
 from letta.schemas.file import FileMetadata as PydanticFileMetadata
 from letta.schemas.passage import Passage as PydanticPassage
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.services.archive_manager import ArchiveManager
 from letta.utils import enforce_types
 
 
@@ -42,6 +44,9 @@ async def get_openai_embedding_async(text: str, model: str, endpoint: str) -> li
 class PassageManager:
     """Manager class to handle business logic related to Passages."""
 
+    def __init__(self):
+        self.archive_manager = ArchiveManager()
+
     # AGENT PASSAGE METHODS
     @enforce_types
     @trace_method
@@ -49,7 +54,7 @@ class PassageManager:
         """Fetch an agent passage by ID."""
         with db_registry.session() as session:
             try:
-                passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                passage = ArchivalPassage.read(db_session=session, identifier=passage_id, actor=actor)
                 return passage.to_pydantic()
             except NoResultFound:
                 raise NoResultFound(f"Agent passage with id {passage_id} not found in database.")
@@ -60,7 +65,7 @@ class PassageManager:
         """Fetch an agent passage by ID."""
         async with db_registry.async_session() as session:
             try:
-                passage = await AgentPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
+                passage = await ArchivalPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
                 return passage.to_pydantic()
             except NoResultFound:
                 raise NoResultFound(f"Agent passage with id {passage_id} not found in database.")
@@ -109,7 +114,7 @@ class PassageManager:
             except NoResultFound:
                 # Try archival passages
                 try:
-                    passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                    passage = ArchivalPassage.read(db_session=session, identifier=passage_id, actor=actor)
                     return passage.to_pydantic()
                 except NoResultFound:
                     raise NoResultFound(f"Passage with id {passage_id} not found in database.")
@@ -134,7 +139,7 @@ class PassageManager:
             except NoResultFound:
                 # Try archival passages
                 try:
-                    passage = await AgentPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
+                    passage = await ArchivalPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
                     return passage.to_pydantic()
                 except NoResultFound:
                     raise NoResultFound(f"Passage with id {passage_id} not found in database.")
@@ -143,8 +148,8 @@ class PassageManager:
     @trace_method
     def create_agent_passage(self, pydantic_passage: PydanticPassage, actor: PydanticUser) -> PydanticPassage:
         """Create a new agent passage."""
-        if not pydantic_passage.agent_id:
-            raise ValueError("Agent passage must have agent_id")
+        if not pydantic_passage.archive_id:
+            raise ValueError("Agent passage must have archive_id")
         if pydantic_passage.source_id:
             raise ValueError("Agent passage cannot have source_id")
 
@@ -159,8 +164,8 @@ class PassageManager:
             "is_deleted": data.get("is_deleted", False),
             "created_at": data.get("created_at", datetime.now(timezone.utc)),
         }
-        agent_fields = {"agent_id": data["agent_id"]}
-        passage = AgentPassage(**common_fields, **agent_fields)
+        agent_fields = {"archive_id": data["archive_id"]}
+        passage = ArchivalPassage(**common_fields, **agent_fields)
 
         with db_registry.session() as session:
             passage.create(session, actor=actor)
@@ -170,8 +175,8 @@ class PassageManager:
     @trace_method
     async def create_agent_passage_async(self, pydantic_passage: PydanticPassage, actor: PydanticUser) -> PydanticPassage:
         """Create a new agent passage."""
-        if not pydantic_passage.agent_id:
-            raise ValueError("Agent passage must have agent_id")
+        if not pydantic_passage.archive_id:
+            raise ValueError("Agent passage must have archive_id")
         if pydantic_passage.source_id:
             raise ValueError("Agent passage cannot have source_id")
 
@@ -186,8 +191,8 @@ class PassageManager:
             "is_deleted": data.get("is_deleted", False),
             "created_at": data.get("created_at", datetime.now(timezone.utc)),
         }
-        agent_fields = {"agent_id": data["agent_id"]}
-        passage = AgentPassage(**common_fields, **agent_fields)
+        agent_fields = {"archive_id": data["archive_id"]}
+        passage = ArchivalPassage(**common_fields, **agent_fields)
 
         async with db_registry.async_session() as session:
             passage = await passage.create_async(session, actor=actor)
@@ -201,8 +206,8 @@ class PassageManager:
         """Create a new source passage."""
         if not pydantic_passage.source_id:
             raise ValueError("Source passage must have source_id")
-        if pydantic_passage.agent_id:
-            raise ValueError("Source passage cannot have agent_id")
+        if pydantic_passage.archive_id:
+            raise ValueError("Source passage cannot have archive_id")
 
         data = pydantic_passage.model_dump(to_orm=True)
         common_fields = {
@@ -234,8 +239,8 @@ class PassageManager:
         """Create a new source passage."""
         if not pydantic_passage.source_id:
             raise ValueError("Source passage must have source_id")
-        if pydantic_passage.agent_id:
-            raise ValueError("Source passage cannot have agent_id")
+        if pydantic_passage.archive_id:
+            raise ValueError("Source passage cannot have archive_id")
 
         data = pydantic_passage.model_dump(to_orm=True)
         common_fields = {
@@ -308,21 +313,21 @@ class PassageManager:
             "created_at": data.get("created_at", datetime.now(timezone.utc)),
         }
 
-        if "agent_id" in data and data["agent_id"]:
-            assert not data.get("source_id"), "Passage cannot have both agent_id and source_id"
+        if "archive_id" in data and data["archive_id"]:
+            assert not data.get("source_id"), "Passage cannot have both archive_id and source_id"
             agent_fields = {
-                "agent_id": data["agent_id"],
+                "archive_id": data["archive_id"],
             }
-            passage = AgentPassage(**common_fields, **agent_fields)
+            passage = ArchivalPassage(**common_fields, **agent_fields)
         elif "source_id" in data and data["source_id"]:
-            assert not data.get("agent_id"), "Passage cannot have both agent_id and source_id"
+            assert not data.get("archive_id"), "Passage cannot have both archive_id and source_id"
             source_fields = {
                 "source_id": data["source_id"],
                 "file_id": data.get("file_id"),
             }
             passage = SourcePassage(**common_fields, **source_fields)
         else:
-            raise ValueError("Passage must have either agent_id or source_id")
+            raise ValueError("Passage must have either archive_id or source_id")
 
         return passage
 
@@ -334,14 +339,14 @@ class PassageManager:
 
     @enforce_types
     @trace_method
-    async def create_many_agent_passages_async(self, passages: List[PydanticPassage], actor: PydanticUser) -> List[PydanticPassage]:
-        """Create multiple agent passages."""
-        agent_passages = []
+    async def create_many_archival_passages_async(self, passages: List[PydanticPassage], actor: PydanticUser) -> List[PydanticPassage]:
+        """Create multiple archival passages."""
+        archival_passages = []
         for p in passages:
-            if not p.agent_id:
-                raise ValueError("Agent passage must have agent_id")
+            if not p.archive_id:
+                raise ValueError("Archival passage must have archive_id")
             if p.source_id:
-                raise ValueError("Agent passage cannot have source_id")
+                raise ValueError("Archival passage cannot have source_id")
 
             data = p.model_dump(to_orm=True)
             common_fields = {
@@ -354,12 +359,12 @@ class PassageManager:
                 "is_deleted": data.get("is_deleted", False),
                 "created_at": data.get("created_at", datetime.now(timezone.utc)),
             }
-            agent_fields = {"agent_id": data["agent_id"]}
-            agent_passages.append(AgentPassage(**common_fields, **agent_fields))
+            archival_fields = {"archive_id": data["archive_id"]}
+            archival_passages.append(ArchivalPassage(**common_fields, **archival_fields))
 
         async with db_registry.async_session() as session:
-            agent_created = await AgentPassage.batch_create_async(items=agent_passages, db_session=session, actor=actor)
-            return [p.to_pydantic() for p in agent_created]
+            archival_created = await ArchivalPassage.batch_create_async(items=archival_passages, db_session=session, actor=actor)
+            return [p.to_pydantic() for p in archival_created]
 
     @enforce_types
     @trace_method
@@ -379,8 +384,8 @@ class PassageManager:
         for p in passages:
             if not p.source_id:
                 raise ValueError("Source passage must have source_id")
-            if p.agent_id:
-                raise ValueError("Source passage cannot have agent_id")
+            if p.archive_id:
+                raise ValueError("Source passage cannot have archive_id")
 
             data = p.model_dump(to_orm=True)
             common_fields = {
@@ -436,7 +441,7 @@ class PassageManager:
 
             for p in passages:
                 model = self._preprocess_passage_for_creation(p)
-                if isinstance(model, AgentPassage):
+                if isinstance(model, ArchivalPassage):
                     agent_passages.append(model)
                 elif isinstance(model, SourcePassage):
                     source_passages.append(model)
@@ -445,7 +450,7 @@ class PassageManager:
 
             results = []
             if agent_passages:
-                agent_created = await AgentPassage.batch_create_async(items=agent_passages, db_session=session, actor=actor)
+                agent_created = await ArchivalPassage.batch_create_async(items=agent_passages, db_session=session, actor=actor)
                 results.extend(agent_created)
             if source_passages:
                 source_created = await SourcePassage.batch_create_async(items=source_passages, db_session=session, actor=actor)
@@ -458,7 +463,6 @@ class PassageManager:
     def insert_passage(
         self,
         agent_state: AgentState,
-        agent_id: str,
         text: str,
         actor: PydanticUser,
     ) -> List[PydanticPassage]:
@@ -494,10 +498,15 @@ class PassageManager:
                         raise TypeError(
                             f"Got back an unexpected payload from text embedding function, type={type(embedding)}, value={embedding}"
                         )
+                # Get or create the default archive for the agent
+                archive = self.archive_manager.get_or_create_default_archive_for_agent(
+                    agent_id=agent_state.id, agent_name=agent_state.name, actor=actor
+                )
+
                 passage = self.create_agent_passage(
                     PydanticPassage(
                         organization_id=actor.organization_id,
-                        agent_id=agent_id,
+                        archive_id=archive.id,
                         text=text,
                         embedding=embedding,
                         embedding_config=agent_state.embedding_config,
@@ -516,12 +525,18 @@ class PassageManager:
     async def insert_passage_async(
         self,
         agent_state: AgentState,
-        agent_id: str,
         text: str,
         actor: PydanticUser,
         image_ids: Optional[List[str]] = None,
     ) -> List[PydanticPassage]:
         """Insert passage(s) into archival memory"""
+        # Get or create default archive for the agent
+        archive = await self.archive_manager.get_or_create_default_archive_for_agent_async(
+            agent_id=agent_state.id,
+            agent_name=agent_state.name,
+            actor=actor,
+        )
+        archive_id = archive.id
 
         embedding_chunk_size = agent_state.embedding_config.embedding_chunk_size
         text_chunks = list(parse_and_chunk_text(text, embedding_chunk_size))
@@ -535,7 +550,7 @@ class PassageManager:
             passages = [
                 PydanticPassage(
                     organization_id=actor.organization_id,
-                    agent_id=agent_id,
+                    archive_id=archive_id,
                     text=chunk_text,
                     embedding=embedding,
                     embedding_config=agent_state.embedding_config,
@@ -543,7 +558,7 @@ class PassageManager:
                 for chunk_text, embedding in zip(text_chunks, embeddings)
             ]
 
-            passages = await self.create_many_agent_passages_async(passages=passages, actor=actor)
+            passages = await self.create_many_archival_passages_async(passages=passages, actor=actor)
 
             return passages
 
@@ -595,7 +610,7 @@ class PassageManager:
 
         with db_registry.session() as session:
             try:
-                curr_passage = AgentPassage.read(
+                curr_passage = ArchivalPassage.read(
                     db_session=session,
                     identifier=passage_id,
                     actor=actor,
@@ -623,7 +638,7 @@ class PassageManager:
 
         async with db_registry.async_session() as session:
             try:
-                curr_passage = await AgentPassage.read_async(
+                curr_passage = await ArchivalPassage.read_async(
                     db_session=session,
                     identifier=passage_id,
                     actor=actor,
@@ -705,7 +720,7 @@ class PassageManager:
 
         with db_registry.session() as session:
             try:
-                passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                passage = ArchivalPassage.read(db_session=session, identifier=passage_id, actor=actor)
                 passage.hard_delete(session, actor=actor)
                 return True
             except NoResultFound:
@@ -720,7 +735,7 @@ class PassageManager:
 
         async with db_registry.async_session() as session:
             try:
-                passage = await AgentPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
+                passage = await ArchivalPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
                 await passage.hard_delete_async(session, actor=actor)
                 return True
             except NoResultFound:
@@ -783,7 +798,7 @@ class PassageManager:
             except NoResultFound:
                 # Try agent passages
                 try:
-                    curr_passage = AgentPassage.read(
+                    curr_passage = ArchivalPassage.read(
                         db_session=session,
                         identifier=passage_id,
                         actor=actor,
@@ -824,7 +839,7 @@ class PassageManager:
             except NoResultFound:
                 # Try archival passages
                 try:
-                    passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                    passage = ArchivalPassage.read(db_session=session, identifier=passage_id, actor=actor)
                     passage.hard_delete(session, actor=actor)
                     return True
                 except NoResultFound:
@@ -854,7 +869,7 @@ class PassageManager:
             except NoResultFound:
                 # Try archival passages
                 try:
-                    passage = await AgentPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
+                    passage = await ArchivalPassage.read_async(db_session=session, identifier=passage_id, actor=actor)
                     await passage.hard_delete_async(session, actor=actor)
                     return True
                 except NoResultFound:
@@ -883,7 +898,7 @@ class PassageManager:
     ) -> bool:
         """Delete multiple agent passages."""
         async with db_registry.async_session() as session:
-            await AgentPassage.bulk_hard_delete_async(db_session=session, identifiers=[p.id for p in passages], actor=actor)
+            await ArchivalPassage.bulk_hard_delete_async(db_session=session, identifiers=[p.id for p in passages], actor=actor)
             return True
 
     @enforce_types
@@ -947,7 +962,21 @@ class PassageManager:
             agent_id: The agent ID of the messages
         """
         with db_registry.session() as session:
-            return AgentPassage.size(db_session=session, actor=actor, agent_id=agent_id)
+            if agent_id:
+                # Count passages through the archives relationship
+                return (
+                    session.query(ArchivalPassage)
+                    .join(ArchivesAgents, ArchivalPassage.archive_id == ArchivesAgents.archive_id)
+                    .filter(
+                        ArchivesAgents.agent_id == agent_id,
+                        ArchivalPassage.organization_id == actor.organization_id,
+                        ArchivalPassage.is_deleted == False,
+                    )
+                    .count()
+                )
+            else:
+                # Count all archival passages in the organization
+                return ArchivalPassage.size(db_session=session, actor=actor)
 
     # DEPRECATED - Use agent_passage_size() instead since this only counted agent passages anyway
     @enforce_types
@@ -961,8 +990,7 @@ class PassageManager:
         import warnings
 
         warnings.warn("size is deprecated. Use agent_passage_size() instead.", DeprecationWarning, stacklevel=2)
-        with db_registry.session() as session:
-            return AgentPassage.size(db_session=session, actor=actor, agent_id=agent_id)
+        return self.agent_passage_size(actor=actor, agent_id=agent_id)
 
     @enforce_types
     @trace_method
@@ -977,7 +1005,23 @@ class PassageManager:
             agent_id: The agent ID of the messages
         """
         async with db_registry.async_session() as session:
-            return await AgentPassage.size_async(db_session=session, actor=actor, agent_id=agent_id)
+            if agent_id:
+                # Count passages through the archives relationship
+                from sqlalchemy import func, select
+
+                result = await session.execute(
+                    select(func.count(ArchivalPassage.id))
+                    .join(ArchivesAgents, ArchivalPassage.archive_id == ArchivesAgents.archive_id)
+                    .where(
+                        ArchivesAgents.agent_id == agent_id,
+                        ArchivalPassage.organization_id == actor.organization_id,
+                        ArchivalPassage.is_deleted == False,
+                    )
+                )
+                return result.scalar() or 0
+            else:
+                # Count all archival passages in the organization
+                return await ArchivalPassage.size_async(db_session=session, actor=actor)
 
     @enforce_types
     @trace_method

letta/services/tool_executor/core_tool_executor.py

@@ -176,7 +176,6 @@ class LettaCoreToolExecutor(ToolExecutor):
         """
         await PassageManager().insert_passage_async(
             agent_state=agent_state,
-            agent_id=agent_state.id,
             text=content,
             actor=actor,
         )

letta/services/tool_executor/sandbox_tool_executor.py

@@ -42,7 +42,7 @@ class SandboxToolExecutor(ToolExecutor):
 
         # Store original memory state
         if agent_state:
-            orig_memory_str = await agent_state.memory.compile_async()
+            orig_memory_str = await agent_state.memory.compile_in_thread_async()
         else:
             orig_memory_str = None
 
@@ -73,7 +73,7 @@ class SandboxToolExecutor(ToolExecutor):
 
             # Verify memory integrity
             if agent_state:
-                new_memory_str = await agent_state.memory.compile_async()
+                new_memory_str = await agent_state.memory.compile_in_thread_async()
                 assert orig_memory_str == new_memory_str, "Memory should not be modified in a sandbox tool"
 
             # Update agent memory if needed

letta/services/tool_executor/tool_execution_manager.py

@@ -4,11 +4,11 @@ from typing import Any, Dict, Optional, Type
 from letta.constants import FUNCTION_RETURN_VALUE_TRUNCATED
 from letta.helpers.datetime_helpers import AsyncTimer
 from letta.log import get_logger
-from letta.orm.enums import ToolType
 from letta.otel.context import get_ctx_attributes
 from letta.otel.metric_registry import MetricRegistry
 from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState
+from letta.schemas.enums import ToolType
 from letta.schemas.sandbox_config import SandboxConfig
 from letta.schemas.tool import Tool
 from letta.schemas.tool_execution_result import ToolExecutionResult

letta/services/tool_manager.py

@@ -22,12 +22,12 @@ from letta.constants import (
 from letta.errors import LettaToolNameConflictError
 from letta.functions.functions import derive_openai_json_schema, load_function_set
 from letta.log import get_logger
-from letta.orm.enums import ToolType
 
 # TODO: Remove this once we translate all of these to the ORM
 from letta.orm.errors import NoResultFound
 from letta.orm.tool import Tool as ToolModel
 from letta.otel.tracing import trace_method
+from letta.schemas.enums import ToolType
 from letta.schemas.tool import Tool as PydanticTool
 from letta.schemas.tool import ToolCreate, ToolUpdate
 from letta.schemas.user import User as PydanticUser
@@ -46,7 +46,7 @@ class ToolManager:
     # TODO: Refactor this across the codebase to use CreateTool instead of passing in a Tool object
     @enforce_types
     @trace_method
-    def create_or_update_tool(self, pydantic_tool: PydanticTool, actor: PydanticUser) -> PydanticTool:
+    def create_or_update_tool(self, pydantic_tool: PydanticTool, actor: PydanticUser, bypass_name_check: bool = False) -> PydanticTool:
         """Create a new tool based on the ToolCreate schema."""
         tool_id = self.get_tool_id_by_name(tool_name=pydantic_tool.name, actor=actor)
         if tool_id:
@@ -60,7 +60,9 @@ class ToolManager:
                 updated_tool_type = None
                 if "tool_type" in update_data:
                     updated_tool_type = update_data.get("tool_type")
-                tool = self.update_tool_by_id(tool_id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type)
+                tool = self.update_tool_by_id(
+                    tool_id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type, bypass_name_check=bypass_name_check
+                )
             else:
                 printd(
                     f"`create_or_update_tool` was called with user_id={actor.id}, organization_id={actor.organization_id}, name={pydantic_tool.name}, but found existing tool with nothing to update."
@@ -73,7 +75,9 @@ class ToolManager:
 
     @enforce_types
     @trace_method
-    async def create_or_update_tool_async(self, pydantic_tool: PydanticTool, actor: PydanticUser) -> PydanticTool:
+    async def create_or_update_tool_async(
+        self, pydantic_tool: PydanticTool, actor: PydanticUser, bypass_name_check: bool = False
+    ) -> PydanticTool:
         """Create a new tool based on the ToolCreate schema."""
         tool_id = await self.get_tool_id_by_name_async(tool_name=pydantic_tool.name, actor=actor)
         if tool_id:
@@ -88,7 +92,9 @@ class ToolManager:
                 updated_tool_type = None
                 if "tool_type" in update_data:
                     updated_tool_type = update_data.get("tool_type")
-                tool = await self.update_tool_by_id_async(tool_id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type)
+                tool = await self.update_tool_by_id_async(
+                    tool_id, ToolUpdate(**update_data), actor, updated_tool_type=updated_tool_type, bypass_name_check=bypass_name_check
+                )
             else:
                 printd(
                     f"`create_or_update_tool` was called with user_id={actor.id}, organization_id={actor.organization_id}, name={pydantic_tool.name}, but found existing tool with nothing to update."
@@ -358,9 +364,13 @@ class ToolManager:
                     results.append(pydantic_tool)
                 except (ValueError, ModuleNotFoundError, AttributeError) as e:
                     tools_to_delete.append(tool)
-                    logger.warning(f"Deleting malformed tool with id={tool.id} and name={tool.name}, error was:\n{e}")
-                    logger.warning("Deleted tool: ")
-                    logger.warning(tool.pretty_print_columns())
+                    logger.warning(
+                        "Deleting malformed tool with id=%s and name=%s. Error was:\n%s\nDeleted tool:%s",
+                        tool.id,
+                        tool.name,
+                        e,
+                        tool.pretty_print_columns(),
+                    )
 
         for tool in tools_to_delete:
             await self.delete_tool_by_id_async(tool.id, actor=actor)
@@ -387,7 +397,12 @@ class ToolManager:
     @enforce_types
     @trace_method
     def update_tool_by_id(
-        self, tool_id: str, tool_update: ToolUpdate, actor: PydanticUser, updated_tool_type: Optional[ToolType] = None
+        self,
+        tool_id: str,
+        tool_update: ToolUpdate,
+        actor: PydanticUser,
+        updated_tool_type: Optional[ToolType] = None,
+        bypass_name_check: bool = False,
     ) -> PydanticTool:
         """Update a tool by its ID with the given ToolUpdate object."""
         # First, check if source code update would cause a name conflict
@@ -395,17 +410,29 @@ class ToolManager:
         new_name = None
         new_schema = None
 
-        if "source_code" in update_data.keys() and "json_schema" not in update_data.keys():
-            # Derive the new schema and name from the source code
-            new_schema = derive_openai_json_schema(source_code=update_data["source_code"])
-            new_name = new_schema["name"]
+        # TODO: Consider this behavior...is this what we want?
+        # TODO: I feel like it's bad if json_schema strays from source code so
+        # if source code is provided, always derive the name from it
+        if "source_code" in update_data.keys() and not bypass_name_check:
+            # derive the schema from source code to get the function name
+            derived_schema = derive_openai_json_schema(source_code=update_data["source_code"])
+            new_name = derived_schema["name"]
+
+            # if json_schema wasn't provided, use the derived schema
+            if "json_schema" not in update_data.keys():
+                new_schema = derived_schema
+            else:
+                # if json_schema was provided, update only its name to match the source code
+                new_schema = update_data["json_schema"].copy()
+                new_schema["name"] = new_name
+                # update the json_schema in update_data so it gets applied in the loop
+                update_data["json_schema"] = new_schema
 
-            # Get current tool to check if name is changing
+            # get current tool to check if name is changing
             current_tool = self.get_tool_by_id(tool_id=tool_id, actor=actor)
-
-            # Check if the name is changing and if so, verify it doesn't conflict
+            # check if the name is changing and if so, verify it doesn't conflict
             if new_name != current_tool.name:
-                # Check if a tool with the new name already exists
+                # check if a tool with the new name already exists
                 existing_tool = self.get_tool_by_name(tool_name=new_name, actor=actor)
                 if existing_tool:
                     raise LettaToolNameConflictError(tool_name=new_name)
@@ -433,7 +460,12 @@ class ToolManager:
     @enforce_types
     @trace_method
     async def update_tool_by_id_async(
-        self, tool_id: str, tool_update: ToolUpdate, actor: PydanticUser, updated_tool_type: Optional[ToolType] = None
+        self,
+        tool_id: str,
+        tool_update: ToolUpdate,
+        actor: PydanticUser,
+        updated_tool_type: Optional[ToolType] = None,
+        bypass_name_check: bool = False,
     ) -> PydanticTool:
         """Update a tool by its ID with the given ToolUpdate object."""
         # First, check if source code update would cause a name conflict
@@ -441,17 +473,29 @@ class ToolManager:
         new_name = None
         new_schema = None
 
-        if "source_code" in update_data.keys() and "json_schema" not in update_data.keys():
-            # Derive the new schema and name from the source code
-            new_schema = derive_openai_json_schema(source_code=update_data["source_code"])
-            new_name = new_schema["name"]
+        # TODO: Consider this behavior...is this what we want?
+        # TODO: I feel like it's bad if json_schema strays from source code so
+        # if source code is provided, always derive the name from it
+        if "source_code" in update_data.keys() and not bypass_name_check:
+            # derive the schema from source code to get the function name
+            derived_schema = derive_openai_json_schema(source_code=update_data["source_code"])
+            new_name = derived_schema["name"]
+
+            # if json_schema wasn't provided, use the derived schema
+            if "json_schema" not in update_data.keys():
+                new_schema = derived_schema
+            else:
+                # if json_schema was provided, update only its name to match the source code
+                new_schema = update_data["json_schema"].copy()
+                new_schema["name"] = new_name
+                # update the json_schema in update_data so it gets applied in the loop
+                update_data["json_schema"] = new_schema
 
-            # Get current tool to check if name is changing
+            # get current tool to check if name is changing
             current_tool = await self.get_tool_by_id_async(tool_id=tool_id, actor=actor)
-
-            # Check if the name is changing and if so, verify it doesn't conflict
+            # check if the name is changing and if so, verify it doesn't conflict
             if new_name != current_tool.name:
-                # Check if a tool with the new name already exists
+                # check if a tool with the new name already exists
                 name_exists = await self.tool_name_exists_async(tool_name=new_name, actor=actor)
                 if name_exists:
                     raise LettaToolNameConflictError(tool_name=new_name)

letta/services/tool_sandbox/base.py

@@ -80,7 +80,7 @@ class AsyncToolSandboxBase(ABC):
         Generate code to run inside of execution sandbox. Serialize the agent state and arguments, call the tool,
         then base64-encode/pickle the result. Runs a jinja2 template constructing the python file.
         """
-        from letta.templates.template_helper import render_template_async
+        from letta.templates.template_helper import render_template_in_thread
 
         # Select the appropriate template based on whether the function is async
         TEMPLATE_NAME = "sandbox_code_file_async.py.j2" if self.is_async_function else "sandbox_code_file.py.j2"
@@ -107,7 +107,7 @@ class AsyncToolSandboxBase(ABC):
 
         agent_state_pickle = pickle.dumps(agent_state) if self.inject_agent_state else None
 
-        return await render_template_async(
+        return await render_template_in_thread(
             TEMPLATE_NAME,
             future_import=future_import,
             inject_agent_state=self.inject_agent_state,