letta-nightly 0.8.0.dev20250606104326__py3-none-any.whl → 0.8.2.dev20250606215616__py3-none-any.whl

letta/services/file_processor/parser/mistral_parser.py

@@ -9,6 +9,16 @@ from letta.settings import settings
 logger = get_logger(__name__)
 
 
+SIMPLE_TEXT_MIME_TYPES = {
+    "text/plain",
+    "text/markdown",
+    "text/x-markdown",
+    "application/json",
+    "application/jsonl",
+    "application/x-jsonlines",
+}
+
+
 class MistralFileParser(FileParser):
     """Mistral-based OCR extraction"""
 
@@ -23,7 +33,7 @@ class MistralFileParser(FileParser):
 
             # TODO: Kind of hacky...we try to exit early here?
             # TODO: Create our internal file parser representation we return instead of OCRResponse
-            if mime_type == "text/plain":
+            if mime_type in SIMPLE_TEXT_MIME_TYPES or mime_type.startswith("text/"):
                 text = content.decode("utf-8", errors="replace")
                 return OCRResponse(
                     model=self.model,

letta/services/files_agents_manager.py

@@ -5,10 +5,11 @@ from sqlalchemy import and_, func, select, update
 
 from letta.orm.errors import NoResultFound
 from letta.orm.files_agents import FileAgent as FileAgentModel
+from letta.otel.tracing import trace_method
+from letta.schemas.block import Block as PydanticBlock
 from letta.schemas.file import FileAgent as PydanticFileAgent
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
-from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 
@@ -22,6 +23,7 @@ class FileAgentManager:
         *,
         agent_id: str,
         file_id: str,
+        file_name: str,
         actor: PydanticUser,
         is_open: bool = True,
         visible_content: Optional[str] = None,
@@ -38,6 +40,7 @@ class FileAgentManager:
                 and_(
                     FileAgentModel.agent_id == agent_id,
                     FileAgentModel.file_id == file_id,
+                    FileAgentModel.file_name == file_name,
                     FileAgentModel.organization_id == actor.organization_id,
                 )
             )
