letta-nightly 0.7.20.dev20250521104258__py3-none-any.whl → 0.7.21.dev20250521233415__py3-none-any.whl

letta/services/block_manager.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Dict, List, Optional
 
 from sqlalchemy import select
@@ -82,39 +83,64 @@ class BlockManager:
             return block.to_pydantic()
 
     @enforce_types
-    def get_blocks(
+    async def get_blocks_async(
         self,
         actor: PydanticUser,
         label: Optional[str] = None,
         is_template: Optional[bool] = None,
         template_name: Optional[str] = None,
-        identifier_keys: Optional[List[str]] = None,
         identity_id: Optional[str] = None,
-        id: Optional[str] = None,
-        after: Optional[str] = None,
+        identifier_keys: Optional[List[str]] = None,
         limit: Optional[int] = 50,
     ) -> List[PydanticBlock]:
-        """Retrieve blocks based on various optional filters."""
-        with db_registry.session() as session:
-            # Prepare filters
-            filters = {"organization_id": actor.organization_id}
+        """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:
-                filters["label"] = label
+                query = query.where(BlockModel.label == label)
+
             if is_template is not None:
-                filters["is_template"] = is_template
+                query = query.where(BlockModel.is_template == is_template)
+
             if template_name:
-                filters["template_name"] = template_name
-            if id:
-                filters["id"] = id
-
-            blocks = BlockModel.list(
-                db_session=session,
-                after=after,
-                limit=limit,
-                identifier_keys=identifier_keys,
-                identity_id=identity_id,
-                **filters,
-            )
+                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]
 
@@ -190,15 +216,6 @@ class BlockManager:
             except NoResultFound:
                 return None
 
-    @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 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
     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."""
@@ -247,16 +264,14 @@ class BlockManager:
             return pydantic_blocks
 
     @enforce_types
-    def get_agents_for_block(self, block_id: str, actor: PydanticUser) -> List[PydanticAgentState]:
+    async def get_agents_for_block_async(self, block_id: str, actor: PydanticUser) -> List[PydanticAgentState]:
         """
         Retrieve all agents associated with a given block.
         """
-        with db_registry.session() as session:
-            block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)
+        async with db_registry.async_session() as session:
+            block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
             agents_orm = block.agents
