letta-nightly 0.8.15.dev20250719104256__py3-none-any.whl → 0.8.16.dev20250721070720__py3-none-any.whl

letta/services/agent_manager.py

@@ -28,7 +28,7 @@ from letta.orm import AgentPassage, AgentsTags
 from letta.orm import Block as BlockModel
 from letta.orm import BlocksAgents
 from letta.orm import Group as GroupModel
-from letta.orm import IdentitiesAgents
+from letta.orm import GroupsAgents, IdentitiesAgents
 from letta.orm import Source as SourceModel
 from letta.orm import SourcePassage, SourcesAgents
 from letta.orm import Tool as ToolModel
@@ -71,6 +71,7 @@ from letta.services.helpers.agent_manager_helper import (
     _apply_identity_filters,
     _apply_pagination,
     _apply_pagination_async,
+    _apply_relationship_filters,
     _apply_tag_filter,
     _process_relationship,
     _process_relationship_async,
@@ -84,12 +85,14 @@ from letta.services.helpers.agent_manager_helper import (
     derive_system_message,
     initialize_message_sequence,
     package_initial_message_sequence,
+    validate_agent_exists_async,
 )
 from letta.services.identity_manager import IdentityManager
 from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
+from letta.settings import DatabaseChoice, settings
 from letta.utils import enforce_types, united_diff
 
 logger = get_logger(__name__)
@@ -448,7 +451,11 @@ class AgentManager:
 
     @trace_method
     async def create_agent_async(
-        self, agent_create: CreateAgent, actor: PydanticUser, _test_only_force_id: Optional[str] = None
+        self,
+        agent_create: CreateAgent,
+        actor: PydanticUser,
+        _test_only_force_id: Optional[str] = None,
+        _init_with_no_messages: bool = False,
     ) -> PydanticAgentState:
         # validate required configs
         if not agent_create.llm_config or not agent_create.embedding_config:
@@ -616,20 +623,36 @@ class AgentManager:
                     ]
                     await session.execute(insert(AgentEnvironmentVariable).values(env_rows))
 
-                # initial message sequence
-                agent_state = await new_agent.to_pydantic_async(include_relationships={"memory"})
-                init_messages = self._generate_initial_message_sequence(
-                    actor,
-                    agent_state=agent_state,
-                    supplied_initial_message_sequence=agent_create.initial_message_sequence,
-                )
-                new_agent.message_ids = [msg.id for msg in init_messages]
-
-            await session.refresh(new_agent)
-
-            result = await new_agent.to_pydantic_async()
+                include_relationships = []
+                if tool_ids:
+                    include_relationships.append("tools")
+                if source_ids:
+                    include_relationships.append("sources")
+                if block_ids:
+                    include_relationships.append("memory")
+                if identity_ids:
+                    include_relationships.append("identity_ids")
+                if tag_values:
+                    include_relationships.append("tags")
+
+                result = await new_agent.to_pydantic_async(include_relationships=include_relationships)
+
+                # initial message sequence (skip if _init_with_no_messages is True)
+                if not _init_with_no_messages:
+                    init_messages = self._generate_initial_message_sequence(
+                        actor,
+                        agent_state=result,
+                        supplied_initial_message_sequence=agent_create.initial_message_sequence,
+                    )
+                    result.message_ids = [msg.id for msg in init_messages]
+                    new_agent.message_ids = [msg.id for msg in init_messages]
+                    await new_agent.update_async(session, no_refresh=True)
+                else:
+                    init_messages = []
 
-        await self.message_manager.create_many_messages_async(pydantic_msgs=init_messages, actor=actor)
+        # Only create messages if we initialized with messages
+        if not _init_with_no_messages:
+            await self.message_manager.create_many_messages_async(pydantic_msgs=init_messages, actor=actor)
         return result
 
     @enforce_types
@@ -920,6 +943,30 @@ class AgentManager:
 
             return await agent.to_pydantic_async()
 
