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

letta/services/block_manager.py

@@ -14,6 +14,7 @@ from letta.schemas.block import Block as PydanticBlock
 from letta.schemas.block import BlockUpdate
 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 @@ logger = get_logger(__name__)
 class BlockManager:
     """Manager class to handle business logic related to Blocks."""
 
+    @trace_method
     @enforce_types
     def create_or_update_block(self, block: PydanticBlock, actor: PydanticUser) -> PydanticBlock:
         """Create a new block based on the Block schema."""
@@ -36,6 +38,7 @@ class BlockManager:
                 block.create(session, actor=actor)
             return block.to_pydantic()
 
+    @trace_method
     @enforce_types
     def batch_create_blocks(self, blocks: List[PydanticBlock], actor: PydanticUser) -> List[PydanticBlock]:
         """
@@ -59,6 +62,7 @@ class BlockManager:
             # Convert back to Pydantic
             return [m.to_pydantic() for m in created_models]
 
+    @trace_method
     @enforce_types
     def update_block(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser) -> PydanticBlock:
         """Update a block by its ID with the given BlockUpdate object."""
@@ -74,6 +78,7 @@ class BlockManager:
             block.update(db_session=session, actor=actor)
             return block.to_pydantic()
 
+    @trace_method
     @enforce_types
     def delete_block(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
         """Delete a block by its ID."""
@@ -82,6 +87,7 @@ class BlockManager:
             block.hard_delete(db_session=session, actor=actor)
             return block.to_pydantic()
 
+    @trace_method
     @enforce_types
     async def get_blocks_async(
         self,
@@ -144,68 +150,7 @@ class BlockManager:
 
             return [block.to_pydantic() for block in blocks]
 
-    @enforce_types
-    async def get_blocks_async(
-        self,
-        actor: PydanticUser,
-        label: Optional[str] = None,
-        is_template: Optional[bool] = None,
-        template_name: Optional[str] = None,
-        identity_id: Optional[str] = None,
-        identifier_keys: Optional[List[str]] = None,
-        limit: Optional[int] = 50,
-    ) -> List[PydanticBlock]:
-        """Async version of get_blocks method. Retrieve blocks based on various optional filters."""
-        from sqlalchemy import select
-        from sqlalchemy.orm import noload
-
-        from letta.orm.sqlalchemy_base import AccessType
-
-        async with db_registry.async_session() as session:
-            # Start with a basic query
-            query = select(BlockModel)
-
-            # Explicitly avoid loading relationships
-            query = query.options(noload(BlockModel.agents), noload(BlockModel.identities), noload(BlockModel.groups))
-
-            # Apply access control
-            query = BlockModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
-
-            # Add filters
-            query = query.where(BlockModel.organization_id == actor.organization_id)
-            if label:
-                query = query.where(BlockModel.label == label)
-
-            if is_template is not None:
-                query = query.where(BlockModel.is_template == is_template)
-
-            if template_name:
-                query = query.where(BlockModel.template_name == template_name)
-
-            if identifier_keys:
-                query = (
-                    query.join(BlockModel.identities)
-                    .filter(BlockModel.identities.property.mapper.class_.identifier_key.in_(identifier_keys))
-                    .distinct(BlockModel.id)
-                )
-
-            if identity_id:
-                query = (
-                    query.join(BlockModel.identities)
-                    .filter(BlockModel.identities.property.mapper.class_.id == identity_id)
-                    .distinct(BlockModel.id)
-                )
-
-            # Add limit
-            if limit:
-                query = query.limit(limit)
-
-            # Execute the query
-            result = await session.execute(query)
-            blocks = result.scalars().all()
-
-            return [block.to_pydantic() for block in blocks]
-
+    @trace_method
     @enforce_types
     def get_block_by_id(self, block_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticBlock]:
         """Retrieve a block by its name."""
@@ -216,6 +161,7 @@ class BlockManager:
             except NoResultFound:
                 return None
 
+    @trace_method
     @enforce_types
     async def get_all_blocks_by_ids_async(self, block_ids: List[str], actor: Optional[PydanticUser] = None) -> List[PydanticBlock]:
         """Retrieve blocks by their ids without loading unnecessary relationships. Async implementation."""
@@ -263,6 +209,7 @@ class BlockManager:
 
             return pydantic_blocks
 
+    @trace_method
     @enforce_types
     async def get_agents_for_block_async(self, block_id: str, actor: PydanticUser) -> List[PydanticAgentState]:
         """
@@ -273,6 +220,7 @@ class BlockManager:
             agents_orm = block.agents
             return await asyncio.gather(*[agent.to_pydantic_async() for agent in agents_orm])
 
