letta-nightly 0.10.0.dev20250806104523__py3-none-any.whl → 0.11.0.dev20250807104511__py3-none-any.whl

letta/services/agent_serialization_manager.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Dict, List
+from typing import Any, Dict, List, Optional
 
 from letta.constants import MCP_TOOL_TAG_NAME_PREFIX
 from letta.errors import AgentFileExportError, AgentFileImportError
@@ -22,6 +22,7 @@ from letta.schemas.agent_file import (
 from letta.schemas.block import Block
 from letta.schemas.enums import FileProcessingStatus
 from letta.schemas.file import FileMetadata
+from letta.schemas.group import Group, GroupCreate
 from letta.schemas.mcp import MCPServer
 from letta.schemas.message import Message
 from letta.schemas.source import Source
@@ -230,6 +231,9 @@ class AgentSerializationManager:
                 file_agent.source_id = self._map_db_to_file_id(file_agent.source_id, SourceSchema.__id_prefix__)
                 file_agent.agent_id = agent_file_id
 
+        if agent_schema.group_ids:
+            agent_schema.group_ids = [self._map_db_to_file_id(group_id, GroupSchema.__id_prefix__) for group_id in agent_schema.group_ids]
+
         return agent_schema
 
     def _convert_tool_to_schema(self, tool) -> ToolSchema:
@@ -308,6 +312,24 @@ class AgentSerializationManager:
             logger.error(f"Failed to convert MCP server {mcp_server.id}: {e}")
             raise
 
+    def _convert_group_to_schema(self, group: Group) -> GroupSchema:
+        """Convert Group to GroupSchema with ID remapping"""
+        try:
+            group_file_id = self._map_db_to_file_id(group.id, GroupSchema.__id_prefix__, allow_new=False)
+            group_schema = GroupSchema.from_group(group)
+            group_schema.id = group_file_id
+            group_schema.agent_ids = [
+                self._map_db_to_file_id(agent_id, AgentSchema.__id_prefix__, allow_new=False) for agent_id in group_schema.agent_ids
+            ]
+            if hasattr(group_schema.manager_config, "manager_agent_id"):
+                group_schema.manager_config.manager_agent_id = self._map_db_to_file_id(
+                    group_schema.manager_config.manager_agent_id, AgentSchema.__id_prefix__, allow_new=False
+                )
+            return group_schema
+        except Exception as e:
+            logger.error(f"Failed to convert group {group.id}: {e}")
+            raise
+
     async def export(self, agent_ids: List[str], actor: User) -> AgentFileSchema:
         """
         Export agents and their related entities to AgentFileSchema format.
@@ -332,6 +354,23 @@ class AgentSerializationManager:
                 missing_ids = [agent_id for agent_id in agent_ids if agent_id not in found_ids]
                 raise AgentFileExportError(f"The following agent IDs were not found: {missing_ids}")
 
+            groups = []
+            group_agent_ids = []
+            for agent_state in agent_states:
+                if agent_state.multi_agent_group != None:
+                    groups.append(agent_state.multi_agent_group)
+                    group_agent_ids.extend(agent_state.multi_agent_group.agent_ids)
+
+            group_agent_ids = list(set(group_agent_ids) - set(agent_ids))
+            if group_agent_ids:
+                group_agent_states = await self.agent_manager.get_agents_by_ids_async(agent_ids=group_agent_ids, actor=actor)
+                if len(group_agent_states) != len(group_agent_ids):
+                    found_ids = {agent.id for agent in group_agent_states}
+                    missing_ids = [agent_id for agent_id in group_agent_ids if agent_id not in found_ids]
+                    raise AgentFileExportError(f"The following agent IDs were not found: {missing_ids}")
+                agent_ids.extend(group_agent_ids)
+                agent_states.extend(group_agent_states)
+
             # cache for file-agent relationships to avoid duplicate queries
             files_agents_cache = {}  # Maps agent_id to list of file_agent relationships
 
@@ -359,13 +398,14 @@ class AgentSerializationManager:
             source_schemas = [self._convert_source_to_schema(source) for source in source_set]
             file_schemas = [self._convert_file_to_schema(file_metadata) for file_metadata in file_set]
             mcp_server_schemas = [self._convert_mcp_server_to_schema(mcp_server) for mcp_server in mcp_server_set]
+            group_schemas = [self._convert_group_to_schema(group) for group in groups]
 
             logger.info(f"Exporting {len(agent_ids)} agents to agent file format")
 
             # Return AgentFileSchema with converted entities
             return AgentFileSchema(
                 agents=agent_schemas,
-                groups=[],  # TODO: Extract and convert groups
+                groups=group_schemas,
                 blocks=block_schemas,
                 files=file_schemas,
                 sources=source_schemas,
@@ -379,7 +419,13 @@ class AgentSerializationManager:
             logger.error(f"Failed to export agent file: {e}")
             raise AgentFileExportError(f"Export failed: {e}") from e
 
-    async def import_file(self, schema: AgentFileSchema, actor: User, dry_run: bool = False) -> ImportResult:
+    async def import_file(
+        self,
+        schema: AgentFileSchema,
+        actor: User,
+        dry_run: bool = False,
+        env_vars: Optional[Dict[str, Any]] = None,
+    ) -> ImportResult:
         """
         Import AgentFileSchema into the database.
 
@@ -546,6 +592,10 @@ class AgentSerializationManager:
                 if agent_data.get("block_ids"):
                     agent_data["block_ids"] = [file_to_db_ids[file_id] for file_id in agent_data["block_ids"]]
 
+                if env_vars:
+                    for var in agent_data["tool_exec_environment_variables"]:
+                        var["value"] = env_vars.get(var["key"], "")
+
                 agent_create = CreateAgent(**agent_data)
                 created_agent = await self.agent_manager.create_agent_async(agent_create, actor, _init_with_no_messages=True)
                 file_to_db_ids[agent_schema.id] = created_agent.id
@@ -607,6 +657,15 @@ class AgentSerializationManager:
                     )
                     imported_count += len(files_for_agent)
 
