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

letta/services/message_manager.py

@@ -12,6 +12,7 @@ from letta.schemas.letta_message import LettaMessageUpdateUnion
 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.utils import enforce_types
 
 logger = get_logger(__name__)
@@ -20,15 +21,10 @@ logger = get_logger(__name__)
 class MessageManager:
     """Manager class to handle business logic related to Messages."""
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def get_message_by_id(self, message_id: str, actor: PydanticUser) -> Optional[PydanticMessage]:
         """Fetch a message by ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 message = MessageModel.read(db_session=session, identifier=message_id, actor=actor)
                 return message.to_pydantic()
@@ -38,7 +34,7 @@ class MessageManager:
     @enforce_types
     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 self.session_maker() as session:
+        with db_registry.session() as session:
             results = MessageModel.list(db_session=session, id=message_ids, organization_id=actor.organization_id, limit=len(message_ids))
 
             if len(results) != len(message_ids):
@@ -53,7 +49,7 @@ class MessageManager:
     @enforce_types
     def create_message(self, pydantic_msg: PydanticMessage, actor: PydanticUser) -> PydanticMessage:
         """Create a new message."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Set the organization id of the Pydantic message
             pydantic_msg.organization_id = actor.organization_id
             msg_data = pydantic_msg.model_dump(to_orm=True)
@@ -86,7 +82,7 @@ class MessageManager:
             orm_messages.append(MessageModel(**msg_data))
 
         # Use the batch_create method for efficient creation
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             created_messages = MessageModel.batch_create(orm_messages, session, actor=actor)
 
             # Convert back to Pydantic models
@@ -173,7 +169,7 @@ class MessageManager:
         """
         Updates an existing record in the database with values from the provided record object.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Fetch existing message from database
             message = MessageModel.read(
                 db_session=session,
@@ -181,31 +177,57 @@ class MessageManager:
                 actor=actor,
             )
 
-            # Some safety checks specific to messages
-            if message_update.tool_calls and message.role != MessageRole.assistant:
-                raise ValueError(
-                    f"Tool calls {message_update.tool_calls} can only be added to assistant messages. Message {message_id} has role {message.role}."
-                )
-            if message_update.tool_call_id and message.role != MessageRole.tool:
-                raise ValueError(
-                    f"Tool call IDs {message_update.tool_call_id} can only be added to tool messages. Message {message_id} has role {message.role}."
-                )
-
-            # get update dictionary
-            update_data = message_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
-            # Remove redundant update fields
-            update_data = {key: value for key, value in update_data.items() if getattr(message, key) != value}
-
-            for key, value in update_data.items():
-                setattr(message, key, value)
+            message = self._update_message_by_id_impl(message_id, message_update, actor, message)
             message.update(db_session=session, actor=actor)
+            return message.to_pydantic()
+
+    @enforce_types
+    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.
+        Async version of the function above.
+        """
+        async with db_registry.async_session() as session:
+            # Fetch existing message from database
+            message = await MessageModel.read_async(
+                db_session=session,
+                identifier=message_id,
+                actor=actor,
+            )
 
+            message = self._update_message_by_id_impl(message_id, message_update, actor, message)
+            await message.update_async(db_session=session, actor=actor)
             return message.to_pydantic()
 
+    def _update_message_by_id_impl(
+        self, message_id: str, message_update: MessageUpdate, actor: PydanticUser, message: MessageModel
+    ) -> MessageModel:
+        """
+        Modifies the existing message object to update the database in the sync/async functions.
+        """
+        # Some safety checks specific to messages
+        if message_update.tool_calls and message.role != MessageRole.assistant:
+            raise ValueError(
+                f"Tool calls {message_update.tool_calls} can only be added to assistant messages. Message {message_id} has role {message.role}."
+            )
+        if message_update.tool_call_id and message.role != MessageRole.tool:
+            raise ValueError(
+                f"Tool call IDs {message_update.tool_call_id} can only be added to tool messages. Message {message_id} has role {message.role}."
+            )
+
+        # get update dictionary
+        update_data = message_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
+        # Remove redundant update fields
+        update_data = {key: value for key, value in update_data.items() if getattr(message, key) != value}
+
+        for key, value in update_data.items():
+            setattr(message, key, value)
+        return message
+
     @enforce_types
     def delete_message_by_id(self, message_id: str, actor: PydanticUser) -> bool:
         """Delete a message."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 msg = MessageModel.read(
                     db_session=session,
@@ -229,7 +251,7 @@ class MessageManager:
             actor: The user requesting the count
             role: The role of the message
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return MessageModel.size(db_session=session, actor=actor, role=role, agent_id=agent_id)
 
     @enforce_types
@@ -293,7 +315,7 @@ class MessageManager:
             NoResultFound: If the provided after/before message IDs do not exist.
         """
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Permission check: raise if the agent doesn't exist or actor is not allowed.
             AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
 