+    @enforce_types
+    @trace_method
+    async def update_message_ids_async(
+        self,
+        agent_id: str,
+        message_ids: List[str],
+        actor: PydanticUser,
+    ) -> None:
+        async with db_registry.async_session() as session:
+            query = select(AgentModel)
+            query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+            query = query.where(AgentModel.id == agent_id)
+            query = _apply_relationship_filters(query, include_relationships=[])
+
+            result = await session.execute(query)
+            agent = result.scalar_one_or_none()
+
+            agent.updated_at = datetime.now(timezone.utc)
+            agent.last_updated_by_id = actor.id
+            agent.message_ids = message_ids
+
+            await agent.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True)
+            await session.commit()
+
     # TODO: Make this general and think about how to roll this into sqlalchemybase
     @trace_method
     def list_agents(
@@ -1019,7 +1066,8 @@ class AgentManager:
             identity_id (Optional[str]): Filter by identifier ID.
             identifier_keys (Optional[List[str]]): Search agents by identifier keys.
             include_relationships (Optional[List[str]]): List of fields to load for performance optimization.
-            ascending
+            ascending (bool): Sort agents in ascending order.
+            sort_by (Optional[str]): Sort agents by this field.
 
         Returns:
             List[PydanticAgentState]: The filtered list of matching agents.
@@ -1032,11 +1080,11 @@ class AgentManager:
             query = _apply_filters(query, name, query_text, project_id, template_id, base_template_id)
             query = _apply_identity_filters(query, identity_id, identifier_keys)
             query = _apply_tag_filter(query, tags, match_all_tags)
+            query = _apply_relationship_filters(query, include_relationships)
             query = await _apply_pagination_async(query, before, after, session, ascending=ascending, sort_by=sort_by)
 
             if limit:
                 query = query.limit(limit)
-
             result = await session.execute(query)
             agents = result.scalars().all()
             return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents])
@@ -1168,10 +1216,23 @@ class AgentManager:
         include_relationships: Optional[List[str]] = None,
     ) -> PydanticAgentState:
         """Fetch an agent by its ID."""
-
         async with db_registry.async_session() as session:
-            agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
-            return await agent.to_pydantic_async(include_relationships=include_relationships)
+            try:
+                query = select(AgentModel)
+                query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+                query = query.where(AgentModel.id == agent_id)
+                query = _apply_relationship_filters(query, include_relationships)
+
+                result = await session.execute(query)
+                agent = result.scalar_one_or_none()
+
+                if agent is None:
+                    raise NoResultFound(f"Agent with ID {agent_id} not found")
+
+                return await agent.to_pydantic_async(include_relationships=include_relationships)
+            except Exception as e:
+                logger.error(f"Error fetching agent {agent_id}: {str(e)}")
+                raise
 
     @enforce_types
     @trace_method
@@ -1183,12 +1244,23 @@ class AgentManager:
     ) -> list[PydanticAgentState]:
         """Fetch a list of agents by their IDs."""
         async with db_registry.async_session() as session:
-            agents = await AgentModel.read_multiple_async(
-                db_session=session,
-                identifiers=agent_ids,
-                actor=actor,
-            )
-            return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents])
+            try:
+                query = select(AgentModel)
+                query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+                query = query.where(AgentModel.id.in_(agent_ids))
+                query = _apply_relationship_filters(query, include_relationships)
+
+                result = await session.execute(query)
+                agents = result.scalars().all()
+
+                if not agents:
+                    logger.warning(f"No agents found with IDs: {agent_ids}")
+                    return []
+
+                return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents])
+            except Exception as e:
+                logger.error(f"Error fetching agents with IDs {agent_ids}: {str(e)}")
+                raise
 
     @enforce_types
     @trace_method
@@ -1334,6 +1406,9 @@ class AgentManager:
             schema = MarshmallowAgentSchema(session=session, actor=actor)
             agent = schema.load(serialized_agent_dict, session=session)
 