+            for group in schema.groups:
+                group_data = group.model_dump(exclude={"id"})
+                group_data["agent_ids"] = [file_to_db_ids[agent_id] for agent_id in group_data["agent_ids"]]
+                if "manager_agent_id" in group_data["manager_config"]:
+                    group_data["manager_config"]["manager_agent_id"] = file_to_db_ids[group_data["manager_config"]["manager_agent_id"]]
+                created_group = await self.group_manager.create_group_async(GroupCreate(**group_data), actor)
+                file_to_db_ids[group.id] = created_group.id
+                imported_count += 1
+
             return ImportResult(
                 success=True,
                 message=f"Import completed successfully. Imported {imported_count} entities.",

letta/services/archive_manager.py

@@ -0,0 +1,269 @@
+from typing import List, Optional
+
+from sqlalchemy import select
+
+from letta.log import get_logger
+from letta.orm import ArchivalPassage
+from letta.orm import Archive as ArchiveModel
+from letta.orm import ArchivesAgents
+from letta.schemas.archive import Archive as PydanticArchive
+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__)
+
+
+class ArchiveManager:
+    """Manager class to handle business logic related to Archives."""
+
+    @enforce_types
+    def create_archive(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        actor: PydanticUser = None,
+    ) -> PydanticArchive:
+        """Create a new archive."""
+        try:
+            with db_registry.session() as session:
+                archive = ArchiveModel(
+                    name=name,
+                    description=description,
+                    organization_id=actor.organization_id,
+                )
+                archive.create(session, actor=actor)
+                return archive.to_pydantic()
+        except Exception as e:
+            logger.exception(f"Failed to create archive {name}. error={e}")
+            raise
+
+    @enforce_types
+    async def create_archive_async(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        actor: PydanticUser = None,
+    ) -> PydanticArchive:
+        """Create a new archive."""
+        try:
+            async with db_registry.async_session() as session:
+                archive = ArchiveModel(
+                    name=name,
+                    description=description,
+                    organization_id=actor.organization_id,
+                )
+                await archive.create_async(session, actor=actor)
+                return archive.to_pydantic()
+        except Exception as e:
+            logger.exception(f"Failed to create archive {name}. error={e}")
+            raise
+
+    @enforce_types
+    async def get_archive_by_id_async(
+        self,
+        archive_id: str,
+        actor: PydanticUser,
+    ) -> PydanticArchive:
+        """Get an archive by ID."""
+        async with db_registry.async_session() as session:
+            archive = await ArchiveModel.read_async(
+                db_session=session,
+                identifier=archive_id,
+                actor=actor,
+            )
+            return archive.to_pydantic()
+
+    @enforce_types
+    def attach_agent_to_archive(
+        self,
+        agent_id: str,
+        archive_id: str,
+        is_owner: bool,
+        actor: PydanticUser,
+    ) -> None:
+        """Attach an agent to an archive."""
+        with db_registry.session() as session:
+            # Check if already attached
+            existing = session.query(ArchivesAgents).filter_by(agent_id=agent_id, archive_id=archive_id).first()
+
+            if existing:
+                # Update ownership if needed
+                if existing.is_owner != is_owner:
+                    existing.is_owner = is_owner
+                    session.commit()
+                return
+
+            # Create new relationship
+            archives_agents = ArchivesAgents(
+                agent_id=agent_id,
+                archive_id=archive_id,
+                is_owner=is_owner,
+            )
+            session.add(archives_agents)
+            session.commit()
+
+    @enforce_types
+    async def attach_agent_to_archive_async(
+        self,
+        agent_id: str,
+        archive_id: str,
+        is_owner: bool = False,
+        actor: PydanticUser = None,
+    ) -> None:
+        """Attach an agent to an archive."""
+        async with db_registry.async_session() as session:
+            # Check if relationship already exists
+            existing = await session.execute(
+                select(ArchivesAgents).where(
+                    ArchivesAgents.agent_id == agent_id,
+                    ArchivesAgents.archive_id == archive_id,
+                )
+            )
+            existing_record = existing.scalar_one_or_none()
+
+            if existing_record:
+                # Update ownership if needed
+                if existing_record.is_owner != is_owner:
+                    existing_record.is_owner = is_owner
+                    await session.commit()
+                return
+
+            # Create the relationship
+            archives_agents = ArchivesAgents(
+                agent_id=agent_id,
+                archive_id=archive_id,
+                is_owner=is_owner,
+            )
+            session.add(archives_agents)
+            await session.commit()
+
+    @enforce_types
+    async def get_or_create_default_archive_for_agent_async(
+        self,
+        agent_id: str,
+        agent_name: Optional[str] = None,
+        actor: PydanticUser = None,
+    ) -> PydanticArchive:
+        """Get the agent's default archive, creating one if it doesn't exist."""
+        # First check if agent has any archives
+        from letta.services.agent_manager import AgentManager
+
+        agent_manager = AgentManager()
+
+        archive_ids = await agent_manager.get_agent_archive_ids_async(
+            agent_id=agent_id,
+            actor=actor,
+        )
+
+        if archive_ids:
+            # TODO: Remove this check once we support multiple archives per agent
+            if len(archive_ids) > 1:
+                raise ValueError(f"Agent {agent_id} has multiple archives, which is not yet supported")
+            # Get the archive
+            archive = await self.get_archive_by_id_async(
+                archive_id=archive_ids[0],
+                actor=actor,
+            )
+            return archive
+
+        # Create a default archive for this agent
+        archive_name = f"{agent_name or f'Agent {agent_id}'}'s Archive"
+        archive = await self.create_archive_async(
+            name=archive_name,
+            description="Default archive created automatically",
+            actor=actor,
+        )
+
+        # Attach the agent to the archive as owner
+        await self.attach_agent_to_archive_async(
+            agent_id=agent_id,
+            archive_id=archive.id,
+            is_owner=True,
+            actor=actor,
+        )
+
+        return archive
+
+    @enforce_types
+    def get_or_create_default_archive_for_agent(
+        self,
+        agent_id: str,
+        agent_name: Optional[str] = None,
+        actor: PydanticUser = None,
+    ) -> PydanticArchive:
+        """Get the agent's default archive, creating one if it doesn't exist."""
+        with db_registry.session() as session:
+            # First check if agent has any archives
+            query = select(ArchivesAgents.archive_id).where(ArchivesAgents.agent_id == agent_id)
+            result = session.execute(query)
+            archive_ids = [row[0] for row in result.fetchall()]
+
+            if archive_ids:
+                # TODO: Remove this check once we support multiple archives per agent
+                if len(archive_ids) > 1:
+                    raise ValueError(f"Agent {agent_id} has multiple archives, which is not yet supported")
+                # Get the archive
+                archive = ArchiveModel.read(db_session=session, identifier=archive_ids[0], actor=actor)
+                return archive.to_pydantic()
+
+            # Create a default archive for this agent
+            archive_name = f"{agent_name or f'Agent {agent_id}'}'s Archive"
+
+            # Create the archive
+            archive_model = ArchiveModel(
+                name=archive_name,
+                description="Default archive created automatically",
+                organization_id=actor.organization_id,
+            )
+            archive_model.create(session, actor=actor)
+
+        # Attach the agent to the archive as owner
+        self.attach_agent_to_archive(
+            agent_id=agent_id,
+            archive_id=archive_model.id,
+            is_owner=True,
+            actor=actor,
+        )
+
+        return archive_model.to_pydantic()
+
+    @enforce_types
+    async def get_agents_for_archive_async(
+        self,
+        archive_id: str,
+        actor: PydanticUser,
+    ) -> List[str]:
+        """Get all agent IDs that have access to an archive."""
+        async with db_registry.async_session() as session:
+            result = await session.execute(select(ArchivesAgents.agent_id).where(ArchivesAgents.archive_id == archive_id))
+            return [row[0] for row in result.fetchall()]
+
+    @enforce_types
+    async def get_agent_from_passage_async(
+        self,
+        passage_id: str,
+        actor: PydanticUser,
+    ) -> Optional[str]:
+        """Get the agent ID that owns a passage (through its archive).
+
+        Returns the first agent found (for backwards compatibility).
+        Returns None if no agent found.
+        """
+        async with db_registry.async_session() as session:
+            # First get the passage to find its archive_id
+            passage = await ArchivalPassage.read_async(
+                db_session=session,
+                identifier=passage_id,
+                actor=actor,
+            )
+
+            # Then find agents connected to that archive
+            result = await session.execute(select(ArchivesAgents.agent_id).where(ArchivesAgents.archive_id == passage.archive_id))
+            agent_ids = [row[0] for row in result.fetchall()]
+
+            if not agent_ids:
+                return None
+
+            # For now, return the first agent (backwards compatibility)
+            return agent_ids[0]