@@ -363,7 +385,7 @@ class MessageManager:
         Efficiently deletes all messages associated with a given agent_id,
         while enforcing permission checks and avoiding any ORM‑level loads.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # 1) verify the agent exists and the actor has access
             AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
 

letta/services/organization_manager.py

@@ -4,6 +4,7 @@ from letta.orm.errors import NoResultFound
 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.utils import enforce_types
 
 
@@ -13,14 +14,6 @@ class OrganizationManager:
     DEFAULT_ORG_ID = "org-00000000-0000-4000-8000-000000000000"
     DEFAULT_ORG_NAME = "default_org"
 
-    def __init__(self):
-        # TODO: Please refactor this out
-        # I am currently working on a ORM refactor and would like to make a more minimal set of changes
-        # - Matt
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def get_default_organization(self) -> PydanticOrganization:
         """Fetch the default organization."""
@@ -29,7 +22,7 @@ class OrganizationManager:
     @enforce_types
     def get_organization_by_id(self, org_id: str) -> Optional[PydanticOrganization]:
         """Fetch an organization by ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             organization = OrganizationModel.read(db_session=session, identifier=org_id)
             return organization.to_pydantic()
 
@@ -44,7 +37,7 @@ class OrganizationManager:
 
     @enforce_types
     def _create_organization(self, pydantic_org: PydanticOrganization) -> PydanticOrganization:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             org = OrganizationModel(**pydantic_org.model_dump(to_orm=True))
             org.create(session)
             return org.to_pydantic()
@@ -57,7 +50,7 @@ class OrganizationManager:
     @enforce_types
     def update_organization_name_using_id(self, org_id: str, name: Optional[str] = None) -> PydanticOrganization:
         """Update an organization."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             org = OrganizationModel.read(db_session=session, identifier=org_id)
             if name:
                 org.name = name
@@ -67,7 +60,7 @@ class OrganizationManager:
     @enforce_types
     def update_organization(self, org_id: str, org_update: OrganizationUpdate) -> PydanticOrganization:
         """Update an organization."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             org = OrganizationModel.read(db_session=session, identifier=org_id)
             if org_update.name:
                 org.name = org_update.name
@@ -79,14 +72,14 @@ class OrganizationManager:
     @enforce_types
     def delete_organization_by_id(self, org_id: str):
         """Delete an organization by marking it as deleted."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             organization = OrganizationModel.read(db_session=session, identifier=org_id)
             organization.hard_delete(session)
 
     @enforce_types
     def list_organizations(self, after: Optional[str] = None, limit: Optional[int] = 50) -> List[PydanticOrganization]:
         """List all organizations with optional pagination."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             organizations = OrganizationModel.list(
                 db_session=session,
                 after=after,

letta/services/passage_manager.py

@@ -10,21 +10,17 @@ from letta.orm.passage import AgentPassage, SourcePassage
 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.utils import enforce_types
 
 
 class PassageManager:
     """Manager class to handle business logic related to Passages."""
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def get_passage_by_id(self, passage_id: str, actor: PydanticUser) -> Optional[PydanticPassage]:
         """Fetch a passage by ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Try source passages first
             try:
                 passage = SourcePassage.read(db_session=session, identifier=passage_id, actor=actor)
@@ -69,7 +65,7 @@ class PassageManager:
         else:
             raise ValueError("Passage must have either agent_id or source_id")
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             passage.create(session, actor=actor)
             return passage.to_pydantic()
 
