letta-nightly 0.7.13.dev20250512104305__py3-none-any.whl → 0.7.14.dev20250513020711__py3-none-any.whl

letta/services/block_manager.py

@@ -1,4 +1,3 @@
-import os
 from typing import Dict, List, Optional
 
 from sqlalchemy.orm import Session
@@ -10,9 +9,10 @@ from letta.orm.enums import ActorType
 from letta.orm.errors import NoResultFound
 from letta.schemas.agent import AgentState as PydanticAgentState
 from letta.schemas.block import Block as PydanticBlock
-from letta.schemas.block import BlockUpdate, Human, Persona
+from letta.schemas.block import BlockUpdate
 from letta.schemas.user import User as PydanticUser
-from letta.utils import enforce_types, list_human_files, list_persona_files
+from letta.server.db import db_registry
+from letta.utils import enforce_types
 
 logger = get_logger(__name__)
 
@@ -20,12 +20,6 @@ logger = get_logger(__name__)
 class BlockManager:
     """Manager class to handle business logic related to Blocks."""
 
-    def __init__(self):
-        # Fetching the db_context similarly as in ToolManager
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def create_or_update_block(self, block: PydanticBlock, actor: PydanticUser) -> PydanticBlock:
         """Create a new block based on the Block schema."""
@@ -34,7 +28,7 @@ class BlockManager:
             update_data = BlockUpdate(**block.model_dump(to_orm=True, exclude_none=True))
             self.update_block(block.id, update_data, actor)
         else:
-            with self.session_maker() as session:
+            with db_registry.session() as session:
                 data = block.model_dump(to_orm=True, exclude_none=True)
                 block = BlockModel(**data, organization_id=actor.organization_id)
                 block.create(session, actor=actor)
@@ -53,7 +47,7 @@ class BlockManager:
         if not blocks:
             return []
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             block_models = [
                 BlockModel(**block.model_dump(to_orm=True, exclude_none=True), organization_id=actor.organization_id) for block in blocks
             ]
@@ -68,7 +62,7 @@ class BlockManager:
         """Update a block by its ID with the given BlockUpdate object."""
         # Safety check for block
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
             update_data = block_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
 