+            agent.organization_id = actor.organization_id
+            for block in agent.core_memory:
+                block.organization_id = actor.organization_id
             if append_copy_suffix:
                 agent.name += "_copy"
             if project_id:
@@ -1435,10 +1510,17 @@ class AgentManager:
     @trace_method
     def list_groups(self, agent_id: str, actor: PydanticUser, manager_type: Optional[str] = None) -> List[PydanticGroup]:
         with db_registry.session() as session:
-            agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+            query = (
+                select(GroupModel)
+                .join(GroupsAgents, GroupModel.id == GroupsAgents.group_id)
+                .where(GroupsAgents.agent_id == agent_id, GroupModel.organization_id == actor.organization_id)
+            )
+
             if manager_type:
-                return [group.to_pydantic() for group in agent.groups if group.manager_type == manager_type]
-            return [group.to_pydantic() for group in agent.groups]
+                query = query.where(GroupModel.manager_type == manager_type)
+
+            result = session.execute(query)
+            return [group.to_pydantic() for group in result.scalars()]
 
     # ======================================================================================================================
     # In Context Messages Management
@@ -1559,15 +1641,9 @@ class AgentManager:
 
         Updates to the memory header should *not* trigger a rebuild, since that will simply flood recall storage with excess messages
         """
-        num_messages_task = self.message_manager.size_async(actor=actor, agent_id=agent_id)
-        num_archival_memories_task = self.passage_manager.agent_passage_size_async(actor=actor, agent_id=agent_id)
-        agent_state_task = self.get_agent_by_id_async(agent_id=agent_id, include_relationships=["memory", "sources", "tools"], actor=actor)
-
-        num_messages, num_archival_memories, agent_state = await asyncio.gather(
-            num_messages_task,
-            num_archival_memories_task,
-            agent_state_task,
-        )
+        num_messages = await self.message_manager.size_async(actor=actor, agent_id=agent_id)
+        num_archival_memories = await self.passage_manager.agent_passage_size_async(actor=actor, agent_id=agent_id)
+        agent_state = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=["memory", "sources", "tools"], actor=actor)
 
         if not tool_rules_solver:
             tool_rules_solver = ToolRulesSolver(agent_state.tool_rules)
@@ -1719,7 +1795,7 @@ class AgentManager:
                 logger.error(
                     f"Agent {agent_id} has no message_ids. Agent details: "
                     f"name={agent.name}, created_at={agent.created_at}, "
-                    f"message_ids={agent.message_ids}, organization_id={agent.organization_id}"
+                    f"message_ids={agent.message_ids}, organization_id={actor.organization_id}"
                 )
                 raise ValueError(f"Agent {agent_id} has no message_ids - cannot preserve system message")
 
@@ -1784,9 +1860,10 @@ class AgentManager:
                     )
 
             # refresh memory from DB (using block ids)
-            blocks = await asyncio.gather(
-                *[self.block_manager.get_block_by_id_async(block.id, actor=actor) for block in agent_state.memory.get_blocks()]
+            blocks = await self.block_manager.get_all_blocks_by_ids_async(
+                block_ids=[b.id for b in agent_state.memory.get_blocks()], actor=actor
             )
+
             agent_state.memory = Memory(
                 blocks=blocks,
                 file_blocks=agent_state.memory.file_blocks,
@@ -1907,7 +1984,7 @@ class AgentManager:
         """
         async with db_registry.async_session() as session:
             # Validate agent exists and user has access
-            await self._validate_agent_exists_async(session, agent_id, actor)
+            await validate_agent_exists_async(session, agent_id, actor)
 
             # Use raw SQL to efficiently fetch sources - much faster than lazy loading
             # Fast query without relationship loading
@@ -1943,7 +2020,7 @@ class AgentManager:
         """
         async with db_registry.async_session() as session:
             # Validate agent exists and user has access
-            await self._validate_agent_exists_async(session, agent_id, actor)
+            await validate_agent_exists_async(session, agent_id, actor)
 
             # Check if the source is actually attached to this agent using junction table
             attachment_check_query = select(SourcesAgents).where(SourcesAgents.agent_id == agent_id, SourcesAgents.source_id == source_id)
@@ -2710,7 +2787,7 @@ class AgentManager:
         """
         async with db_registry.async_session() as session:
             # lightweight check for agent access