+    @trace_method
     @enforce_types
     def size(
         self,
@@ -286,6 +234,7 @@ class BlockManager:
 
     # Block History Functions
 
+    @trace_method
     @enforce_types
     def checkpoint_block(
         self,
@@ -389,6 +338,7 @@ class BlockManager:
         updated_block = block.update(db_session=session, actor=actor, no_commit=True)
         return updated_block
 
+    @trace_method
     @enforce_types
     def undo_checkpoint_block(self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None) -> PydanticBlock:
         """
@@ -431,6 +381,7 @@ class BlockManager:
             session.commit()
             return block.to_pydantic()
 
+    @trace_method
     @enforce_types
     def redo_checkpoint_block(self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None) -> PydanticBlock:
         """
@@ -469,6 +420,7 @@ class BlockManager:
             session.commit()
             return block.to_pydantic()
 
+    @trace_method
     @enforce_types
     async def bulk_update_block_values_async(
         self, updates: Dict[str, str], actor: PydanticUser, return_hydrated: bool = False

letta/services/group_manager.py

@@ -12,11 +12,13 @@ from letta.schemas.letta_message import LettaMessage
 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
 
 
 class GroupManager:
 
+    @trace_method
     @enforce_types
     def list_groups(
         self,
@@ -42,12 +44,14 @@ class GroupManager:
             )
             return [group.to_pydantic() for group in groups]
 
+    @trace_method
     @enforce_types
     def retrieve_group(self, group_id: str, actor: PydanticUser) -> PydanticGroup:
         with db_registry.session() as session:
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
             return group.to_pydantic()
 
+    @trace_method
     @enforce_types
     def create_group(self, group: GroupCreate, actor: PydanticUser) -> PydanticGroup:
         with db_registry.session() as session:
@@ -93,6 +97,7 @@ class GroupManager:
             new_group.create(session, actor=actor)
             return new_group.to_pydantic()
 
+    @trace_method
     @enforce_types
     def modify_group(self, group_id: str, group_update: GroupUpdate, actor: PydanticUser) -> PydanticGroup:
         with db_registry.session() as session:
@@ -155,6 +160,7 @@ class GroupManager:
             group.update(session, actor=actor)
             return group.to_pydantic()
 
+    @trace_method
     @enforce_types
     def delete_group(self, group_id: str, actor: PydanticUser) -> None:
         with db_registry.session() as session:
@@ -162,6 +168,7 @@ class GroupManager:
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
             group.hard_delete(session)
 
+    @trace_method
     @enforce_types
     def list_group_messages(
         self,
@@ -198,6 +205,7 @@ class GroupManager:
 
             return messages
 
+    @trace_method
     @enforce_types
     def reset_messages(self, group_id: str, actor: PydanticUser) -> None:
         with db_registry.session() as session:
@@ -211,6 +219,7 @@ class GroupManager:
 
             session.commit()
 
+    @trace_method
     @enforce_types
     def bump_turns_counter(self, group_id: str, actor: PydanticUser) -> int:
         with db_registry.session() as session:
@@ -222,6 +231,18 @@ class GroupManager:
             group.update(session, actor=actor)
             return group.turns_counter
 
+    @trace_method
+    @enforce_types
+    async def bump_turns_counter_async(self, group_id: str, actor: PydanticUser) -> int:
+        async with db_registry.async_session() as session:
+            # Ensure group is loadable by user
+            group = await GroupModel.read_async(session, identifier=group_id, actor=actor)
+
+            # Update turns counter
+            group.turns_counter = (group.turns_counter + 1) % group.sleeptime_agent_frequency
+            await group.update_async(session, actor=actor)
+            return group.turns_counter
+
     @enforce_types
     def get_last_processed_message_id_and_update(self, group_id: str, last_processed_message_id: str, actor: PydanticUser) -> str:
         with db_registry.session() as session:
@@ -235,6 +256,22 @@ class GroupManager:
 
             return prev_last_processed_message_id
 
+    @trace_method
+    @enforce_types
+    async def get_last_processed_message_id_and_update_async(
+        self, group_id: str, last_processed_message_id: str, actor: PydanticUser
+    ) -> str:
+        async with db_registry.async_session() as session:
+            # Ensure group is loadable by user
+            group = await GroupModel.read_async(session, identifier=group_id, actor=actor)
+
+            # Update last processed message id
+            prev_last_processed_message_id = group.last_processed_message_id
+            group.last_processed_message_id = last_processed_message_id
+            await group.update_async(session, actor=actor)
+
+            return prev_last_processed_message_id
+
     @enforce_types
     def size(
         self,

letta/services/identity_manager.py

@@ -12,12 +12,14 @@ from letta.schemas.identity import Identity as PydanticIdentity
 from letta.schemas.identity import IdentityCreate, IdentityProperty, IdentityType, IdentityUpdate, IdentityUpsert
 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
 
 
 class IdentityManager:
 
     @enforce_types
+    @trace_method
     async def list_identities_async(
         self,
         name: Optional[str] = None,
@@ -48,12 +50,14 @@ class IdentityManager:
             return [identity.to_pydantic() for identity in identities]
 
     @enforce_types
+    @trace_method
     async def get_identity_async(self, identity_id: str, actor: PydanticUser) -> PydanticIdentity:
         async with db_registry.async_session() as session:
             identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
             return identity.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def create_identity_async(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
         async with db_registry.async_session() as session:
             new_identity = IdentityModel(**identity.model_dump(exclude={"agent_ids", "block_ids"}, exclude_unset=True))
@@ -78,6 +82,7 @@ class IdentityManager:
             return new_identity.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def upsert_identity_async(self, identity: IdentityUpsert, actor: PydanticUser) -> PydanticIdentity:
         async with db_registry.async_session() as session:
             existing_identity = await IdentityModel.read_async(
@@ -103,6 +108,7 @@ class IdentityManager:
             )
 
     @enforce_types
+    @trace_method
     async def update_identity_async(
         self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False
     ) -> PydanticIdentity:
@@ -165,6 +171,7 @@ class IdentityManager:
         return existing_identity.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def upsert_identity_properties_async(
         self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser
     ) -> PydanticIdentity:
@@ -181,6 +188,7 @@ class IdentityManager:
             )
 
     @enforce_types
+    @trace_method
     async def delete_identity_async(self, identity_id: str, actor: PydanticUser) -> None:
         async with db_registry.async_session() as session:
             identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
@@ -192,6 +200,7 @@ class IdentityManager:
             await session.commit()
 
     @enforce_types
+    @trace_method
     async def size_async(
         self,
         actor: PydanticUser,

letta/services/job_manager.py

@@ -25,6 +25,7 @@ from letta.schemas.step import Step as PydanticStep
 from letta.schemas.usage import LettaUsageStatistics
 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
 
 
@@ -32,6 +33,7 @@ class JobManager:
     """Manager class to handle business logic related to Jobs."""
 
     @enforce_types
+    @trace_method
     def create_job(
         self, pydantic_job: Union[PydanticJob, PydanticRun, PydanticBatchJob], actor: PydanticUser
     ) -> Union[PydanticJob, PydanticRun, PydanticBatchJob]:
@@ -45,6 +47,7 @@ class JobManager:
         return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def create_job_async(
         self, pydantic_job: Union[PydanticJob, PydanticRun, PydanticBatchJob], actor: PydanticUser
     ) -> Union[PydanticJob, PydanticRun, PydanticBatchJob]:
@@ -58,6 +61,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     def update_job_by_id(self, job_id: str, job_update: JobUpdate, actor: PydanticUser) -> PydanticJob:
         """Update a job by its ID with the given JobUpdate object."""
         with db_registry.session() as session:
@@ -82,6 +86,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def update_job_by_id_async(self, job_id: str, job_update: JobUpdate, actor: PydanticUser) -> PydanticJob:
         """Update a job by its ID with the given JobUpdate object asynchronously."""
         async with db_registry.async_session() as session:
@@ -106,6 +111,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     def get_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Fetch a job by its ID."""
         with db_registry.session() as session:
@@ -114,6 +120,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def get_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Fetch a job by its ID asynchronously."""
         async with db_registry.async_session() as session:
@@ -122,6 +129,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     def list_jobs(
         self,
         actor: PydanticUser,
@@ -151,6 +159,7 @@ class JobManager:
             return [job.to_pydantic() for job in jobs]
 
     @enforce_types
+    @trace_method
     async def list_jobs_async(
         self,
         actor: PydanticUser,
@@ -180,6 +189,7 @@ class JobManager:
             return [job.to_pydantic() for job in jobs]
 
     @enforce_types
+    @trace_method
     def delete_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Delete a job by its ID."""
         with db_registry.session() as session:
@@ -188,6 +198,7 @@ class JobManager:
             return job.to_pydantic()
 
     @enforce_types
+    @trace_method
     def get_job_messages(
         self,
         job_id: str,
@@ -238,6 +249,7 @@ class JobManager:
         return [message.to_pydantic() for message in messages]
 
     @enforce_types
+    @trace_method
     def get_job_steps(
         self,
         job_id: str,
@@ -283,6 +295,7 @@ class JobManager:
         return [step.to_pydantic() for step in steps]
 
     @enforce_types
+    @trace_method
     def add_message_to_job(self, job_id: str, message_id: str, actor: PydanticUser) -> None:
         """
         Associate a message with a job by creating a JobMessage record.
@@ -306,6 +319,7 @@ class JobManager:
             session.commit()
 
     @enforce_types
+    @trace_method
     def get_job_usage(self, job_id: str, actor: PydanticUser) -> LettaUsageStatistics:
         """
         Get usage statistics for a job.
@@ -343,6 +357,7 @@ class JobManager:
             )
 
     @enforce_types
+    @trace_method
     def add_job_usage(
         self,
         job_id: str,
@@ -383,6 +398,7 @@ class JobManager:
             session.commit()
 
     @enforce_types
+    @trace_method
     def get_run_messages(
         self,
         run_id: str,
@@ -434,6 +450,7 @@ class JobManager:
         return messages
 
     @enforce_types
+    @trace_method
     def get_step_messages(
         self,
         run_id: str,