@@ -81,7 +75,7 @@ class BlockManager:
     @enforce_types
     def delete_block(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
         """Delete a block by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             block = BlockModel.read(db_session=session, identifier=block_id)
             block.hard_delete(db_session=session, actor=actor)
             return block.to_pydantic()
@@ -100,7 +94,7 @@ class BlockManager:
         limit: Optional[int] = 50,
     ) -> List[PydanticBlock]:
         """Retrieve blocks based on various optional filters."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Prepare filters
             filters = {"organization_id": actor.organization_id}
             if label:
@@ -126,7 +120,7 @@ class BlockManager:
     @enforce_types
     def get_block_by_id(self, block_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticBlock]:
         """Retrieve a block by its name."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
                 return block.to_pydantic()
@@ -136,32 +130,65 @@ class BlockManager:
     @enforce_types
     def get_all_blocks_by_ids(self, block_ids: List[str], actor: Optional[PydanticUser] = None) -> List[PydanticBlock]:
         """Retrieve blocks by their ids."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             blocks = [block.to_pydantic() for block in BlockModel.read_multiple(db_session=session, identifiers=block_ids, actor=actor)]
             # backwards compatibility. previous implementation added None for every block not found.
             blocks.extend([None for _ in range(len(block_ids) - len(blocks))])
             return blocks
 
     @enforce_types
-    def add_default_blocks(self, actor: PydanticUser):
-        for persona_file in list_persona_files():
-            with open(persona_file, "r", encoding="utf-8") as f:
-                text = f.read()
-            name = os.path.basename(persona_file).replace(".txt", "")
-            self.create_or_update_block(Persona(template_name=name, value=text, is_template=True), actor=actor)
-
-        for human_file in list_human_files():
-            with open(human_file, "r", encoding="utf-8") as f:
-                text = f.read()
-            name = os.path.basename(human_file).replace(".txt", "")
-            self.create_or_update_block(Human(template_name=name, value=text, is_template=True), actor=actor)
+    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."""
+        from sqlalchemy import select
+        from sqlalchemy.orm import noload
+
+        from letta.orm.sqlalchemy_base import AccessType
+
+        if not block_ids:
+            return []
+
+        async with db_registry.async_session() as session:
+            # Start with a basic query
+            query = select(BlockModel)
+
+            # Add ID filter
+            query = query.where(BlockModel.id.in_(block_ids))
+
+            # Explicitly avoid loading relationships
+            query = query.options(noload(BlockModel.agents), noload(BlockModel.identities), noload(BlockModel.groups))
+
+            # Apply access control if actor is provided
+            if actor:
+                query = BlockModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+
+            # Add soft delete filter if applicable
+            if hasattr(BlockModel, "is_deleted"):
+                query = query.where(BlockModel.is_deleted == False)
+
+            # Execute the query
+            result = await session.execute(query)
+            blocks = result.scalars().all()
+
+            # Convert to Pydantic models
+            pydantic_blocks = [block.to_pydantic() for block in blocks]
+
+            # For backward compatibility, add None for missing blocks
+            if len(pydantic_blocks) < len(block_ids):
+                {block.id for block in pydantic_blocks}
+                result_blocks = []
+                for block_id in block_ids:
+                    block = next((b for b in pydantic_blocks if b.id == block_id), None)
+                    result_blocks.append(block)
+                return result_blocks
+
+            return pydantic_blocks
 
     @enforce_types
     def get_agents_for_block(self, block_id: str, actor: PydanticUser) -> List[PydanticAgentState]:
         """
         Retrieve all agents associated with a given block.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
             agents_orm = block.agents
             agents_pydantic = [agent.to_pydantic() for agent in agents_orm]
@@ -176,7 +203,7 @@ class BlockManager:
         """
         Get the total count of blocks for the given user.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return BlockModel.size(db_session=session, actor=actor)
 
     # Block History Functions
@@ -199,7 +226,7 @@ class BlockManager:
           strictly linear history.
         - A single commit at the end ensures atomicity.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # 1) Load the Block
             if use_preloaded_block is not None:
                 block = session.merge(use_preloaded_block)
@@ -291,7 +318,7 @@ class BlockManager:
         If older sequences have been pruned, we jump to the largest sequence
         number that is still < current_seq.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # 1) Load the current block
             block = (
                 session.merge(use_preloaded_block)
@@ -333,7 +360,7 @@ class BlockManager:
         If some middle checkpoints have been pruned, we jump to the smallest
         sequence > current_seq that remains.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             block = (
                 session.merge(use_preloaded_block)
                 if use_preloaded_block
@@ -383,7 +410,7 @@ class BlockManager:
             NoResultFound if any block_id doesn’t exist or isn’t visible to this actor
             ValueError     if any new value exceeds its block’s limit
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             q = session.query(BlockModel).filter(BlockModel.id.in_(updates.keys()), BlockModel.organization_id == actor.organization_id)
             blocks = q.all()
 

letta/services/group_manager.py

@@ -11,16 +11,12 @@ from letta.schemas.group import GroupCreate, GroupUpdate, ManagerType
 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.utils import enforce_types
 
 
 class GroupManager:
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def list_groups(
         self,
@@ -31,7 +27,7 @@ class GroupManager:
         after: Optional[str] = None,
         limit: Optional[int] = 50,
     ) -> list[PydanticGroup]:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             filters = {"organization_id": actor.organization_id}
             if project_id:
                 filters["project_id"] = project_id
@@ -48,13 +44,13 @@ class GroupManager:
 
     @enforce_types
     def retrieve_group(self, group_id: str, actor: PydanticUser) -> PydanticGroup:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
             return group.to_pydantic()
 
     @enforce_types
     def create_group(self, group: GroupCreate, actor: PydanticUser) -> PydanticGroup:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             new_group = GroupModel()
             new_group.organization_id = actor.organization_id
             new_group.description = group.description
@@ -99,7 +95,7 @@ class GroupManager:
 
     @enforce_types
     def modify_group(self, group_id: str, group_update: GroupUpdate, actor: PydanticUser) -> PydanticGroup:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
 
             sleeptime_agent_frequency = None
@@ -161,7 +157,7 @@ class GroupManager:
 
     @enforce_types
     def delete_group(self, group_id: str, actor: PydanticUser) -> None:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Retrieve the agent
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
             group.hard_delete(session)
@@ -178,7 +174,7 @@ class GroupManager:
         assistant_message_tool_name: str = "send_message",
         assistant_message_tool_kwarg: str = "message",
     ) -> list[LettaMessage]:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             filters = {
                 "organization_id": actor.organization_id,
                 "group_id": group_id,
@@ -204,7 +200,7 @@ class GroupManager:
 
     @enforce_types
     def reset_messages(self, group_id: str, actor: PydanticUser) -> None:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Ensure group is loadable by user
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
 
@@ -217,7 +213,7 @@ class GroupManager:
 
     @enforce_types
     def bump_turns_counter(self, group_id: str, actor: PydanticUser) -> int:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Ensure group is loadable by user
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
 
@@ -228,7 +224,7 @@ class GroupManager:
 
     @enforce_types
     def get_last_processed_message_id_and_update(self, group_id: str, last_processed_message_id: str, actor: PydanticUser) -> str:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Ensure group is loadable by user
             group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
 
@@ -247,7 +243,7 @@ class GroupManager:
         """
         Get the total count of groups for the given user.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return GroupModel.size(db_session=session, actor=actor)
 
     def _process_agent_relationship(self, session: Session, group: GroupModel, agent_ids: List[str], allow_partial=False, replace=True):

letta/services/identity_manager.py

@@ -10,16 +10,12 @@ from letta.orm.identity import Identity as IdentityModel
 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.utils import enforce_types
 
 
 class IdentityManager:
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def list_identities(
         self,
@@ -32,7 +28,7 @@ class IdentityManager:
         limit: Optional[int] = 50,
         actor: PydanticUser = None,
     ) -> list[PydanticIdentity]:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             filters = {"organization_id": actor.organization_id}
             if project_id:
                 filters["project_id"] = project_id
@@ -52,13 +48,13 @@ class IdentityManager:
 
     @enforce_types
     def get_identity(self, identity_id: str, actor: PydanticUser) -> PydanticIdentity:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
             return identity.to_pydantic()
 
     @enforce_types
     def create_identity(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             new_identity = IdentityModel(**identity.model_dump(exclude={"agent_ids", "block_ids"}, exclude_unset=True))
             new_identity.organization_id = actor.organization_id
             self._process_relationship(
@@ -82,7 +78,7 @@ class IdentityManager:
 
     @enforce_types
     def upsert_identity(self, identity: IdentityUpsert, actor: PydanticUser) -> PydanticIdentity:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             existing_identity = IdentityModel.read(
                 db_session=session,
                 identifier_key=identity.identifier_key,
@@ -107,7 +103,7 @@ class IdentityManager:
 
     @enforce_types
     def update_identity(self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False) -> PydanticIdentity:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 existing_identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
             except NoResultFound:
@@ -167,7 +163,7 @@ class IdentityManager:
 
     @enforce_types
     def upsert_identity_properties(self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser) -> PydanticIdentity:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             existing_identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
             if existing_identity is None:
                 raise HTTPException(status_code=404, detail="Identity not found")
@@ -181,7 +177,7 @@ class IdentityManager:
 
     @enforce_types
     def delete_identity(self, identity_id: str, actor: PydanticUser) -> None:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             identity = IdentityModel.read(db_session=session, identifier=identity_id)
             if identity is None:
                 raise HTTPException(status_code=404, detail="Identity not found")
@@ -198,7 +194,7 @@ class IdentityManager:
         """
         Get the total count of identities for the given user.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return IdentityModel.size(db_session=session, actor=actor)
 
     def _process_relationship(

letta/services/job_manager.py

@@ -24,24 +24,19 @@ from letta.schemas.run import Run as PydanticRun
 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.utils import enforce_types
 
 
 class JobManager:
     """Manager class to handle business logic related to Jobs."""
 
-    def __init__(self):
-        # Fetching the db_context similarly as in OrganizationManager
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def create_job(
         self, pydantic_job: Union[PydanticJob, PydanticRun, PydanticBatchJob], actor: PydanticUser
     ) -> Union[PydanticJob, PydanticRun, PydanticBatchJob]:
         """Create a new job based on the JobCreate schema."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Associate the job with the user
             pydantic_job.user_id = actor.id
             job_data = pydantic_job.model_dump(to_orm=True)
@@ -52,7 +47,7 @@ class JobManager:
     @enforce_types
     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 self.session_maker() as session:
+        with db_registry.session() as session:
             # Fetch the job by ID
             job = self._verify_job_access(session=session, job_id=job_id, actor=actor, access=["write"])
 
@@ -76,7 +71,7 @@ class JobManager:
     @enforce_types
     def get_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Fetch a job by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Retrieve job by ID using the Job model's read method
             job = JobModel.read(db_session=session, identifier=job_id, actor=actor, access_type=AccessType.USER)
             return job.to_pydantic()
@@ -93,7 +88,7 @@ class JobManager:
         ascending: bool = True,
     ) -> List[PydanticJob]:
         """List all jobs with optional pagination and status filter."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             filter_kwargs = {"user_id": actor.id, "job_type": job_type}
 
             # Add status filter if provided
@@ -113,7 +108,7 @@ class JobManager:
     @enforce_types
     def delete_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Delete a job by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             job = self._verify_job_access(session=session, job_id=job_id, actor=actor)
             job.hard_delete(db_session=session, actor=actor)
             return job.to_pydantic()
@@ -147,7 +142,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Build filters
             filters = {}
             if role is not None:
@@ -195,7 +190,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Build filters
             filters = {}
             filters["job_id"] = job_id
@@ -227,7 +222,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # First verify job exists and user has access
             self._verify_job_access(session, job_id, actor, access=["write"])
 
@@ -251,7 +246,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # First verify job exists and user has access
             self._verify_job_access(session, job_id, actor)
 
@@ -293,7 +288,7 @@ class JobManager:
         Raises:
             NoResultFound: If the job does not exist or user does not have access
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # First verify job exists and user has access
             self._verify_job_access(session, job_id, actor, access=["write"])
 
@@ -453,7 +448,7 @@ class JobManager:
         Returns:
             The request config for the job
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             job = session.query(JobModel).filter(JobModel.id == run_id).first()
             request_config = job.request_config or LettaRequestConfig()
         return request_config

letta/services/llm_batch_manager.py

@@ -16,6 +16,7 @@ from letta.schemas.llm_batch_job import LLMBatchJob as PydanticLLMBatchJob
 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.utils import enforce_types
 
 logger = get_logger(__name__)
@@ -24,11 +25,6 @@ logger = get_logger(__name__)
 class LLMBatchManager:
     """Manager for handling both LLMBatchJob and LLMBatchItem operations."""
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def create_llm_batch_job(
         self,
@@ -39,7 +35,7 @@ class LLMBatchManager:
         status: JobStatus = JobStatus.created,
     ) -> PydanticLLMBatchJob:
         """Create a new LLM batch job."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             batch = LLMBatchJob(
                 status=status,
                 llm_provider=llm_provider,
@@ -53,7 +49,7 @@ class LLMBatchManager:
     @enforce_types
     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:
+        with db_registry.session() as session:
             batch = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
             return batch.to_pydantic()
 
@@ -66,7 +62,7 @@ class LLMBatchManager:
         latest_polling_response: Optional[BetaMessageBatch] = None,
     ) -> PydanticLLMBatchJob:
         """Update a batch job’s status and optionally its polling response."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             batch = LLMBatchJob.read(db_session=session, identifier=llm_batch_id, actor=actor)
             batch.status = status
             batch.latest_polling_response = latest_polling_response
@@ -85,7 +81,7 @@ class LLMBatchManager:
         """
         now = datetime.datetime.now(datetime.timezone.utc)
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             mappings = []
             for llm_batch_id, status, response in updates:
                 mappings.append(
@@ -119,7 +115,7 @@ class LLMBatchManager:
 
         The results are ordered by their id in ascending order.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             query = session.query(LLMBatchJob).filter(LLMBatchJob.letta_batch_job_id == letta_batch_id)
 
             if actor is not None:
@@ -140,7 +136,7 @@ class LLMBatchManager:
     @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:
+        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)
 
@@ -158,7 +154,7 @@ 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 self.session_maker() as session:
+        with db_registry.session() as session:
             # If cursor is provided, get sequence_id for that message
             cursor_sequence_id = None
             if cursor:
@@ -203,7 +199,7 @@ class LLMBatchManager:
     @enforce_types
     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:
+        with db_registry.session() as session:
             query = session.query(LLMBatchJob).filter(LLMBatchJob.status == JobStatus.running)
 
             if actor is not None:
@@ -224,7 +220,7 @@ class LLMBatchManager:
         step_state: Optional[AgentStepState] = None,
     ) -> PydanticLLMBatchItem:
         """Create a new batch item."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             item = LLMBatchItem(
                 llm_batch_id=llm_batch_id,
                 agent_id=agent_id,
@@ -249,7 +245,7 @@ class LLMBatchManager:
         Returns:
             List of created batch items as Pydantic models
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Convert Pydantic models to ORM objects
             orm_items = []
             for item in llm_batch_items:
@@ -274,7 +270,7 @@ class LLMBatchManager:
     @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:
+        with db_registry.session() as session:
             item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
             return item.to_pydantic()
 
@@ -289,7 +285,7 @@ class LLMBatchManager:
         step_state: Optional[AgentStepState] = None,
     ) -> PydanticLLMBatchItem:
         """Update fields on a batch item."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             item = LLMBatchItem.read(db_session=session, identifier=item_id, actor=actor)
 
             if request_status:
@@ -325,7 +321,7 @@ class LLMBatchManager:
 
         The results are ordered by their id in ascending order.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             query = session.query(LLMBatchItem).filter(LLMBatchItem.llm_batch_id == llm_batch_id)
 
             if actor is not None:
@@ -367,7 +363,7 @@ 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 self.session_maker() as session:
+        with db_registry.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)
@@ -434,7 +430,7 @@ class LLMBatchManager:
     @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:
+        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)
 
@@ -449,6 +445,6 @@ class LLMBatchManager:
         Returns:
             int: The total number of batch items associated with the given llm_batch_id.
         """
-        with self.session_maker() as session:
+        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