@@ -61,6 +64,7 @@ class FileAgentManager:
             assoc = FileAgentModel(
                 agent_id=agent_id,
                 file_id=file_id,
+                file_name=file_name,
                 organization_id=actor.organization_id,
                 is_open=is_open,
                 visible_content=visible_content,
@@ -71,7 +75,7 @@ class FileAgentManager:
 
     @enforce_types
     @trace_method
-    async def update_file_agent(
+    async def update_file_agent_by_id(
         self,
         *,
         agent_id: str,
@@ -82,7 +86,33 @@ class FileAgentManager:
     ) -> PydanticFileAgent:
         """Patch an existing association row."""
         async with db_registry.async_session() as session:
-            assoc = await self._get_association(session, agent_id, file_id, actor)
+            assoc = await self._get_association_by_file_id(session, agent_id, file_id, actor)
+
+            if is_open is not None:
+                assoc.is_open = is_open
+            if visible_content is not None:
+                assoc.visible_content = visible_content
+
+            # touch timestamp
+            assoc.last_accessed_at = datetime.now(timezone.utc)
+
+            await assoc.update_async(session, actor=actor)
+            return assoc.to_pydantic()
+
+    @enforce_types
+    @trace_method
+    async def update_file_agent_by_name(
+        self,
+        *,
+        agent_id: str,
+        file_name: str,
+        actor: PydanticUser,
+        is_open: Optional[bool] = None,
+        visible_content: Optional[str] = None,
+    ) -> PydanticFileAgent:
+        """Patch an existing association row."""
+        async with db_registry.async_session() as session:
+            assoc = await self._get_association_by_file_name(session, agent_id, file_name, actor)
 
             if is_open is not None:
                 assoc.is_open = is_open
@@ -100,15 +130,61 @@ class FileAgentManager:
     async def detach_file(self, *, agent_id: str, file_id: str, actor: PydanticUser) -> None:
         """Hard-delete the association."""
         async with db_registry.async_session() as session:
-            assoc = await self._get_association(session, agent_id, file_id, actor)
+            assoc = await self._get_association_by_file_id(session, agent_id, file_id, actor)
             await assoc.hard_delete_async(session, actor=actor)
 
     @enforce_types
     @trace_method
-    async def get_file_agent(self, *, agent_id: str, file_id: str, actor: PydanticUser) -> Optional[PydanticFileAgent]:
+    async def get_file_agent_by_id(self, *, agent_id: str, file_id: str, actor: PydanticUser) -> Optional[PydanticFileAgent]:
+        async with db_registry.async_session() as session:
+            try:
+                assoc = await self._get_association_by_file_id(session, agent_id, file_id, actor)
+                return assoc.to_pydantic()
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    @trace_method
+    async def get_all_file_blocks_by_name(
+        self,
+        *,
+        file_names: List[str],
+        actor: PydanticUser,
+    ) -> List[PydanticBlock]:
+        """
+        Retrieve multiple FileAgent associations by their IDs in a single query.
+
+        Args:
+            file_names: List of file names to retrieve
+            actor: The user making the request
+
+        Returns:
+            List of PydanticFileAgent objects found (may be fewer than requested if some IDs don't exist)
+        """
+        if not file_names:
+            return []
+
+        async with db_registry.async_session() as session:
+            # Use IN clause for efficient bulk retrieval
+            query = select(FileAgentModel).where(
+                and_(
+                    FileAgentModel.file_name.in_(file_names),
+                    FileAgentModel.organization_id == actor.organization_id,
+                )
+            )
+
+            # Execute query and get all results
+            rows = (await session.execute(query)).scalars().all()
+
+            # Convert to Pydantic models
+            return [row.to_pydantic_block() for row in rows]
+
+    @enforce_types
+    @trace_method
+    async def get_file_agent_by_file_name(self, *, agent_id: str, file_name: str, actor: PydanticUser) -> Optional[PydanticFileAgent]:
         async with db_registry.async_session() as session:
             try:
-                assoc = await self._get_association(session, agent_id, file_id, actor)
+                assoc = await self._get_association_by_file_name(session, agent_id, file_name, actor)
                 return assoc.to_pydantic()
             except NoResultFound:
                 return None
@@ -170,7 +246,7 @@ class FileAgentManager:
             await session.execute(stmt)
             await session.commit()
 
-    async def _get_association(self, session, agent_id: str, file_id: str, actor: PydanticUser) -> FileAgentModel:
+    async def _get_association_by_file_id(self, session, agent_id: str, file_id: str, actor: PydanticUser) -> FileAgentModel:
         q = select(FileAgentModel).where(
             and_(
                 FileAgentModel.agent_id == agent_id,
@@ -182,3 +258,16 @@ class FileAgentManager:
         if not assoc:
             raise NoResultFound(f"FileAgent(agent_id={agent_id}, file_id={file_id}) not found in org {actor.organization_id}")
         return assoc
+
+    async def _get_association_by_file_name(self, session, agent_id: str, file_name: str, actor: PydanticUser) -> FileAgentModel:
+        q = select(FileAgentModel).where(
+            and_(
+                FileAgentModel.agent_id == agent_id,
+                FileAgentModel.file_name == file_name,
+                FileAgentModel.organization_id == actor.organization_id,
+            )
+        )
+        assoc = await session.scalar(q)
+        if not assoc:
+            raise NoResultFound(f"FileAgent(agent_id={agent_id}, file_name={file_name}) not found in org {actor.organization_id}")
+        return assoc

letta/services/group_manager.py

@@ -7,13 +7,13 @@ from letta.orm.agent import Agent as AgentModel
 from letta.orm.errors import NoResultFound
 from letta.orm.group import Group as GroupModel
 from letta.orm.message import Message as MessageModel
+from letta.otel.tracing import trace_method
 from letta.schemas.group import Group as PydanticGroup
 from letta.schemas.group import GroupCreate, GroupUpdate, ManagerType
 from letta.schemas.letta_message import LettaMessage
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
-from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 
@@ -152,9 +152,9 @@ class GroupManager:
 
     @trace_method
     @enforce_types
-    def modify_group(self, group_id: str, group_update: GroupUpdate, actor: PydanticUser) -> PydanticGroup:
-        with db_registry.session() as session:
-            group = GroupModel.read(db_session=session, identifier=group_id, actor=actor)
+    async def modify_group_async(self, group_id: str, group_update: GroupUpdate, actor: PydanticUser) -> PydanticGroup:
+        async with db_registry.async_session() as session:
+            group = await GroupModel.read_async(db_session=session, identifier=group_id, actor=actor)
 
             sleeptime_agent_frequency = None
             max_message_buffer_length = None
@@ -206,11 +206,11 @@ class GroupManager:
             if group_update.description:
                 group.description = group_update.description
             if group_update.agent_ids:
-                self._process_agent_relationship(
+                await self._process_agent_relationship_async(
                     session=session, group=group, agent_ids=group_update.agent_ids, allow_partial=False, replace=True
                 )
 
-            group.update(session, actor=actor)
+            await group.update_async(session, actor=actor)
             return group.to_pydantic()
 
     @trace_method

letta/services/helpers/agent_manager_helper.py

@@ -1,27 +1,33 @@
 import datetime
 from typing import List, Literal, Optional
 
-from sqlalchemy import and_, asc, desc, or_, select
+import numpy as np
+from sqlalchemy import Select, and_, asc, desc, func, literal, or_, select, union_all
 from sqlalchemy.sql.expression import exists
 
 from letta import system
-from letta.constants import IN_CONTEXT_MEMORY_KEYWORD, STRUCTURED_OUTPUT_MODELS
+from letta.constants import IN_CONTEXT_MEMORY_KEYWORD, MAX_EMBEDDING_DIM, STRUCTURED_OUTPUT_MODELS
+from letta.embeddings import embedding_model
 from letta.helpers import ToolRulesSolver
 from letta.helpers.datetime_helpers import get_local_time, get_local_time_fast
+from letta.orm import AgentPassage, SourcePassage, SourcesAgents
 from letta.orm.agent import Agent as AgentModel
 from letta.orm.agents_tags import AgentsTags
 from letta.orm.errors import NoResultFound
 from letta.orm.identity import Identity
+from letta.orm.sqlite_functions import adapt_array
+from letta.otel.tracing import trace_method
 from letta.prompts import gpt_system
 from letta.schemas.agent import AgentState, AgentType
+from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import MessageRole
 from letta.schemas.letta_message_content import TextContent
 from letta.schemas.memory import Memory
 from letta.schemas.message import Message, MessageCreate
 from letta.schemas.tool_rule import ToolRule
 from letta.schemas.user import User
+from letta.settings import settings
 from letta.system import get_initial_boot_messages, get_login_event, package_function_response
-from letta.tracing import trace_method
 
 
 # Static methods
@@ -566,3 +572,367 @@ def _apply_filters(
     if base_template_id:
         query = query.where(AgentModel.base_template_id == base_template_id)
     return query
+
+
+def build_passage_query(
+    actor: User,
+    agent_id: Optional[str] = None,
+    file_id: Optional[str] = None,
+    query_text: Optional[str] = None,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    before: Optional[str] = None,
+    after: Optional[str] = None,
+    source_id: Optional[str] = None,
+    embed_query: bool = False,
+    ascending: bool = True,
+    embedding_config: Optional[EmbeddingConfig] = None,
+    agent_only: bool = False,
+) -> Select:
+    """Helper function to build the base passage query with all filters applied.
+    Supports both before and after pagination across merged source and agent passages.
+
+    Returns the query before any limit or count operations are applied.
+    """
+    embedded_text = None
+    if embed_query:
+        assert embedding_config is not None, "embedding_config must be specified for vector search"
+        assert query_text is not None, "query_text must be specified for vector search"
+        embedded_text = embedding_model(embedding_config).get_text_embedding(query_text)
+        embedded_text = np.array(embedded_text)
+        embedded_text = np.pad(embedded_text, (0, MAX_EMBEDDING_DIM - embedded_text.shape[0]), mode="constant").tolist()
+
+    # Start with base query for source passages
+    source_passages = None
+    if not agent_only:  # Include source passages
+        if agent_id is not None:
+            source_passages = (
+                select(SourcePassage, literal(None).label("agent_id"))
+                .join(SourcesAgents, SourcesAgents.source_id == SourcePassage.source_id)
+                .where(SourcesAgents.agent_id == agent_id)
+                .where(SourcePassage.organization_id == actor.organization_id)
+            )
+        else:
+            source_passages = select(SourcePassage, literal(None).label("agent_id")).where(
+                SourcePassage.organization_id == actor.organization_id
+            )
+
+        if source_id:
+            source_passages = source_passages.where(SourcePassage.source_id == source_id)
+        if file_id:
+            source_passages = source_passages.where(SourcePassage.file_id == file_id)
+
+    # Add agent passages query
+    agent_passages = None
+    if agent_id is not None:
+        agent_passages = (
+            select(
+                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,
+                literal(None).label("file_id"),
+                literal(None).label("source_id"),
+                AgentPassage.agent_id,
+            )
+            .where(AgentPassage.agent_id == agent_id)
+            .where(AgentPassage.organization_id == actor.organization_id)
+        )
+
+    # Combine queries
+    if source_passages is not None and agent_passages is not None:
+        combined_query = union_all(source_passages, agent_passages).cte("combined_passages")
+    elif agent_passages is not None:
+        combined_query = agent_passages.cte("combined_passages")
+    elif source_passages is not None:
+        combined_query = source_passages.cte("combined_passages")
+    else:
+        raise ValueError("No passages found")
+
+    # Build main query from combined CTE
+    main_query = select(combined_query)
+
+    # Apply filters
+    if start_date:
+        main_query = main_query.where(combined_query.c.created_at >= start_date)
+    if end_date:
+        main_query = main_query.where(combined_query.c.created_at <= end_date)
+    if source_id:
+        main_query = main_query.where(combined_query.c.source_id == source_id)
+    if file_id:
+        main_query = main_query.where(combined_query.c.file_id == file_id)
+
+    # Vector search
+    if embedded_text:
+        if settings.letta_pg_uri_no_default:
+            # PostgreSQL with pgvector
+            main_query = main_query.order_by(combined_query.c.embedding.cosine_distance(embedded_text).asc())
+        else:
+            # SQLite with custom vector type
+            query_embedding_binary = adapt_array(embedded_text)
+            main_query = main_query.order_by(
+                func.cosine_distance(combined_query.c.embedding, query_embedding_binary).asc(),
+                combined_query.c.created_at.asc() if ascending else combined_query.c.created_at.desc(),
+                combined_query.c.id.asc(),
+            )
+    else:
+        if query_text:
+            main_query = main_query.where(func.lower(combined_query.c.text).contains(func.lower(query_text)))
+
+    # Handle pagination
+    if before or after:
+        # Create reference CTEs
+        if before:
+            before_ref = select(combined_query.c.created_at, combined_query.c.id).where(combined_query.c.id == before).cte("before_ref")
+        if after:
+            after_ref = select(combined_query.c.created_at, combined_query.c.id).where(combined_query.c.id == after).cte("after_ref")
+
+        if before and after:
+            # Window-based query (get records between before and after)
+            main_query = main_query.where(
+                or_(
+                    combined_query.c.created_at < select(before_ref.c.created_at).scalar_subquery(),
+                    and_(
+                        combined_query.c.created_at == select(before_ref.c.created_at).scalar_subquery(),
+                        combined_query.c.id < select(before_ref.c.id).scalar_subquery(),
+                    ),
+                )
+            )
+            main_query = main_query.where(
+                or_(
+                    combined_query.c.created_at > select(after_ref.c.created_at).scalar_subquery(),
+                    and_(
+                        combined_query.c.created_at == select(after_ref.c.created_at).scalar_subquery(),
+                        combined_query.c.id > select(after_ref.c.id).scalar_subquery(),
+                    ),
+                )
+            )
+        else:
+            # Pure pagination (only before or only after)
+            if before:
+                main_query = main_query.where(
+                    or_(
+                        combined_query.c.created_at < select(before_ref.c.created_at).scalar_subquery(),
+                        and_(
+                            combined_query.c.created_at == select(before_ref.c.created_at).scalar_subquery(),
+                            combined_query.c.id < select(before_ref.c.id).scalar_subquery(),
+                        ),
+                    )
+                )
+            if after:
+                main_query = main_query.where(
+                    or_(
+                        combined_query.c.created_at > select(after_ref.c.created_at).scalar_subquery(),
+                        and_(
+                            combined_query.c.created_at == select(after_ref.c.created_at).scalar_subquery(),
+                            combined_query.c.id > select(after_ref.c.id).scalar_subquery(),
+                        ),
+                    )
+                )
+
+    # Add ordering if not already ordered by similarity
+    if not embed_query:
+        if ascending:
+            main_query = main_query.order_by(
+                combined_query.c.created_at.asc(),
+                combined_query.c.id.asc(),
+            )
+        else:
+            main_query = main_query.order_by(
+                combined_query.c.created_at.desc(),
+                combined_query.c.id.asc(),
+            )
+
+    return main_query
+
+
+def build_source_passage_query(
+    actor: User,
+    agent_id: Optional[str] = None,
+    file_id: Optional[str] = None,
+    query_text: Optional[str] = None,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    before: Optional[str] = None,
+    after: Optional[str] = None,
+    source_id: Optional[str] = None,
+    embed_query: bool = False,
+    ascending: bool = True,
+    embedding_config: Optional[EmbeddingConfig] = None,
+) -> Select:
+    """Build query for source passages with all filters applied."""
+
+    # Handle embedding for vector search
+    embedded_text = None
+    if embed_query:
+        assert embedding_config is not None, "embedding_config must be specified for vector search"
+        assert query_text is not None, "query_text must be specified for vector search"
+        embedded_text = embedding_model(embedding_config).get_text_embedding(query_text)
+        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 source passages
+    query = select(SourcePassage).where(SourcePassage.organization_id == actor.organization_id)
+
+    # If agent_id is specified, join with SourcesAgents to get only passages linked to that agent
+    if agent_id is not None:
+        query = query.join(SourcesAgents, SourcesAgents.source_id == SourcePassage.source_id)
+        query = query.where(SourcesAgents.agent_id == agent_id)
+
+    # Apply filters
+    if source_id:
+        query = query.where(SourcePassage.source_id == source_id)
+    if file_id:
+        query = query.where(SourcePassage.file_id == file_id)
+    if start_date:
+        query = query.where(SourcePassage.created_at >= start_date)
+    if end_date:
+        query = query.where(SourcePassage.created_at <= end_date)
+
+    # Handle text search or vector search
+    if embedded_text:
+        if settings.letta_pg_uri_no_default:
+            # PostgreSQL with pgvector
+            query = query.order_by(SourcePassage.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(SourcePassage.embedding, query_embedding_binary).asc(),
+                SourcePassage.created_at.asc() if ascending else SourcePassage.created_at.desc(),
+                SourcePassage.id.asc(),
+            )
+    else:
+        if query_text:
+            query = query.where(func.lower(SourcePassage.text).contains(func.lower(query_text)))
+
+    # Handle pagination
+    if before or after:
+        if before:
+            # Get the reference record
+            before_subq = select(SourcePassage.created_at, SourcePassage.id).where(SourcePassage.id == before).subquery()
+            query = query.where(
+                or_(
+                    SourcePassage.created_at < before_subq.c.created_at,
+                    and_(
+                        SourcePassage.created_at == before_subq.c.created_at,
+                        SourcePassage.id < before_subq.c.id,
+                    ),
+                )
+            )
+
+        if after:
+            # Get the reference record
+            after_subq = select(SourcePassage.created_at, SourcePassage.id).where(SourcePassage.id == after).subquery()
+            query = query.where(
+                or_(
+                    SourcePassage.created_at > after_subq.c.created_at,
+                    and_(
+                        SourcePassage.created_at == after_subq.c.created_at,
+                        SourcePassage.id > after_subq.c.id,
+                    ),
+                )
+            )
+
+    # Apply ordering if not already ordered by similarity
+    if not embed_query:
+        if ascending:
+            query = query.order_by(SourcePassage.created_at.asc(), SourcePassage.id.asc())
+        else:
+            query = query.order_by(SourcePassage.created_at.desc(), SourcePassage.id.asc())
+
+    return query
+
+
+def build_agent_passage_query(
+    actor: User,
+    agent_id: str,  # Required for agent passages
+    query_text: Optional[str] = None,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+    before: Optional[str] = None,
+    after: Optional[str] = None,
+    embed_query: bool = False,
+    ascending: bool = True,
+    embedding_config: Optional[EmbeddingConfig] = None,
+) -> Select:
+    """Build query for agent passages with all filters applied."""
+
+    # Handle embedding for vector search
+    embedded_text = None
+    if embed_query:
+        assert embedding_config is not None, "embedding_config must be specified for vector search"
+        assert query_text is not None, "query_text must be specified for vector search"
+        embedded_text = embedding_model(embedding_config).get_text_embedding(query_text)
+        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)
+
+    # Apply filters
+    if start_date:
+        query = query.where(AgentPassage.created_at >= start_date)
+    if end_date:
+        query = query.where(AgentPassage.created_at <= end_date)
+
+    # Handle text search or vector search
+    if embedded_text:
+        if settings.letta_pg_uri_no_default:
+            # PostgreSQL with pgvector
+            query = query.order_by(AgentPassage.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(),
+            )
+    else:
+        if query_text:
+            query = query.where(func.lower(AgentPassage.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()
+            query = query.where(
+                or_(
+                    AgentPassage.created_at < before_subq.c.created_at,
+                    and_(
+                        AgentPassage.created_at == before_subq.c.created_at,
+                        AgentPassage.id < before_subq.c.id,
+                    ),
+                )
+            )
+
+        if after:
+            # Get the reference record
+            after_subq = select(AgentPassage.created_at, AgentPassage.id).where(AgentPassage.id == after).subquery()
+            query = query.where(
+                or_(
+                    AgentPassage.created_at > after_subq.c.created_at,
+                    and_(
+                        AgentPassage.created_at == after_subq.c.created_at,
+                        AgentPassage.id > after_subq.c.id,
+                    ),
+                )
+            )
+
+    # 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())
+        else:
+            query = query.order_by(AgentPassage.created_at.desc(), AgentPassage.id.asc())
+
+    return query

letta/services/identity_manager.py

@@ -7,11 +7,11 @@ from sqlalchemy.exc import NoResultFound
 from letta.orm.agent import Agent as AgentModel
 from letta.orm.block import Block as BlockModel
 from letta.orm.identity import Identity as IdentityModel
+from letta.otel.tracing import trace_method
 from letta.schemas.identity import Identity as PydanticIdentity
 from letta.schemas.identity import IdentityCreate, IdentityProperty, IdentityType, IdentityUpdate, IdentityUpsert
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
-from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 

letta/services/job_manager.py

@@ -14,6 +14,7 @@ from letta.orm.message import Message as MessageModel
 from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step
 from letta.orm.step import Step as StepModel
+from letta.otel.tracing import trace_method
 from letta.schemas.enums import JobStatus, MessageRole
 from letta.schemas.job import BatchJob as PydanticBatchJob
 from letta.schemas.job import Job as PydanticJob
@@ -25,7 +26,6 @@ from letta.schemas.step import Step as PydanticStep
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
-from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 

letta/services/llm_batch_manager.py

@@ -9,6 +9,7 @@ from letta.log import get_logger
 from letta.orm import Message as MessageModel
 from letta.orm.llm_batch_items import LLMBatchItem
 from letta.orm.llm_batch_job import LLMBatchJob
+from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentStepState
 from letta.schemas.enums import AgentStepStatus, JobStatus, ProviderType
 from letta.schemas.llm_batch_job import LLMBatchItem as PydanticLLMBatchItem
@@ -17,7 +18,6 @@ from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
-from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 logger = get_logger(__name__)