letta-nightly 0.7.21.dev20250522104246__py3-none-any.whl → 0.7.22.dev20250523104244__py3-none-any.whl

letta/services/llm_batch_manager.py

@@ -17,6 +17,7 @@ from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 logger = get_logger(__name__)
@@ -26,6 +27,7 @@ class LLMBatchManager:
     """Manager for handling both LLMBatchJob and LLMBatchItem operations."""
 
     @enforce_types
+    @trace_method
     async def create_llm_batch_job_async(
         self,
         llm_provider: ProviderType,
@@ -47,6 +49,7 @@ class LLMBatchManager:
             return batch.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def get_llm_batch_job_by_id_async(self, llm_batch_id: str, actor: Optional[PydanticUser] = None) -> PydanticLLMBatchJob:
         """Retrieve a single batch job by ID."""
         async with db_registry.async_session() as session:
@@ -54,7 +57,8 @@ class LLMBatchManager:
             return batch.to_pydantic()
 
     @enforce_types
-    def update_llm_batch_status(
+    @trace_method
+    async def update_llm_batch_status_async(
         self,
         llm_batch_id: str,
         status: JobStatus,
@@ -62,15 +66,15 @@ class LLMBatchManager:
         latest_polling_response: Optional[BetaMessageBatch] = None,
     ) -> PydanticLLMBatchJob:
         """Update a batch job’s status and optionally its polling response."""
-        with db_registry.session() as session:
-            batch = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
+        async with db_registry.async_session() as session:
+            batch = await LLMBatchJob.read_async(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)
+            batch = await batch.update_async(db_session=session, actor=actor)
             return batch.to_pydantic()
 
-    def bulk_update_llm_batch_statuses(
+    async def bulk_update_llm_batch_statuses_async(
         self,
         updates: List[BatchPollingResult],
     ) -> None:
@@ -81,7 +85,7 @@ class LLMBatchManager:
         """
         now = datetime.datetime.now(datetime.timezone.utc)
 
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             mappings = []
             for llm_batch_id, status, response in updates:
                 mappings.append(
@@ -93,17 +97,18 @@ class LLMBatchManager:
                     }
                 )
 
-            session.bulk_update_mappings(LLMBatchJob, mappings)
-            session.commit()
+            await session.run_sync(lambda ses: ses.bulk_update_mappings(LLMBatchJob, mappings))
+            await session.commit()
 
     @enforce_types
-    def list_llm_batch_jobs(
+    @trace_method
+    async def list_llm_batch_jobs_async(
         self,
         letta_batch_id: str,
         limit: Optional[int] = None,
         actor: Optional[PydanticUser] = None,
         after: Optional[str] = None,
-    ) -> List[PydanticLLMBatchItem]:
+    ) -> List[PydanticLLMBatchJob]:
         """
         List all batch items for a given llm_batch_id, optionally filtered by additional criteria and limited in count.
 
@@ -115,33 +120,35 @@ class LLMBatchManager:
 
         The results are ordered by their id in ascending order.
         """
-        with db_registry.session() as session:
-            query = session.query(LLMBatchJob).filter(LLMBatchJob.letta_batch_job_id == letta_batch_id)
+        async with db_registry.async_session() as session:
+            query = select(LLMBatchJob).where(LLMBatchJob.letta_batch_job_id == letta_batch_id)
 
             if actor is not None:
-                query = query.filter(LLMBatchJob.organization_id == actor.organization_id)
+                query = query.where(LLMBatchJob.organization_id == actor.organization_id)
 
             # Additional optional filters
             if after is not None:
-                query = query.filter(LLMBatchJob.id > after)
+                query = query.where(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]
+            results = await session.execute(query)
+            return [item.to_pydantic() for item in results.scalars().all()]
 
     @enforce_types
-    def delete_llm_batch_request(self, llm_batch_id: str, actor: PydanticUser) -> None:
+    @trace_method
+    async def delete_llm_batch_request_async(self, llm_batch_id: str, actor: PydanticUser) -> None:
         """Hard delete a batch job by ID."""
-        with db_registry.session() as session:
-            batch = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
-            batch.hard_delete(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            batch = await LLMBatchJob.read_async(db_session=session, identifier=llm_batch_id, actor=actor)
+            await batch.hard_delete_async(db_session=session, actor=actor)
 
     @enforce_types
-    def get_messages_for_letta_batch(
+    @trace_method
+    async def get_messages_for_letta_batch_async(
         self,
         letta_batch_job_id: str,
         limit: int = 100,
@@ -154,12 +161,12 @@ class LLMBatchManager:
         Retrieve messages across all LLM batch jobs associated with a Letta batch job.
         Optimized for PostgreSQL performance using ID-based keyset pagination.
         """
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             # If cursor is provided, get sequence_id for that message
             cursor_sequence_id = None
             if cursor:
-                cursor_query = session.query(MessageModel.sequence_id).filter(MessageModel.id == cursor).limit(1)
-                cursor_result = cursor_query.first()
+                cursor_query = select(MessageModel.sequence_id).where(MessageModel.id == cursor).limit(1)
+                cursor_result = await session.execute(cursor_query)
                 if cursor_result:
                     cursor_sequence_id = cursor_result[0]
                 else:
@@ -167,24 +174,24 @@ class LLMBatchManager:
                     pass
 
             query = (
-                session.query(MessageModel)
+                select(MessageModel)
                 .join(LLMBatchItem, MessageModel.batch_item_id == LLMBatchItem.id)
                 .join(LLMBatchJob, LLMBatchItem.llm_batch_id == LLMBatchJob.id)
-                .filter(LLMBatchJob.letta_batch_job_id == letta_batch_job_id)
+                .where(LLMBatchJob.letta_batch_job_id == letta_batch_job_id)
             )
 
             if actor is not None:
-                query = query.filter(MessageModel.organization_id == actor.organization_id)
+                query = query.where(MessageModel.organization_id == actor.organization_id)
 
             if agent_id is not None:
-                query = query.filter(MessageModel.agent_id == agent_id)
+                query = query.where(MessageModel.agent_id == agent_id)
 
             # Apply cursor-based pagination if cursor exists
             if cursor_sequence_id is not None:
                 if sort_descending:
-                    query = query.filter(MessageModel.sequence_id < cursor_sequence_id)
+                    query = query.where(MessageModel.sequence_id < cursor_sequence_id)
                 else:
-                    query = query.filter(MessageModel.sequence_id > cursor_sequence_id)
+                    query = query.where(MessageModel.sequence_id > cursor_sequence_id)
 
             if sort_descending:
                 query = query.order_by(desc(MessageModel.sequence_id))
@@ -193,10 +200,11 @@ class LLMBatchManager:
 
             query = query.limit(limit)
 
-            results = query.all()
-            return [message.to_pydantic() for message in results]
+            results = await session.execute(query)
+            return [message.to_pydantic() for message in results.scalars().all()]
 
     @enforce_types
+    @trace_method
     async def list_running_llm_batches_async(self, actor: Optional[PydanticUser] = None) -> List[PydanticLLMBatchJob]:
         """Return all running LLM batch jobs, optionally filtered by actor's organization."""
         async with db_registry.async_session() as session:
@@ -209,7 +217,8 @@ class LLMBatchManager:
             return [batch.to_pydantic() for batch in results.scalars().all()]
 
     @enforce_types
-    def create_llm_batch_item(
+    @trace_method
+    async def create_llm_batch_item_async(
         self,
         llm_batch_id: str,
         agent_id: str,
@@ -220,7 +229,7 @@ class LLMBatchManager:
         step_state: Optional[AgentStepState] = None,
     ) -> PydanticLLMBatchItem:
         """Create a new batch item."""
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             item = LLMBatchItem(
                 llm_batch_id=llm_batch_id,
                 agent_id=agent_id,
@@ -230,10 +239,11 @@ class LLMBatchManager:
                 step_state=step_state,
                 organization_id=actor.organization_id,
             )
-            item.create(session, actor=actor)
+            await item.create_async(session, actor=actor)
             return item.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def create_llm_batch_items_bulk_async(
         self, llm_batch_items: List[PydanticLLMBatchItem], actor: PydanticUser
     ) -> List[PydanticLLMBatchItem]:
@@ -269,14 +279,16 @@ class LLMBatchManager:
             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:
+    @trace_method
+    async def get_llm_batch_item_by_id_async(self, item_id: str, actor: PydanticUser) -> PydanticLLMBatchItem:
         """Retrieve a single batch item by ID."""
-        with db_registry.session() as session:
-            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
+        async with db_registry.async_session() as session:
+            item = await LLMBatchItem.read_async(db_session=session, identifier=item_id, actor=actor)
             return item.to_pydantic()
 
     @enforce_types
-    def update_llm_batch_item(
+    @trace_method
+    async def update_llm_batch_item_async(
         self,
         item_id: str,
         actor: PydanticUser,
@@ -286,8 +298,8 @@ class LLMBatchManager:
         step_state: Optional[AgentStepState] = None,
     ) -> PydanticLLMBatchItem:
         """Update fields on a batch item."""
-        with db_registry.session() as session:
-            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
+        async with db_registry.async_session() as session:
+            item = await LLMBatchItem.read_async(db_session=session, identifier=item_id, actor=actor)
 
             if request_status:
                 item.request_status = request_status
@@ -298,9 +310,11 @@ class LLMBatchManager:
             if step_state:
                 item.step_state = step_state
 
-            return item.update(db_session=session, actor=actor).to_pydantic()
+            result = await item.update_async(db_session=session, actor=actor)
+            return result.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def list_llm_batch_items_async(
         self,
         llm_batch_id: str,
@@ -346,7 +360,8 @@ class LLMBatchManager:
             results = await session.execute(query)
             return [item.to_pydantic() for item in results.scalars()]
 
-    def bulk_update_llm_batch_items(
+    @trace_method
+    async def bulk_update_llm_batch_items_async(
         self, llm_batch_id_agent_id_pairs: List[Tuple[str, str]], field_updates: List[Dict[str, Any]], strict: bool = True
     ) -> None:
         """
@@ -364,13 +379,13 @@ class LLMBatchManager:
         if len(llm_batch_id_agent_id_pairs) != len(field_updates):
             raise ValueError("llm_batch_id_agent_id_pairs and field_updates must have the same length")
 
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             # Lookup primary keys for all requested (batch_id, agent_id) pairs
-            items = (
-                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()
+            query = select(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)
             )
+            result = await session.execute(query)
+            items = result.all()
             pair_to_pk = {(batch_id, agent_id): pk for pk, batch_id, agent_id in items}
 
             if strict:
@@ -395,11 +410,12 @@ class LLMBatchManager:
                 mappings.append(update_fields)
 
             if mappings:
-                session.bulk_update_mappings(LLMBatchItem, mappings)
-                session.commit()
+                await session.run_sync(lambda ses: ses.bulk_update_mappings(LLMBatchItem, mappings))
+                await session.commit()
 
     @enforce_types
-    def bulk_update_batch_llm_items_results_by_agent(self, updates: List[ItemUpdateInfo], strict: bool = True) -> None:
+    @trace_method
+    async def bulk_update_batch_llm_items_results_by_agent_async(self, updates: List[ItemUpdateInfo], strict: bool = True) -> 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 = [
@@ -410,33 +426,41 @@ class LLMBatchManager:
             for update in updates
         ]
 
-        self.bulk_update_llm_batch_items(batch_id_agent_id_pairs, field_updates, strict=strict)
+        await self.bulk_update_llm_batch_items_async(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
-    def bulk_update_llm_batch_items_step_status_by_agent(self, updates: List[StepStatusUpdateInfo], strict: bool = True) -> None:
+    @trace_method
+    async def bulk_update_llm_batch_items_step_status_by_agent_async(
+        self, updates: List[StepStatusUpdateInfo], strict: bool = True
+    ) -> 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, strict=strict)
+        await self.bulk_update_llm_batch_items_async(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
-    def bulk_update_llm_batch_items_request_status_by_agent(self, updates: List[RequestStatusUpdateInfo], strict: bool = True) -> None:
+    @trace_method
+    async def bulk_update_llm_batch_items_request_status_by_agent_async(
+        self, updates: List[RequestStatusUpdateInfo], strict: bool = True
+    ) -> 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, strict=strict)
+        await self.bulk_update_llm_batch_items_async(batch_id_agent_id_pairs, field_updates, strict=strict)
 
     @enforce_types
-    def delete_llm_batch_item(self, item_id: str, actor: PydanticUser) -> None:
+    @trace_method
+    async def delete_llm_batch_item_async(self, item_id: str, actor: PydanticUser) -> None:
         """Hard delete a batch item by ID."""
-        with db_registry.session() as session:
-            item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
-            item.hard_delete(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            item = await LLMBatchItem.read_async(db_session=session, identifier=item_id, actor=actor)
+            await item.hard_delete_async(db_session=session, actor=actor)
 
     @enforce_types
-    def count_llm_batch_items(self, llm_batch_id: str) -> int:
+    @trace_method
+    async def count_llm_batch_items_async(self, llm_batch_id: str) -> int:
         """
         Efficiently count the number of batch items for a given llm_batch_id.
 
@@ -446,6 +470,6 @@ class LLMBatchManager:
         Returns:
             int: The total number of batch items associated with the given llm_batch_id.
         """
-        with db_registry.session() as session:
-            count = session.query(func.count(LLMBatchItem.id)).filter(LLMBatchItem.llm_batch_id == llm_batch_id).scalar()
-            return count or 0
+        async with db_registry.async_session() as session:
+            count = await session.execute(select(func.count(LLMBatchItem.id)).where(LLMBatchItem.llm_batch_id == llm_batch_id))
+            return count.scalar() or 0

letta/services/message_manager.py

@@ -13,6 +13,7 @@ from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.message import MessageUpdate
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 logger = get_logger(__name__)
@@ -22,6 +23,7 @@ class MessageManager:
     """Manager class to handle business logic related to Messages."""
 
     @enforce_types
+    @trace_method
     def get_message_by_id(self, message_id: str, actor: PydanticUser) -> Optional[PydanticMessage]:
         """Fetch a message by ID."""
         with db_registry.session() as session:
@@ -32,6 +34,7 @@ class MessageManager:
                 return None
 
     @enforce_types
+    @trace_method
     async def get_message_by_id_async(self, message_id: str, actor: PydanticUser) -> Optional[PydanticMessage]:
         """Fetch a message by ID."""
         async with db_registry.async_session() as session:
@@ -42,6 +45,7 @@ class MessageManager:
                 return None
 
     @enforce_types
+    @trace_method
     def get_messages_by_ids(self, message_ids: List[str], actor: PydanticUser) -> List[PydanticMessage]:
         """Fetch messages by ID and return them in the requested order."""
         with db_registry.session() as session:
@@ -49,6 +53,7 @@ class MessageManager:
         return self._get_messages_by_id_postprocess(results, message_ids)
 
     @enforce_types
+    @trace_method
     async def get_messages_by_ids_async(self, message_ids: List[str], actor: PydanticUser) -> List[PydanticMessage]:
         """Fetch messages by ID and return them in the requested order. Async version of above function."""
         async with db_registry.async_session() as session:
@@ -71,6 +76,7 @@ class MessageManager:
         return list(filter(lambda x: x is not None, [result_dict.get(msg_id, None) for msg_id in message_ids]))
 
     @enforce_types
+    @trace_method
     def create_message(self, pydantic_msg: PydanticMessage, actor: PydanticUser) -> PydanticMessage:
         """Create a new message."""
         with db_registry.session() as session:
@@ -92,6 +98,7 @@ class MessageManager:
         return orm_messages
 
     @enforce_types
+    @trace_method
     def create_many_messages(self, pydantic_msgs: List[PydanticMessage], actor: PydanticUser) -> List[PydanticMessage]:
         """
         Create multiple messages in a single database transaction.
@@ -111,6 +118,7 @@ class MessageManager:
             return [msg.to_pydantic() for msg in created_messages]
 
     @enforce_types
+    @trace_method
     async def create_many_messages_async(self, pydantic_msgs: List[PydanticMessage], actor: PydanticUser) -> List[PydanticMessage]:
         """
         Create multiple messages in a single database transaction asynchronously.
@@ -131,6 +139,7 @@ class MessageManager:
             return [msg.to_pydantic() for msg in created_messages]
 
     @enforce_types
+    @trace_method
     def update_message_by_letta_message(
         self, message_id: str, letta_message_update: LettaMessageUpdateUnion, actor: PydanticUser
     ) -> PydanticMessage:
@@ -169,6 +178,7 @@ class MessageManager:
         raise ValueError(f"Message type got modified: {letta_message_update.message_type}")
 
     @enforce_types
+    @trace_method
     def update_message_by_letta_message(
         self, message_id: str, letta_message_update: LettaMessageUpdateUnion, actor: PydanticUser
     ) -> PydanticMessage:
@@ -207,6 +217,7 @@ class MessageManager:
         raise ValueError(f"Message type got modified: {letta_message_update.message_type}")
 
     @enforce_types
+    @trace_method
     def update_message_by_id(self, message_id: str, message_update: MessageUpdate, actor: PydanticUser) -> PydanticMessage:
         """
         Updates an existing record in the database with values from the provided record object.
@@ -224,6 +235,7 @@ class MessageManager:
             return message.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def update_message_by_id_async(self, message_id: str, message_update: MessageUpdate, actor: PydanticUser) -> PydanticMessage:
         """
         Updates an existing record in the database with values from the provided record object.
@@ -267,6 +279,7 @@ class MessageManager:
         return message
 
     @enforce_types
+    @trace_method
     def delete_message_by_id(self, message_id: str, actor: PydanticUser) -> bool:
         """Delete a message."""
         with db_registry.session() as session:
@@ -281,6 +294,7 @@ class MessageManager:
                 raise ValueError(f"Message with id {message_id} not found.")
 
     @enforce_types
+    @trace_method
     def size(
         self,
         actor: PydanticUser,
@@ -297,6 +311,7 @@ class MessageManager:
             return MessageModel.size(db_session=session, actor=actor, role=role, agent_id=agent_id)
 
     @enforce_types
+    @trace_method
     async def size_async(
         self,
         actor: PydanticUser,
@@ -312,6 +327,7 @@ class MessageManager:
             return await MessageModel.size_async(db_session=session, actor=actor, role=role, agent_id=agent_id)
 
     @enforce_types
+    @trace_method
     def list_user_messages_for_agent(
         self,
         agent_id: str,
@@ -334,6 +350,7 @@ class MessageManager:
         )
 
     @enforce_types
+    @trace_method
     def list_messages_for_agent(
         self,
         agent_id: str,
@@ -437,6 +454,7 @@ class MessageManager:
             return [msg.to_pydantic() for msg in results]
 
     @enforce_types
+    @trace_method
     async def list_messages_for_agent_async(
         self,
         agent_id: str,
@@ -538,6 +556,7 @@ class MessageManager:
             return [msg.to_pydantic() for msg in results]
 
     @enforce_types
+    @trace_method
     def delete_all_messages_for_agent(self, agent_id: str, actor: PydanticUser) -> int:
         """
         Efficiently deletes all messages associated with a given agent_id,

letta/services/organization_manager.py

@@ -5,6 +5,7 @@ from letta.orm.organization import Organization as OrganizationModel
 from letta.schemas.organization import Organization as PydanticOrganization
 from letta.schemas.organization import OrganizationUpdate
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 
@@ -15,11 +16,13 @@ class OrganizationManager:
     DEFAULT_ORG_NAME = "default_org"
 
     @enforce_types
+    @trace_method
     def get_default_organization(self) -> PydanticOrganization:
         """Fetch the default organization."""
         return self.get_organization_by_id(self.DEFAULT_ORG_ID)
 
     @enforce_types
+    @trace_method
     def get_organization_by_id(self, org_id: str) -> Optional[PydanticOrganization]:
         """Fetch an organization by ID."""
         with db_registry.session() as session:
@@ -27,6 +30,7 @@ class OrganizationManager:
             return organization.to_pydantic()
 
     @enforce_types
+    @trace_method
     def create_organization(self, pydantic_org: PydanticOrganization) -> PydanticOrganization:
         """Create a new organization."""
         try:
@@ -36,6 +40,7 @@ class OrganizationManager:
             return self._create_organization(pydantic_org=pydantic_org)
 
     @enforce_types
+    @trace_method
     def _create_organization(self, pydantic_org: PydanticOrganization) -> PydanticOrganization:
         with db_registry.session() as session:
             org = OrganizationModel(**pydantic_org.model_dump(to_orm=True))
@@ -43,11 +48,13 @@ class OrganizationManager:
             return org.to_pydantic()
 
     @enforce_types
+    @trace_method
     def create_default_organization(self) -> PydanticOrganization:
         """Create the default organization."""
         return self.create_organization(PydanticOrganization(name=self.DEFAULT_ORG_NAME, id=self.DEFAULT_ORG_ID))
 
     @enforce_types
+    @trace_method
     def update_organization_name_using_id(self, org_id: str, name: Optional[str] = None) -> PydanticOrganization:
         """Update an organization."""
         with db_registry.session() as session:
@@ -58,6 +65,7 @@ class OrganizationManager:
             return org.to_pydantic()
 
     @enforce_types
+    @trace_method
     def update_organization(self, org_id: str, org_update: OrganizationUpdate) -> PydanticOrganization:
         """Update an organization."""
         with db_registry.session() as session:
@@ -70,6 +78,7 @@ class OrganizationManager:
             return org.to_pydantic()
 
     @enforce_types
+    @trace_method
     def delete_organization_by_id(self, org_id: str):
         """Delete an organization by marking it as deleted."""
         with db_registry.session() as session:
@@ -77,6 +86,7 @@ class OrganizationManager:
             organization.hard_delete(session)
 
     @enforce_types
+    @trace_method
     def list_organizations(self, after: Optional[str] = None, limit: Optional[int] = 50) -> List[PydanticOrganization]:
         """List all organizations with optional pagination."""
         with db_registry.session() as session:

letta/services/passage_manager.py

@@ -11,6 +11,7 @@ from letta.schemas.agent import AgentState
 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.tracing import trace_method
 from letta.utils import enforce_types
 
 
@@ -18,6 +19,7 @@ class PassageManager:
     """Manager class to handle business logic related to Passages."""
 
     @enforce_types
+    @trace_method
     def get_passage_by_id(self, passage_id: str, actor: PydanticUser) -> Optional[PydanticPassage]:
         """Fetch a passage by ID."""
         with db_registry.session() as session:
@@ -34,6 +36,7 @@ class PassageManager:
                     raise NoResultFound(f"Passage with id {passage_id} not found in database.")
 
     @enforce_types
+    @trace_method
     def create_passage(self, pydantic_passage: PydanticPassage, actor: PydanticUser) -> PydanticPassage:
         """Create a new passage in the appropriate table based on whether it has agent_id or source_id."""
         # Common fields for both passage types
@@ -70,11 +73,13 @@ class PassageManager:
             return passage.to_pydantic()
 
     @enforce_types
+    @trace_method
     def create_many_passages(self, passages: List[PydanticPassage], actor: PydanticUser) -> List[PydanticPassage]:
         """Create multiple passages."""
         return [self.create_passage(p, actor) for p in passages]
 
     @enforce_types
+    @trace_method
     def insert_passage(
         self,
         agent_state: AgentState,
@@ -136,6 +141,7 @@ class PassageManager:
             raise e
 
     @enforce_types
+    @trace_method
     def update_passage_by_id(self, passage_id: str, passage: PydanticPassage, actor: PydanticUser, **kwargs) -> Optional[PydanticPassage]:
         """Update a passage."""
         if not passage_id:
@@ -170,6 +176,7 @@ class PassageManager:
             return curr_passage.to_pydantic()
 
     @enforce_types
+    @trace_method
     def delete_passage_by_id(self, passage_id: str, actor: PydanticUser) -> bool:
         """Delete a passage from either source or archival passages."""
         if not passage_id:
@@ -190,6 +197,8 @@ class PassageManager:
                 except NoResultFound:
                     raise NoResultFound(f"Passage with id {passage_id} not found.")
 
+    @enforce_types
+    @trace_method
     def delete_passages(
         self,
         actor: PydanticUser,
@@ -202,6 +211,7 @@ class PassageManager:
         return True
 
     @enforce_types
+    @trace_method
     def size(
         self,
         actor: PydanticUser,
@@ -217,6 +227,7 @@ class PassageManager:
             return AgentPassage.size(db_session=session, actor=actor, agent_id=agent_id)
 
     @enforce_types
+    @trace_method
     async def size_async(
         self,
         actor: PydanticUser,
@@ -230,6 +241,8 @@ class PassageManager:
         async with db_registry.async_session() as session:
             return await AgentPassage.size_async(db_session=session, actor=actor, agent_id=agent_id)
 
+    @enforce_types
+    @trace_method
     def estimate_embeddings_size(
         self,
         actor: PydanticUser,