-            await self._validate_agent_exists_async(session, agent_id, actor)
+            await validate_agent_exists_async(session, agent_id, actor)
 
             # direct query for tools via join - much more performant
             query = (
@@ -2752,7 +2829,12 @@ class AgentManager:
             )
 
             if query_text:
-                query = query.filter(AgentsTags.tag.ilike(f"%{query_text}%"))
+                if settings.database_engine is DatabaseChoice.POSTGRES:
+                    # PostgreSQL: Use ILIKE for case-insensitive search
+                    query = query.filter(AgentsTags.tag.ilike(f"%{query_text}%"))
+                else:
+                    # SQLite: Use LIKE with LOWER for case-insensitive search
+                    query = query.filter(func.lower(AgentsTags.tag).like(func.lower(f"%{query_text}%")))
 
             if after:
                 query = query.filter(AgentsTags.tag > after)
@@ -2788,7 +2870,12 @@ class AgentManager:
             )
 
             if query_text:
-                query = query.where(AgentsTags.tag.ilike(f"%{query_text}%"))
+                if settings.database_engine is DatabaseChoice.POSTGRES:
+                    # PostgreSQL: Use ILIKE for case-insensitive search
+                    query = query.where(AgentsTags.tag.ilike(f"%{query_text}%"))
+                else:
+                    # SQLite: Use LIKE with LOWER for case-insensitive search
+                    query = query.where(func.lower(AgentsTags.tag).like(func.lower(f"%{query_text}%")))
 
             if after:
                 query = query.where(AgentsTags.tag > after)

letta/services/block_manager.py

@@ -1,12 +1,14 @@
 import asyncio
 from typing import Dict, List, Optional
 
-from sqlalchemy import select
+from sqlalchemy import delete, select
 from sqlalchemy.orm import Session
 
 from letta.log import get_logger
+from letta.orm.agent import Agent as AgentModel
 from letta.orm.block import Block as BlockModel
 from letta.orm.block_history import BlockHistory
+from letta.orm.blocks_agents import BlocksAgents
 from letta.orm.errors import NoResultFound
 from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState as PydanticAgentState
@@ -50,8 +52,10 @@ class BlockManager:
             async with db_registry.async_session() as session:
                 data = block.model_dump(to_orm=True, exclude_none=True)
                 block = BlockModel(**data, organization_id=actor.organization_id)
-                await block.create_async(session, actor=actor)
-                return block.to_pydantic()
+                await block.create_async(session, actor=actor, no_commit=True, no_refresh=True)
+                pydantic_block = block.to_pydantic()
+                await session.commit()
+                return pydantic_block
 
     @enforce_types
     @trace_method
@@ -95,11 +99,12 @@ class BlockManager:
             block_models = [
                 BlockModel(**block.model_dump(to_orm=True, exclude_none=True), organization_id=actor.organization_id) for block in blocks
             ]
-
-            created_models = await BlockModel.batch_create_async(items=block_models, db_session=session, actor=actor)
-
-            # Convert back to Pydantic
-            return [m.to_pydantic() for m in created_models]
+            created_models = await BlockModel.batch_create_async(
+                items=block_models, db_session=session, actor=actor, no_commit=True, no_refresh=True
+            )
+            result = [m.to_pydantic() for m in created_models]
+            await session.commit()
+            return result
 
     @enforce_types
     @trace_method
@@ -130,26 +135,36 @@ class BlockManager:
             for key, value in update_data.items():
                 setattr(block, key, value)
 
-            await block.update_async(db_session=session, actor=actor)
-            return block.to_pydantic()
+            await block.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True)
+            pydantic_block = block.to_pydantic()
+            await session.commit()
+            return pydantic_block
 
     @enforce_types
     @trace_method