@@ -145,7 +141,7 @@ class PassageManager:
         if not passage_id:
             raise ValueError("Passage ID must be provided.")
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Try source passages first
             try:
                 curr_passage = SourcePassage.read(
@@ -179,7 +175,7 @@ class PassageManager:
         if not passage_id:
             raise ValueError("Passage ID must be provided.")
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Try source passages first
             try:
                 passage = SourcePassage.read(db_session=session, identifier=passage_id, actor=actor)
@@ -217,7 +213,7 @@ class PassageManager:
             actor: The user requesting the count
             agent_id: The agent ID of the messages
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return AgentPassage.size(db_session=session, actor=actor, agent_id=agent_id)
 
     def estimate_embeddings_size(

letta/services/provider_manager.py

@@ -5,20 +5,16 @@ from letta.schemas.enums import ProviderCategory, ProviderType
 from letta.schemas.providers import Provider as PydanticProvider
 from letta.schemas.providers import ProviderCheck, ProviderCreate, ProviderUpdate
 from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
 from letta.utils import enforce_types
 
 
 class ProviderManager:
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def create_provider(self, request: ProviderCreate, actor: PydanticUser) -> PydanticProvider:
         """Create a new provider if it doesn't already exist."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             provider_create_args = {**request.model_dump(), "provider_category": ProviderCategory.byok}
             provider = PydanticProvider(**provider_create_args)
 
@@ -38,7 +34,7 @@ class ProviderManager:
     @enforce_types
     def update_provider(self, provider_id: str, provider_update: ProviderUpdate, actor: PydanticUser) -> PydanticProvider:
         """Update provider details."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Retrieve the existing provider by ID
             existing_provider = ProviderModel.read(db_session=session, identifier=provider_id, actor=actor)
 
@@ -54,7 +50,7 @@ class ProviderManager:
     @enforce_types
     def delete_provider_by_id(self, provider_id: str, actor: PydanticUser):
         """Delete a provider."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Clear api key field
             existing_provider = ProviderModel.read(db_session=session, identifier=provider_id, actor=actor)
             existing_provider.api_key = None
@@ -80,7 +76,7 @@ class ProviderManager:
             filter_kwargs["name"] = name
         if provider_type:
             filter_kwargs["provider_type"] = provider_type
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             providers = ProviderModel.list(
                 db_session=session,
                 after=after,

letta/services/sandbox_config_manager.py

@@ -11,6 +11,7 @@ from letta.schemas.sandbox_config import LocalSandboxConfig
 from letta.schemas.sandbox_config import SandboxConfig as PydanticSandboxConfig
 from letta.schemas.sandbox_config import SandboxConfigCreate, SandboxConfigUpdate, SandboxType
 from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
 from letta.utils import enforce_types, printd
 
 logger = get_logger(__name__)
@@ -19,11 +20,6 @@ logger = get_logger(__name__)
 class SandboxConfigManager:
     """Manager class to handle business logic related to SandboxConfig and SandboxEnvironmentVariable."""
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def get_or_create_default_sandbox_config(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
         sandbox_config = self.get_sandbox_config_by_type(sandbox_type, actor=actor)
@@ -69,7 +65,7 @@ class SandboxConfigManager:
             return db_sandbox
         else:
             # If the sandbox configuration doesn't exist, create a new one
-            with self.session_maker() as session:
+            with db_registry.session() as session:
                 db_sandbox = SandboxConfigModel(**sandbox_config.model_dump(exclude_none=True))
                 db_sandbox.create(session, actor=actor)
                 return db_sandbox.to_pydantic()
@@ -79,7 +75,7 @@ class SandboxConfigManager:
         self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
     ) -> PydanticSandboxConfig:
         """Update an existing sandbox configuration."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
             # We need to check that the sandbox_update provided is the same type as the original sandbox
             if sandbox.type != sandbox_update.config.type:
@@ -104,7 +100,7 @@ class SandboxConfigManager:
     @enforce_types
     def delete_sandbox_config(self, sandbox_config_id: str, actor: PydanticUser) -> PydanticSandboxConfig:
         """Delete a sandbox configuration by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
             sandbox.hard_delete(db_session=session, actor=actor)
             return sandbox.to_pydantic()
@@ -122,14 +118,14 @@ class SandboxConfigManager:
         if sandbox_type:
             kwargs.update({"type": sandbox_type})
 
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             sandboxes = SandboxConfigModel.list(db_session=session, after=after, limit=limit, **kwargs)
             return [sandbox.to_pydantic() for sandbox in sandboxes]
 
     @enforce_types
     def get_sandbox_config_by_id(self, sandbox_config_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
         """Retrieve a sandbox configuration by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 sandbox = SandboxConfigModel.read(db_session=session, identifier=sandbox_config_id, actor=actor)
                 return sandbox.to_pydantic()
@@ -139,7 +135,7 @@ class SandboxConfigManager:
     @enforce_types
     def get_sandbox_config_by_type(self, type: SandboxType, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
         """Retrieve a sandbox config by its type."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 sandboxes = SandboxConfigModel.list(
                     db_session=session,
@@ -175,7 +171,7 @@ class SandboxConfigManager:
 
             return db_env_var
         else:
-            with self.session_maker() as session:
+            with db_registry.session() as session:
                 env_var = SandboxEnvVarModel(**env_var.model_dump(to_orm=True, exclude_none=True))
                 env_var.create(session, actor=actor)
             return env_var.to_pydantic()
@@ -185,7 +181,7 @@ class SandboxConfigManager:
         self, env_var_id: str, env_var_update: SandboxEnvironmentVariableUpdate, actor: PydanticUser
     ) -> PydanticEnvVar:
         """Update an existing sandbox environment variable."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             env_var = SandboxEnvVarModel.read(db_session=session, identifier=env_var_id, actor=actor)
             update_data = env_var_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
             update_data = {key: value for key, value in update_data.items() if getattr(env_var, key) != value}
@@ -204,7 +200,7 @@ class SandboxConfigManager:
     @enforce_types
     def delete_sandbox_env_var(self, env_var_id: str, actor: PydanticUser) -> PydanticEnvVar:
         """Delete a sandbox environment variable by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             env_var = SandboxEnvVarModel.read(db_session=session, identifier=env_var_id, actor=actor)
             env_var.hard_delete(db_session=session, actor=actor)
             return env_var.to_pydantic()
@@ -218,7 +214,7 @@ class SandboxConfigManager:
         limit: Optional[int] = 50,
     ) -> List[PydanticEnvVar]:
         """List all sandbox environment variables with optional pagination."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             env_vars = SandboxEnvVarModel.list(
                 db_session=session,
                 after=after,
@@ -233,7 +229,7 @@ class SandboxConfigManager:
         self, key: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
     ) -> List[PydanticEnvVar]:
         """List all sandbox environment variables with optional pagination."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             env_vars = SandboxEnvVarModel.list(
                 db_session=session,
                 after=after,
@@ -258,7 +254,7 @@ class SandboxConfigManager:
         self, key: str, sandbox_config_id: str, actor: Optional[PydanticUser] = None
     ) -> Optional[PydanticEnvVar]:
         """Retrieve a sandbox environment variable by its key and sandbox_config_id."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 env_var = SandboxEnvVarModel.list(
                     db_session=session,

letta/services/source_manager.py

@@ -8,17 +8,13 @@ from letta.schemas.file import FileMetadata as PydanticFileMetadata
 from letta.schemas.source import Source as PydanticSource
 from letta.schemas.source import SourceUpdate
 from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
 from letta.utils import enforce_types, printd
 
 
 class SourceManager:
     """Manager class to handle business logic related to Sources."""
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def create_source(self, source: PydanticSource, actor: PydanticUser) -> PydanticSource:
         """Create a new source based on the PydanticSource schema."""
@@ -27,7 +23,7 @@ class SourceManager:
         if db_source:
             return db_source
         else:
-            with self.session_maker() as session:
+            with db_registry.session() as session:
                 # Provide default embedding config if not given
                 source.organization_id = actor.organization_id
                 source = SourceModel(**source.model_dump(to_orm=True, exclude_none=True))
@@ -37,7 +33,7 @@ class SourceManager:
     @enforce_types
     def update_source(self, source_id: str, source_update: SourceUpdate, actor: PydanticUser) -> PydanticSource:
         """Update a source by its ID with the given SourceUpdate object."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
 
             # get update dictionary
@@ -59,7 +55,7 @@ class SourceManager:
     @enforce_types
     def delete_source(self, source_id: str, actor: PydanticUser) -> PydanticSource:
         """Delete a source by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             source = SourceModel.read(db_session=session, identifier=source_id)
             source.hard_delete(db_session=session, actor=actor)
             return source.to_pydantic()
@@ -67,7 +63,7 @@ class SourceManager:
     @enforce_types
     def list_sources(self, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50, **kwargs) -> List[PydanticSource]:
         """List all sources with optional pagination."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             sources = SourceModel.list(
                 db_session=session,
                 after=after,
@@ -85,7 +81,7 @@ class SourceManager:
         """
         Get the total count of sources for the given user.
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             return SourceModel.size(db_session=session, actor=actor)
 
     @enforce_types
@@ -100,7 +96,7 @@ class SourceManager:
         Returns:
             List[PydanticAgentState]: List of agents that have this source attached
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             # Verify source exists and user has permission to access it
             source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
 
@@ -112,7 +108,7 @@ class SourceManager:
     @enforce_types
     def get_source_by_id(self, source_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSource]:
         """Retrieve a source by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
                 return source.to_pydantic()
@@ -122,7 +118,7 @@ class SourceManager:
     @enforce_types
     def get_source_by_name(self, source_name: str, actor: PydanticUser) -> Optional[PydanticSource]:
         """Retrieve a source by its name."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             sources = SourceModel.list(
                 db_session=session,
                 name=source_name,
@@ -141,7 +137,7 @@ class SourceManager:
         if db_file:
             return db_file
         else:
-            with self.session_maker() as session:
+            with db_registry.session() as session:
                 file_metadata.organization_id = actor.organization_id
                 file_metadata = FileMetadataModel(**file_metadata.model_dump(to_orm=True, exclude_none=True))
                 file_metadata.create(session, actor=actor)
@@ -151,7 +147,7 @@ class SourceManager:
     @enforce_types
     def get_file_by_id(self, file_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticFileMetadata]:
         """Retrieve a file by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             try:
                 file = FileMetadataModel.read(db_session=session, identifier=file_id, actor=actor)
                 return file.to_pydantic()
@@ -163,7 +159,7 @@ class SourceManager:
         self, source_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
     ) -> List[PydanticFileMetadata]:
         """List all files with optional pagination."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             files = FileMetadataModel.list(
                 db_session=session, after=after, limit=limit, organization_id=actor.organization_id, source_id=source_id
             )
@@ -172,7 +168,7 @@ class SourceManager:
     @enforce_types
     def delete_file(self, file_id: str, actor: PydanticUser) -> PydanticFileMetadata:
         """Delete a file by its ID."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             file = FileMetadataModel.read(db_session=session, identifier=file_id)
             file.hard_delete(db_session=session, actor=actor)
             return file.to_pydantic()

letta/services/step_manager.py

@@ -11,17 +11,13 @@ from letta.orm.step import Step as StepModel
 from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.step import Step as PydanticStep
 from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
 from letta.tracing import get_trace_id
 from letta.utils import enforce_types
 
 
 class StepManager:
 
-    def __init__(self):
-        from letta.server.db import db_context
-
-        self.session_maker = db_context
-
     @enforce_types
     def list_steps(
         self,
@@ -36,7 +32,7 @@ class StepManager:
         agent_id: Optional[str] = None,
     ) -> List[PydanticStep]:
         """List all jobs with optional pagination and status filter."""
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             filter_kwargs = {"organization_id": actor.organization_id}
             if model:
                 filter_kwargs["model"] = model
@@ -85,7 +81,7 @@ class StepManager:
             "tid": None,
             "trace_id": get_trace_id(),  # Get the current trace ID
         }
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             if job_id:
                 self._verify_job_access(session, job_id, actor, access=["write"])
             new_step = StepModel(**step_data)
@@ -94,7 +90,7 @@ class StepManager:
 
     @enforce_types
     def get_step(self, step_id: str, actor: PydanticUser) -> PydanticStep:
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             step = StepModel.read(db_session=session, identifier=step_id, actor=actor)
             return step.to_pydantic()
 
@@ -113,7 +109,7 @@ class StepManager:
         Raises:
             NoResultFound: If the step does not exist
         """
-        with self.session_maker() as session:
+        with db_registry.session() as session:
             step = session.get(StepModel, step_id)
             if not step:
                 raise NoResultFound(f"Step with id {step_id} does not exist")