-            agents_pydantic = [agent.to_pydantic() for agent in agents_orm]
-
-            return agents_pydantic
+            return await asyncio.gather(*[agent.to_pydantic_async() for agent in agents_orm])
 
     @enforce_types
     def size(

letta/services/helpers/noop_helper.py

@@ -0,0 +1,10 @@
+def singleton(cls):
+    """Decorator to make a class a Singleton class."""
+    instances = {}
+
+    def get_instance(*args, **kwargs):
+        if cls not in instances:
+            instances[cls] = cls(*args, **kwargs)
+        return instances[cls]
+
+    return get_instance

letta/services/identity_manager.py

@@ -1,6 +1,7 @@
 from typing import List, Optional
 
 from fastapi import HTTPException
+from sqlalchemy import select
 from sqlalchemy.exc import NoResultFound
 from sqlalchemy.orm import Session
 
@@ -17,7 +18,7 @@ from letta.utils import enforce_types
 class IdentityManager:
 
     @enforce_types
-    def list_identities(
+    async def list_identities_async(
         self,
         name: Optional[str] = None,
         project_id: Optional[str] = None,
@@ -28,7 +29,7 @@ class IdentityManager:
         limit: Optional[int] = 50,
         actor: PydanticUser = None,
     ) -> list[PydanticIdentity]:
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             filters = {"organization_id": actor.organization_id}
             if project_id:
                 filters["project_id"] = project_id
@@ -36,7 +37,7 @@ class IdentityManager:
                 filters["identifier_key"] = identifier_key
             if identity_type:
                 filters["identity_type"] = identity_type
-            identities = IdentityModel.list(
+            identities = await IdentityModel.list_async(
                 db_session=session,
                 query_text=name,
                 before=before,
@@ -47,17 +48,17 @@ class IdentityManager:
             return [identity.to_pydantic() for identity in identities]
 
     @enforce_types
-    def get_identity(self, identity_id: str, actor: PydanticUser) -> PydanticIdentity:
-        with db_registry.session() as session:
-            identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
+    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
-    def create_identity(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
-        with db_registry.session() as session:
+    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))
             new_identity.organization_id = actor.organization_id
-            self._process_relationship(
+            await self._process_relationship_async(
                 session=session,
                 identity=new_identity,
                 relationship_name="agents",
@@ -65,7 +66,7 @@ class IdentityManager:
                 item_ids=identity.agent_ids,
                 allow_partial=False,
             )
-            self._process_relationship(
+            await self._process_relationship_async(
                 session=session,
                 identity=new_identity,
                 relationship_name="blocks",
@@ -73,13 +74,13 @@ class IdentityManager:
                 item_ids=identity.block_ids,
                 allow_partial=False,
             )
-            new_identity.create(session, actor=actor)
+            await new_identity.create_async(session, actor=actor)
             return new_identity.to_pydantic()
 
     @enforce_types
-    def upsert_identity(self, identity: IdentityUpsert, actor: PydanticUser) -> PydanticIdentity:
-        with db_registry.session() as session:
-            existing_identity = IdentityModel.read(
+    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(
                 db_session=session,
                 identifier_key=identity.identifier_key,
                 project_id=identity.project_id,
@@ -88,7 +89,7 @@ class IdentityManager:
             )
 
         if existing_identity is None:
-            return self.create_identity(identity=IdentityCreate(**identity.model_dump()), actor=actor)
+            return await self.create_identity_async(identity=IdentityCreate(**identity.model_dump()), actor=actor)
         else:
             identity_update = IdentityUpdate(
                 name=identity.name,
@@ -97,25 +98,27 @@ class IdentityManager:
                 agent_ids=identity.agent_ids,
                 properties=identity.properties,
             )
-            return self._update_identity(
+            return await self._update_identity_async(
                 session=session, existing_identity=existing_identity, identity=identity_update, actor=actor, replace=True
             )
 
     @enforce_types
-    def update_identity(self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False) -> PydanticIdentity:
-        with db_registry.session() as session:
+    async def update_identity_async(
+        self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False
+    ) -> PydanticIdentity:
+        async with db_registry.async_session() as session:
             try:
-                existing_identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
+                existing_identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
             except NoResultFound:
                 raise HTTPException(status_code=404, detail="Identity not found")
             if existing_identity.organization_id != actor.organization_id:
                 raise HTTPException(status_code=403, detail="Forbidden")
 
-            return self._update_identity(
+            return await self._update_identity_async(
                 session=session, existing_identity=existing_identity, identity=identity, actor=actor, replace=replace
             )
 
-    def _update_identity(
+    async def _update_identity_async(
         self,
         session: Session,
         existing_identity: IdentityModel,
@@ -139,7 +142,7 @@ class IdentityManager:
                 existing_identity.properties = list(new_properties.values())
 
         if identity.agent_ids is not None:
-            self._process_relationship(
+            await self._process_relationship_async(
                 session=session,
                 identity=existing_identity,
                 relationship_name="agents",
@@ -149,7 +152,7 @@ class IdentityManager:
                 replace=replace,
             )
         if identity.block_ids is not None:
-            self._process_relationship(
+            await self._process_relationship_async(
                 session=session,
                 identity=existing_identity,
                 relationship_name="blocks",
@@ -158,16 +161,18 @@ class IdentityManager:
                 allow_partial=False,
                 replace=replace,
             )
-        existing_identity.update(session, actor=actor)
+        await existing_identity.update_async(session, actor=actor)
         return existing_identity.to_pydantic()
 
     @enforce_types
-    def upsert_identity_properties(self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser) -> PydanticIdentity:
-        with db_registry.session() as session:
-            existing_identity = IdentityModel.read(db_session=session, identifier=identity_id, actor=actor)
+    async def upsert_identity_properties_async(
+        self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser
+    ) -> PydanticIdentity:
+        async with db_registry.async_session() as session:
+            existing_identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
             if existing_identity is None:
                 raise HTTPException(status_code=404, detail="Identity not found")
-            return self._update_identity(
+            return await self._update_identity_async(
                 session=session,
                 existing_identity=existing_identity,
                 identity=IdentityUpdate(properties=properties),
@@ -176,28 +181,28 @@ class IdentityManager:
             )
 
     @enforce_types
-    def delete_identity(self, identity_id: str, actor: PydanticUser) -> None:
-        with db_registry.session() as session:
-            identity = IdentityModel.read(db_session=session, identifier=identity_id)
+    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)
             if identity is None:
                 raise HTTPException(status_code=404, detail="Identity not found")
             if identity.organization_id != actor.organization_id:
                 raise HTTPException(status_code=403, detail="Forbidden")
-            session.delete(identity)
-            session.commit()
+            await session.delete(identity)
+            await session.commit()
 
     @enforce_types
-    def size(
+    async def size_async(
         self,
         actor: PydanticUser,
     ) -> int:
         """
         Get the total count of identities for the given user.
         """
-        with db_registry.session() as session:
-            return IdentityModel.size(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            return await IdentityModel.size_async(db_session=session, actor=actor)
 
-    def _process_relationship(
+    async def _process_relationship_async(
         self,
         session: Session,
         identity: PydanticIdentity,
@@ -214,7 +219,7 @@ class IdentityManager:
             return
 
         # Retrieve models for the provided IDs
-        found_items = session.query(model_class).filter(model_class.id.in_(item_ids)).all()
+        found_items = (await session.execute(select(model_class).where(model_class.id.in_(item_ids)))).scalars().all()
 
         # Validate all items are found if allow_partial is False
         if not allow_partial and len(found_items) != len(item_ids):

letta/services/job_manager.py

@@ -150,6 +150,35 @@ class JobManager:
             )
             return [job.to_pydantic() for job in jobs]
 
+    @enforce_types
+    async def list_jobs_async(
+        self,
+        actor: PydanticUser,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+        statuses: Optional[List[JobStatus]] = None,
+        job_type: JobType = JobType.JOB,
+        ascending: bool = True,
+    ) -> List[PydanticJob]:
+        """List all jobs with optional pagination and status filter."""
+        async with db_registry.async_session() as session:
+            filter_kwargs = {"user_id": actor.id, "job_type": job_type}
+
+            # Add status filter if provided
+            if statuses:
+                filter_kwargs["status"] = statuses
+
+            jobs = await JobModel.list_async(
+                db_session=session,
+                before=before,
+                after=after,
+                limit=limit,
+                ascending=ascending,
+                **filter_kwargs,
+            )
+            return [job.to_pydantic() for job in jobs]
+
     @enforce_types
     def delete_job_by_id(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Delete a job by its ID."""

letta/services/message_manager.py

@@ -31,6 +31,16 @@ class MessageManager:
             except NoResultFound:
                 return None
 
+    @enforce_types
+    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:
+            try:
+                message = await MessageModel.read_async(db_session=session, identifier=message_id, actor=actor)
+                return message.to_pydantic()
+            except NoResultFound:
+                return None
+
     @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."""
@@ -426,6 +436,107 @@ class MessageManager:
             results = query.all()
             return [msg.to_pydantic() for msg in results]
 
+    @enforce_types
+    async def list_messages_for_agent_async(
+        self,
+        agent_id: str,
+        actor: PydanticUser,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+        query_text: Optional[str] = None,
+        roles: Optional[Sequence[MessageRole]] = None,
+        limit: Optional[int] = 50,
+        ascending: bool = True,
+        group_id: Optional[str] = None,
+    ) -> List[PydanticMessage]:
+        """
+        Most performant query to list messages for an agent by directly querying the Message table.
+
+        This function filters by the agent_id (leveraging the index on messages.agent_id)
+        and applies pagination using sequence_id as the cursor.
+        If query_text is provided, it will filter messages whose text content partially matches the query.
+        If role is provided, it will filter messages by the specified role.
+
+        Args:
+            agent_id: The ID of the agent whose messages are queried.
+            actor: The user performing the action (used for permission checks).
+            after: A message ID; if provided, only messages *after* this message (by sequence_id) are returned.
+            before: A message ID; if provided, only messages *before* this message (by sequence_id) are returned.
+            query_text: Optional string to partially match the message text content.
+            roles: Optional MessageRole to filter messages by role.
+            limit: Maximum number of messages to return.
+            ascending: If True, sort by sequence_id ascending; if False, sort descending.
+            group_id: Optional group ID to filter messages by group_id.
+
+        Returns:
+            List[PydanticMessage]: A list of messages (converted via .to_pydantic()).
+
+        Raises:
+            NoResultFound: If the provided after/before message IDs do not exist.
+        """
+
+        async with db_registry.async_session() as session:
+            # Permission check: raise if the agent doesn't exist or actor is not allowed.
+            await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
+
+            # Build a query that directly filters the Message table by agent_id.
+            query = select(MessageModel).where(MessageModel.agent_id == agent_id)
+
+            # If group_id is provided, filter messages by group_id.
+            if group_id:
+                query = query.where(MessageModel.group_id == group_id)
+
+            # If query_text is provided, filter messages using subquery + json_array_elements.
+            if query_text:
+                content_element = func.json_array_elements(MessageModel.content).alias("content_element")
+                query = query.where(
+                    exists(
+                        select(1)
+                        .select_from(content_element)
+                        .where(text("content_element->>'type' = 'text' AND content_element->>'text' ILIKE :query_text"))
+                        .params(query_text=f"%{query_text}%")
+                    )
+                )
+
+            # If role(s) are provided, filter messages by those roles.
+            if roles:
+                role_values = [r.value for r in roles]
+                query = query.where(MessageModel.role.in_(role_values))
+
+            # Apply 'after' pagination if specified.
+            if after:
+                after_query = select(MessageModel.sequence_id).where(MessageModel.id == after)
+                after_result = await session.execute(after_query)
+                after_ref = after_result.one_or_none()
+                if not after_ref:
+                    raise NoResultFound(f"No message found with id '{after}' for agent '{agent_id}'.")
+                # Filter out any messages with a sequence_id <= after_ref.sequence_id
+                query = query.where(MessageModel.sequence_id > after_ref.sequence_id)
+
+            # Apply 'before' pagination if specified.
+            if before:
+                before_query = select(MessageModel.sequence_id).where(MessageModel.id == before)
+                before_result = await session.execute(before_query)
+                before_ref = before_result.one_or_none()
+                if not before_ref:
+                    raise NoResultFound(f"No message found with id '{before}' for agent '{agent_id}'.")
+                # Filter out any messages with a sequence_id >= before_ref.sequence_id
+                query = query.where(MessageModel.sequence_id < before_ref.sequence_id)
+
+            # Apply ordering based on the ascending flag.
+            if ascending:
+                query = query.order_by(MessageModel.sequence_id.asc())
+            else:
+                query = query.order_by(MessageModel.sequence_id.desc())
+
+            # Limit the number of results.
+            query = query.limit(limit)
+
+            # Execute and convert each Message to its Pydantic representation.
+            result = await session.execute(query)
+            results = result.scalars().all()
+            return [msg.to_pydantic() for msg in results]
+
     @enforce_types
     def delete_all_messages_for_agent(self, agent_id: str, actor: PydanticUser) -> int:
         """

letta/services/sandbox_config_manager.py

@@ -122,6 +122,23 @@ class SandboxConfigManager:
             sandboxes = SandboxConfigModel.list(db_session=session, after=after, limit=limit, **kwargs)
             return [sandbox.to_pydantic() for sandbox in sandboxes]
 
+    @enforce_types
+    async def list_sandbox_configs_async(
+        self,
+        actor: PydanticUser,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+        sandbox_type: Optional[SandboxType] = None,
+    ) -> List[PydanticSandboxConfig]:
+        """List all sandbox configurations with optional pagination."""
+        kwargs = {"organization_id": actor.organization_id}
+        if sandbox_type:
+            kwargs.update({"type": sandbox_type})
+
+        async with db_registry.async_session() as session:
+            sandboxes = await SandboxConfigModel.list_async(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."""
@@ -224,6 +241,25 @@ class SandboxConfigManager:
             )
             return [env_var.to_pydantic() for env_var in env_vars]
 
+    @enforce_types
+    async def list_sandbox_env_vars_async(
+        self,
+        sandbox_config_id: str,
+        actor: PydanticUser,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+    ) -> List[PydanticEnvVar]:
+        """List all sandbox environment variables with optional pagination."""
+        async with db_registry.async_session() as session:
+            env_vars = await SandboxEnvVarModel.list_async(
+                db_session=session,
+                after=after,
+                limit=limit,
+                organization_id=actor.organization_id,
+                sandbox_config_id=sandbox_config_id,
+            )
+            return [env_var.to_pydantic() for env_var in env_vars]
+
     @enforce_types
     def list_sandbox_env_vars_by_key(
         self, key: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50