letta/services/helpers/agent_manager_helper.py

@@ -25,9 +25,10 @@ from letta.helpers import ToolRulesSolver
 from letta.helpers.datetime_helpers import format_datetime, get_local_time, get_local_time_fast
 from letta.orm.agent import Agent as AgentModel
 from letta.orm.agents_tags import AgentsTags
+from letta.orm.archives_agents import ArchivesAgents
 from letta.orm.errors import NoResultFound
 from letta.orm.identity import Identity
-from letta.orm.passage import AgentPassage, SourcePassage
+from letta.orm.passage import ArchivalPassage, SourcePassage
 from letta.orm.sources_agents import SourcesAgents
 from letta.orm.sqlite_functions import adapt_array
 from letta.otel.tracing import trace_method
@@ -328,6 +329,74 @@ def compile_system_message(
     return formatted_prompt
 
 
+@trace_method
+def get_system_message_from_compiled_memory(
+    system_prompt: str,
+    memory_with_sources: str,
+    in_context_memory_last_edit: datetime,  # TODO move this inside of BaseMemory?
+    timezone: str,
+    user_defined_variables: Optional[dict] = None,
+    append_icm_if_missing: bool = True,
+    template_format: Literal["f-string", "mustache", "jinja2"] = "f-string",
+    previous_message_count: int = 0,
+    archival_memory_size: int = 0,
+) -> str:
+    """Prepare the final/full system message that will be fed into the LLM API
+
+    The base system message may be templated, in which case we need to render the variables.
+
+    The following are reserved variables:
+      - CORE_MEMORY: the in-context memory of the LLM
+    """
+    if user_defined_variables is not None:
+        # TODO eventually support the user defining their own variables to inject
+        raise NotImplementedError
+    else:
+        variables = {}
+
+    # Add the protected memory variable
+    if IN_CONTEXT_MEMORY_KEYWORD in variables:
+        raise ValueError(f"Found protected variable '{IN_CONTEXT_MEMORY_KEYWORD}' in user-defined vars: {str(user_defined_variables)}")
+    else:
+        # TODO should this all put into the memory.__repr__ function?
+        memory_metadata_string = compile_memory_metadata_block(
+            memory_edit_timestamp=in_context_memory_last_edit,
+            previous_message_count=previous_message_count,
+            archival_memory_size=archival_memory_size,
+            timezone=timezone,
+        )
+
+        full_memory_string = memory_with_sources + "\n\n" + memory_metadata_string
+
+        # Add to the variables list to inject
+        variables[IN_CONTEXT_MEMORY_KEYWORD] = full_memory_string
+
+    if template_format == "f-string":
+        memory_variable_string = "{" + IN_CONTEXT_MEMORY_KEYWORD + "}"
+
+        # Catch the special case where the system prompt is unformatted
+        if append_icm_if_missing:
+            if memory_variable_string not in system_prompt:
+                # In this case, append it to the end to make sure memory is still injected
+                # warnings.warn(f"{IN_CONTEXT_MEMORY_KEYWORD} variable was missing from system prompt, appending instead")
+                system_prompt += "\n\n" + memory_variable_string
+
+        # render the variables using the built-in templater
+        try:
+            if user_defined_variables:
+                formatted_prompt = safe_format(system_prompt, variables)
+            else:
+                formatted_prompt = system_prompt.replace(memory_variable_string, full_memory_string)
+        except Exception as e:
+            raise ValueError(f"Failed to format system prompt - {str(e)}. System prompt value:\n{system_prompt}")
+
+    else:
+        # TODO support for mustache and jinja2
+        raise NotImplementedError(template_format)
+
+    return formatted_prompt
+
+
 @trace_method
 async def compile_system_message_async(
     system_prompt: str,
@@ -374,7 +443,7 @@ async def compile_system_message_async(
             timezone=timezone,
         )
 
-        memory_with_sources = await in_context_memory.compile_async(
+        memory_with_sources = await in_context_memory.compile_in_thread_async(
             tool_usage_rules=tool_constraint_block, sources=sources, max_files_open=max_files_open
         )
         full_memory_string = memory_with_sources + "\n\n" + memory_metadata_string
@@ -918,7 +987,7 @@ def build_passage_query(
                     SourcePassage.organization_id,
                     SourcePassage.file_id,
                     SourcePassage.source_id,
-                    literal(None).label("agent_id"),
+                    literal(None).label("archive_id"),
                 )
                 .join(SourcesAgents, SourcesAgents.source_id == SourcePassage.source_id)
                 .where(SourcesAgents.agent_id == agent_id)
@@ -940,7 +1009,7 @@ def build_passage_query(
                 SourcePassage.organization_id,
                 SourcePassage.file_id,
                 SourcePassage.source_id,
-                literal(None).label("agent_id"),
+                literal(None).label("archive_id"),
             ).where(SourcePassage.organization_id == actor.organization_id)
 
         if source_id:
@@ -954,23 +1023,24 @@ def build_passage_query(
         agent_passages = (
             select(
                 literal(None).label("file_name"),
-                AgentPassage.id,
-                AgentPassage.text,
-                AgentPassage.embedding_config,
-                AgentPassage.metadata_,
-                AgentPassage.embedding,
-                AgentPassage.created_at,
-                AgentPassage.updated_at,
-                AgentPassage.is_deleted,
-                AgentPassage._created_by_id,
-                AgentPassage._last_updated_by_id,
-                AgentPassage.organization_id,
+                ArchivalPassage.id,
+                ArchivalPassage.text,
+                ArchivalPassage.embedding_config,
+                ArchivalPassage.metadata_,
+                ArchivalPassage.embedding,
+                ArchivalPassage.created_at,
+                ArchivalPassage.updated_at,
+                ArchivalPassage.is_deleted,
+                ArchivalPassage._created_by_id,
+                ArchivalPassage._last_updated_by_id,
+                ArchivalPassage.organization_id,
                 literal(None).label("file_id"),
                 literal(None).label("source_id"),
-                AgentPassage.agent_id,
+                ArchivalPassage.archive_id,
             )
-            .where(AgentPassage.agent_id == agent_id)
-            .where(AgentPassage.organization_id == actor.organization_id)
+            .join(ArchivesAgents, ArchivalPassage.archive_id == ArchivesAgents.archive_id)
+            .where(ArchivesAgents.agent_id == agent_id)
+            .where(ArchivalPassage.organization_id == actor.organization_id)
         )
 
     # Combine queries
@@ -1201,56 +1271,60 @@ def build_agent_passage_query(
         embedded_text = np.array(embedded_text)
         embedded_text = np.pad(embedded_text, (0, MAX_EMBEDDING_DIM - embedded_text.shape[0]), mode="constant").tolist()
 
-    # Base query for agent passages
-    query = select(AgentPassage).where(AgentPassage.agent_id == agent_id, AgentPassage.organization_id == actor.organization_id)
+    # Base query for agent passages - join through archives_agents
+    query = (
+        select(ArchivalPassage)
+        .join(ArchivesAgents, ArchivalPassage.archive_id == ArchivesAgents.archive_id)
+        .where(ArchivesAgents.agent_id == agent_id, ArchivalPassage.organization_id == actor.organization_id)
+    )
 
     # Apply filters
     if start_date:
-        query = query.where(AgentPassage.created_at >= start_date)
+        query = query.where(ArchivalPassage.created_at >= start_date)
     if end_date:
-        query = query.where(AgentPassage.created_at <= end_date)
+        query = query.where(ArchivalPassage.created_at <= end_date)
 
     # Handle text search or vector search
     if embedded_text:
         if settings.database_engine is DatabaseChoice.POSTGRES:
             # PostgreSQL with pgvector
-            query = query.order_by(AgentPassage.embedding.cosine_distance(embedded_text).asc())
+            query = query.order_by(ArchivalPassage.embedding.cosine_distance(embedded_text).asc())
         else:
             # SQLite with custom vector type
             query_embedding_binary = adapt_array(embedded_text)
             query = query.order_by(
-                func.cosine_distance(AgentPassage.embedding, query_embedding_binary).asc(),
-                AgentPassage.created_at.asc() if ascending else AgentPassage.created_at.desc(),
-                AgentPassage.id.asc(),
+                func.cosine_distance(ArchivalPassage.embedding, query_embedding_binary).asc(),
+                ArchivalPassage.created_at.asc() if ascending else ArchivalPassage.created_at.desc(),
+                ArchivalPassage.id.asc(),
             )
     else:
         if query_text:
-            query = query.where(func.lower(AgentPassage.text).contains(func.lower(query_text)))
+            query = query.where(func.lower(ArchivalPassage.text).contains(func.lower(query_text)))
 
     # Handle pagination
     if before or after:
         if before:
             # Get the reference record
-            before_subq = select(AgentPassage.created_at, AgentPassage.id).where(AgentPassage.id == before).subquery()
+            before_subq = select(ArchivalPassage.created_at, ArchivalPassage.id).where(ArchivalPassage.id == before).subquery()
             query = query.where(
                 or_(
-                    AgentPassage.created_at < before_subq.c.created_at,
+                    ArchivalPassage.created_at < before_subq.c.created_at,
                     and_(
-                        AgentPassage.created_at == before_subq.c.created_at,
-                        AgentPassage.id < before_subq.c.id,
+                        ArchivalPassage.created_at == before_subq.c.created_at,
+                        ArchivalPassage.id < before_subq.c.id,
                     ),
                 )
             )
 
         if after:
             # Get the reference record
-            after_subq = select(AgentPassage.created_at, AgentPassage.id).where(AgentPassage.id == after).subquery()
+            after_subq = select(ArchivalPassage.created_at, ArchivalPassage.id).where(ArchivalPassage.id == after).subquery()
             query = query.where(
                 or_(
-                    AgentPassage.created_at > after_subq.c.created_at,
+                    ArchivalPassage.created_at > after_subq.c.created_at,
                     and_(
-                        AgentPassage.created_at == after_subq.c.created_at,
-                        AgentPassage.id > after_subq.c.id,
+                        ArchivalPassage.created_at == after_subq.c.created_at,
+                        ArchivalPassage.id > after_subq.c.id,
                     ),
                 )
             )
@@ -1258,9 +1332,9 @@ def build_agent_passage_query(
     # Apply ordering if not already ordered by similarity
     if not embed_query:
         if ascending:
-            query = query.order_by(AgentPassage.created_at.asc(), AgentPassage.id.asc())
+            query = query.order_by(ArchivalPassage.created_at.asc(), ArchivalPassage.id.asc())
         else:
-            query = query.order_by(AgentPassage.created_at.desc(), AgentPassage.id.asc())
+            query = query.order_by(ArchivalPassage.created_at.desc(), ArchivalPassage.id.asc())
 
     return query
 

letta/services/job_manager.py

@@ -806,6 +806,30 @@ class JobManager:
             request_config = job.request_config or LettaRequestConfig()
         return request_config
 
+    @enforce_types
+    async def record_ttft(self, job_id: str, ttft_ns: int, actor: PydanticUser) -> None:
+        """Record time to first token for a run"""
+        try:
+            async with db_registry.async_session() as session:
+                job = await self._verify_job_access_async(session=session, job_id=job_id, actor=actor, access=["write"])
+                job.ttft_ns = ttft_ns
+                await job.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True)
+                await session.commit()
+        except Exception as e:
+            logger.warning(f"Failed to record TTFT for job {job_id}: {e}")
+
+    @enforce_types
+    async def record_response_duration(self, job_id: str, total_duration_ns: int, actor: PydanticUser) -> None:
+        """Record total response duration for a run"""
+        try:
+            async with db_registry.async_session() as session:
+                job = await self._verify_job_access_async(session=session, job_id=job_id, actor=actor, access=["write"])
+                job.total_duration_ns = total_duration_ns
+                await job.update_async(db_session=session, actor=actor, no_commit=True, no_refresh=True)
+                await session.commit()
+        except Exception as e:
+            logger.warning(f"Failed to record response duration for job {job_id}: {e}")
+
     @trace_method
     def _dispatch_callback_sync(self, callback_info: dict) -> dict:
         """