-    def delete_block(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
+    def delete_block(self, block_id: str, actor: PydanticUser) -> None:
         """Delete a block by its ID."""
         with db_registry.session() as session:
+            # First, delete all references in blocks_agents table
+            session.execute(delete(BlocksAgents).where(BlocksAgents.block_id == block_id))
+            session.flush()
+
+            # Then delete the block itself
             block = BlockModel.read(db_session=session, identifier=block_id)
             block.hard_delete(db_session=session, actor=actor)
-            return block.to_pydantic()
 
     @enforce_types
     @trace_method
-    async def delete_block_async(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
+    async def delete_block_async(self, block_id: str, actor: PydanticUser) -> None:
         """Delete a block by its ID."""
         async with db_registry.async_session() as session:
+            # First, delete all references in blocks_agents table
+            await session.execute(delete(BlocksAgents).where(BlocksAgents.block_id == block_id))
+            await session.flush()
+
+            # Then delete the block itself
             block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
             await block.hard_delete_async(db_session=session, actor=actor)
-            return block.to_pydantic()
 
     @enforce_types
     @trace_method
@@ -296,8 +311,14 @@ class BlockManager:
         Retrieve all agents associated with a given block.
         """
         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
+            query = (
+                select(AgentModel)
+                .where(AgentModel.id.in_(select(BlocksAgents.agent_id).where(BlocksAgents.block_id == block_id)))
+                .where(AgentModel.organization_id == actor.organization_id)
+            )
+
+            result = await session.execute(query)
+            agents_orm = result.scalars().all()
             agents = await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents_orm])
             return agents
 
@@ -531,7 +552,7 @@ class BlockManager:
             for block in blocks:
                 new_val = updates[block.id]
                 if len(new_val) > block.limit:
-                    logger.warning(f"Value length ({len(new_val)}) exceeds limit " f"({block.limit}) for block {block.id!r}, truncating...")
+                    logger.warning(f"Value length ({len(new_val)}) exceeds limit ({block.limit}) for block {block.id!r}, truncating...")
                     new_val = new_val[: block.limit]
                 block.value = new_val
 

letta/services/file_manager.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 from datetime import datetime
 from typing import List, Optional
@@ -357,7 +358,9 @@ class FileManager:
 
     @enforce_types
     @trace_method
-    async def get_organization_sources_metadata(self, actor: PydanticUser) -> OrganizationSourcesStats:
+    async def get_organization_sources_metadata(
+        self, actor: PydanticUser, include_detailed_per_source_metadata: bool = False
+    ) -> OrganizationSourcesStats:
         """
         Get aggregated metadata for all sources in an organization with optimized queries.
 
@@ -365,7 +368,7 @@ class FileManager:
         - Total number of sources
         - Total number of files across all sources
         - Total size of all files
-        - Per-source breakdown with file details
+        - Per-source breakdown with file details (if include_detailed_per_source_metadata is True)
         """
         async with db_registry.async_session() as session:
             # Import here to avoid circular imports
@@ -395,31 +398,113 @@ class FileManager:
             for row in source_aggregations:
                 source_id, source_name, file_count, total_size = row
 
-                # Get individual file details for this source
-                files_query = (
-                    select(FileMetadataModel.id, FileMetadataModel.file_name, FileMetadataModel.file_size)
-                    .where(
-                        FileMetadataModel.source_id == source_id,
-                        FileMetadataModel.organization_id == actor.organization_id,
-                        FileMetadataModel.is_deleted == False,
+                if include_detailed_per_source_metadata:
+                    # Get individual file details for this source
+                    files_query = (
+                        select(FileMetadataModel.id, FileMetadataModel.file_name, FileMetadataModel.file_size)
+                        .where(
+                            FileMetadataModel.source_id == source_id,
+                            FileMetadataModel.organization_id == actor.organization_id,
+                            FileMetadataModel.is_deleted == False,
+                        )
+                        .order_by(FileMetadataModel.file_name)
                     )
-                    .order_by(FileMetadataModel.file_name)
-                )
 
-                files_result = await session.execute(files_query)
-                files_rows = files_result.fetchall()
+                    files_result = await session.execute(files_query)
+                    files_rows = files_result.fetchall()
 
-                # Build file stats
-                files = [FileStats(file_id=file_row[0], file_name=file_row[1], file_size=file_row[2]) for file_row in files_rows]
+                    # Build file stats
+                    files = [FileStats(file_id=file_row[0], file_name=file_row[1], file_size=file_row[2]) for file_row in files_rows]
 
-                # Build source metadata
-                source_metadata = SourceStats(
-                    source_id=source_id, source_name=source_name, file_count=file_count, total_size=total_size, files=files
-                )
+                    # Build source metadata
+                    source_metadata = SourceStats(
+                        source_id=source_id, source_name=source_name, file_count=file_count, total_size=total_size, files=files
+                    )
+
+                    metadata.sources.append(source_metadata)
 
-                metadata.sources.append(source_metadata)
                 metadata.total_files += file_count
                 metadata.total_size += total_size
 
-            metadata.total_sources = len(metadata.sources)
+            metadata.total_sources = len(source_aggregations)
             return metadata
+
+    @enforce_types
+    @trace_method
+    async def get_files_by_ids_async(
+        self, file_ids: List[str], actor: PydanticUser, *, include_content: bool = False
+    ) -> List[PydanticFileMetadata]:
+        """
+        Get multiple files by their IDs in a single query.
+
+        Args:
+            file_ids: List of file IDs to retrieve
+            actor: User performing the action
+            include_content: Whether to include file content in the response
+
+        Returns:
+            List[PydanticFileMetadata]: List of files (may be fewer than requested if some don't exist)
+        """
+        if not file_ids:
+            return []
+
+        async with db_registry.async_session() as session:
+            query = select(FileMetadataModel).where(
+                FileMetadataModel.id.in_(file_ids),
+                FileMetadataModel.organization_id == actor.organization_id,
+                FileMetadataModel.is_deleted == False,
+            )
+
+            # Eagerly load content if requested
+            if include_content:
+                query = query.options(selectinload(FileMetadataModel.content))
+
+            result = await session.execute(query)
+            files_orm = result.scalars().all()
+
+            return await asyncio.gather(*[file.to_pydantic_async(include_content=include_content) for file in files_orm])
+
+    @enforce_types
+    @trace_method
+    async def get_files_for_agents_async(
+        self, agent_ids: List[str], actor: PydanticUser, *, include_content: bool = False
+    ) -> List[PydanticFileMetadata]:
+        """
+        Get all files associated with the given agents via file-agent relationships.
+
+        Args:
+            agent_ids: List of agent IDs to find files for
+            actor: User performing the action
+            include_content: Whether to include file content in the response
+
+        Returns:
+            List[PydanticFileMetadata]: List of unique files associated with these agents
+        """
+        if not agent_ids:
+            return []
+
+        async with db_registry.async_session() as session:
+            # We need to import FileAgent here to avoid circular imports
+            from letta.orm.file_agent import FileAgent as FileAgentModel
+
+            # Join through file-agent relationships
+            query = (
+                select(FileMetadataModel)
+                .join(FileAgentModel, FileMetadataModel.id == FileAgentModel.file_id)
+                .where(
+                    FileAgentModel.agent_id.in_(agent_ids),
+                    FileMetadataModel.organization_id == actor.organization_id,
+                    FileMetadataModel.is_deleted == False,
+                    FileAgentModel.is_deleted == False,
+                )
+                .distinct()  # Ensure we don't get duplicate files
+            )
+
+            # Eagerly load content if requested
+            if include_content:
+                query = query.options(selectinload(FileMetadataModel.content))
+
+            result = await session.execute(query)
+            files_orm = result.scalars().all()
+
+            return await asyncio.gather(*[file.to_pydantic_async(include_content=include_content) for file in files_orm])