letta-nightly 0.8.0.dev20250606195656__py3-none-any.whl → 0.8.3.dev20250607000559__py3-none-any.whl

letta/llm_api/helpers.py

@@ -63,11 +63,11 @@ def _convert_to_structured_output_helper(property: dict) -> dict:
 
 
 def convert_to_structured_output(openai_function: dict, allow_optional: bool = False) -> dict:
-    """Convert function call objects to structured output objects
+    """Convert function call objects to structured output objects.
 
     See: https://platform.openai.com/docs/guides/structured-outputs/supported-schemas
     """
-    description = openai_function["description"] if "description" in openai_function else ""
+    description = openai_function.get("description", "")
 
     structured_output = {
         "name": openai_function["name"],
@@ -81,54 +81,58 @@ def convert_to_structured_output(openai_function: dict, allow_optional: bool = F
         },
     }
 
-    # This code needs to be able to handle nested properties
-    # For example, the param details may have "type" + "description",
-    # but if "type" is "object" we expected "properties", where each property has details
-    # and if "type" is "array" we expect "items": <type>
     for param, details in openai_function["parameters"]["properties"].items():
         param_type = details["type"]
-        description = details.get("description", "")
+        param_description = details.get("description", "")
 
         if param_type == "object":
             if "properties" not in details:
-                # Structured outputs requires the properties on dicts be specified ahead of time
-                raise ValueError(f"Property {param} of type object is missing properties")
+                raise ValueError(f"Property {param} of type object is missing 'properties'")
             structured_output["parameters"]["properties"][param] = {
                 "type": "object",
-                "description": description,
+                "description": param_description,
                 "properties": {k: _convert_to_structured_output_helper(v) for k, v in details["properties"].items()},
                 "additionalProperties": False,
                 "required": list(details["properties"].keys()),
             }
 
         elif param_type == "array":
-            structured_output["parameters"]["properties"][param] = {
-                "type": "array",
-                "description": description,
-                "items": _convert_to_structured_output_helper(details["items"]),
-            }
+            items_schema = details.get("items")
+            prefix_items_schema = details.get("prefixItems")
+
+            if prefix_items_schema:
+                # assume fixed-length tuple — safe fallback to use first type for items
+                fallback_item = prefix_items_schema[0] if isinstance(prefix_items_schema, list) else prefix_items_schema
+                structured_output["parameters"]["properties"][param] = {
+                    "type": "array",
+                    "description": param_description,
+                    "prefixItems": [_convert_to_structured_output_helper(item) for item in prefix_items_schema],
+                    "items": _convert_to_structured_output_helper(fallback_item),
+                    "minItems": details.get("minItems", len(prefix_items_schema)),
+                    "maxItems": details.get("maxItems", len(prefix_items_schema)),
+                }
+            elif items_schema:
+                structured_output["parameters"]["properties"][param] = {
+                    "type": "array",
+                    "description": param_description,
+                    "items": _convert_to_structured_output_helper(items_schema),
+                }
+            else:
+                raise ValueError(f"Array param '{param}' is missing both 'items' and 'prefixItems'")
 
         else:
-            structured_output["parameters"]["properties"][param] = {
-                "type": param_type,  # simple type
-                "description": description,
+            prop = {
+                "type": param_type,
+                "description": param_description,
             }
-
-        if "enum" in details:
-            structured_output["parameters"]["properties"][param]["enum"] = details["enum"]
+            if "enum" in details:
+                prop["enum"] = details["enum"]
+            structured_output["parameters"]["properties"][param] = prop
 
     if not allow_optional:
-        # Add all properties to required list
         structured_output["parameters"]["required"] = list(structured_output["parameters"]["properties"].keys())
-
     else:
-        # See what parameters exist that aren't required
-        # Those are implied "optional" types
-        # For those types, turn each of them into a union type with "null"
-        # e.g.
-        # "type": "string" -> "type": ["string", "null"]
-        # TODO
-        raise NotImplementedError
+        raise NotImplementedError("Optional parameter handling is not implemented.")
 
     return structured_output
 
@@ -292,6 +296,8 @@ def unpack_inner_thoughts_from_kwargs(choice: Choice, inner_thoughts_key: str) -
 
         except json.JSONDecodeError as e:
             warnings.warn(f"Failed to strip inner thoughts from kwargs: {e}")
+            print(f"\nFailed to strip inner thoughts from kwargs: {e}")
+            print(f"\nTool call arguments: {tool_call.function.arguments}")
             raise e
     else:
         warnings.warn(f"Did not find tool call in message: {str(message)}")

letta/llm_api/llm_api_tools.py

@@ -26,6 +26,7 @@ from letta.local_llm.chat_completion_proxy import get_chat_completion
 from letta.local_llm.constants import INNER_THOUGHTS_KWARG, INNER_THOUGHTS_KWARG_DESCRIPTION
 from letta.local_llm.utils import num_tokens_from_functions, num_tokens_from_messages
 from letta.orm.user import User
+from letta.otel.tracing import log_event, trace_method
 from letta.schemas.enums import ProviderCategory
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message
@@ -35,7 +36,6 @@ from letta.schemas.provider_trace import ProviderTraceCreate
 from letta.services.telemetry_manager import TelemetryManager
 from letta.settings import ModelSettings
 from letta.streaming_interface import AgentChunkStreamingInterface, AgentRefreshStreamingInterface
-from letta.tracing import log_event, trace_method
 
 LLM_API_PROVIDER_OPTIONS = ["openai", "azure", "anthropic", "google_ai", "cohere", "local", "groq", "deepseek"]
 

letta/llm_api/llm_client_base.py

@@ -1,3 +1,4 @@
+import json
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
@@ -6,13 +7,13 @@ from openai import AsyncStream, Stream
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from letta.errors import LLMError
+from letta.otel.tracing import log_event, trace_method
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message
 from letta.schemas.openai.chat_completion_response import ChatCompletionResponse
 from letta.schemas.provider_trace import ProviderTraceCreate
 from letta.services.telemetry_manager import TelemetryManager
-from letta.tracing import log_event, trace_method
 
 if TYPE_CHECKING:
     from letta.orm import User
@@ -186,3 +187,30 @@ class LLMClientBase:
             An LLMError subclass that represents the error in a provider-agnostic way
         """
         return LLMError(f"Unhandled LLM error: {str(e)}")
+
+    def _fix_truncated_json_response(self, response: ChatCompletionResponse) -> ChatCompletionResponse:
+        """
+        Fixes truncated JSON responses by ensuring the content is properly formatted.
+        This is a workaround for some providers that may return incomplete JSON.
+        """
+        if response.choices and response.choices[0].message and response.choices[0].message.tool_calls:
+            tool_call_args_str = response.choices[0].message.tool_calls[0].function.arguments
+            try:
+                json.loads(tool_call_args_str)
+            except json.JSONDecodeError:
+                try:
+                    json_str_end = ""
+                    quote_count = tool_call_args_str.count('"')
+                    if quote_count % 2 != 0:
+                        json_str_end = json_str_end + '"'
+
+                    open_braces = tool_call_args_str.count("{")
+                    close_braces = tool_call_args_str.count("}")
+                    missing_braces = open_braces - close_braces
+                    json_str_end += "}" * missing_braces
+                    fixed_tool_call_args_str = tool_call_args_str[: -len(json_str_end)] + json_str_end
+                    json.loads(fixed_tool_call_args_str)
+                    response.choices[0].message.tool_calls[0].function.arguments = fixed_tool_call_args_str
+                except json.JSONDecodeError:
+                    pass
+        return response

letta/llm_api/openai.py

@@ -19,6 +19,7 @@ from letta.llm_api.openai_client import (
 from letta.local_llm.constants import INNER_THOUGHTS_KWARG, INNER_THOUGHTS_KWARG_DESCRIPTION, INNER_THOUGHTS_KWARG_DESCRIPTION_GO_FIRST
 from letta.local_llm.utils import num_tokens_from_functions, num_tokens_from_messages
 from letta.log import get_logger
+from letta.otel.tracing import log_event
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message as _Message
 from letta.schemas.message import MessageRole as _MessageRole
@@ -36,7 +37,6 @@ from letta.schemas.openai.chat_completion_response import (
 )
 from letta.schemas.openai.embedding_response import EmbeddingResponse
 from letta.streaming_interface import AgentChunkStreamingInterface, AgentRefreshStreamingInterface
-from letta.tracing import log_event
 from letta.utils import get_tool_call_id, smart_urljoin
 
 logger = get_logger(__name__)

letta/llm_api/openai_client.py

@@ -8,11 +8,11 @@ from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 
 from letta.constants import LETTA_MODEL_ENDPOINT
 from letta.errors import (
+    ContextWindowExceededError,
     ErrorCode,
     LLMAuthenticationError,
     LLMBadRequestError,
     LLMConnectionError,
-    LLMContextWindowExceededError,
     LLMNotFoundError,
     LLMPermissionDeniedError,
     LLMRateLimitError,
@@ -23,6 +23,7 @@ from letta.llm_api.helpers import add_inner_thoughts_to_functions, convert_to_st
 from letta.llm_api.llm_client_base import LLMClientBase
 from letta.local_llm.constants import INNER_THOUGHTS_KWARG, INNER_THOUGHTS_KWARG_DESCRIPTION, INNER_THOUGHTS_KWARG_DESCRIPTION_GO_FIRST
 from letta.log import get_logger
+from letta.otel.tracing import trace_method
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import ProviderCategory, ProviderType
 from letta.schemas.llm_config import LLMConfig
@@ -34,7 +35,6 @@ from letta.schemas.openai.chat_completion_request import Tool as OpenAITool
 from letta.schemas.openai.chat_completion_request import ToolFunctionChoice, cast_message_to_subtype
 from letta.schemas.openai.chat_completion_response import ChatCompletionResponse
 from letta.settings import model_settings
-from letta.tracing import trace_method
 
 logger = get_logger(__name__)
 
@@ -280,7 +280,7 @@ class OpenAIClient(LLMClientBase):
         # OpenAI's response structure directly maps to ChatCompletionResponse
         # We just need to instantiate the Pydantic model for validation and type safety.
         chat_completion_response = ChatCompletionResponse(**response_data)
-
+        chat_completion_response = self._fix_truncated_json_response(chat_completion_response)
         # Unpack inner thoughts if they were embedded in function arguments
         if llm_config.put_inner_thoughts_in_kwargs:
             chat_completion_response = unpack_all_inner_thoughts_from_kwargs(
@@ -342,11 +342,9 @@ class OpenAIClient(LLMClientBase):
             # Check message content if finer-grained errors are needed
             # Example: if "context_length_exceeded" in str(e): return LLMContextLengthExceededError(...)
             # TODO: This is a super soft check. Not sure if we can do better, needs more investigation.
-            if "context" in str(e):
-                return LLMContextWindowExceededError(
-                    message=f"Bad request to OpenAI (context length exceeded): {str(e)}",
-                    code=ErrorCode.INVALID_ARGUMENT,  # Or more specific if detectable
-                    details=e.body,
+            if "This model's maximum context length is" in str(e):
+                return ContextWindowExceededError(
+                    message=f"Bad request to OpenAI (context window exceeded): {str(e)}",
                 )
             else:
                 return LLMBadRequestError(

letta/local_llm/chat_completion_proxy.py

@@ -20,9 +20,9 @@ from letta.local_llm.utils import count_tokens, get_available_wrappers
 from letta.local_llm.vllm.api import get_vllm_completion
 from letta.local_llm.webui.api import get_webui_completion
 from letta.local_llm.webui.legacy_api import get_webui_completion as get_webui_completion_legacy
+from letta.otel.tracing import log_event
 from letta.prompts.gpt_summarize import SYSTEM as SUMMARIZE_SYSTEM_MESSAGE
 from letta.schemas.openai.chat_completion_response import ChatCompletionResponse, Choice, Message, ToolCall, UsageStatistics
-from letta.tracing import log_event
 from letta.utils import get_tool_call_id
 
 has_shown_warning = False

letta/memory.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Callable, Dict, List
 from letta.constants import MESSAGE_SUMMARY_REQUEST_ACK
 from letta.llm_api.llm_api_tools import create
 from letta.llm_api.llm_client import LLMClient
+from letta.otel.tracing import trace_method
 from letta.prompts.gpt_summarize import SYSTEM as SUMMARY_PROMPT_SYSTEM
 from letta.schemas.agent import AgentState
 from letta.schemas.enums import MessageRole
@@ -10,7 +11,6 @@ from letta.schemas.letta_message_content import TextContent
 from letta.schemas.memory import Memory
 from letta.schemas.message import Message
 from letta.settings import summarizer_settings
-from letta.tracing import trace_method
 from letta.utils import count_tokens, printd
 
 if TYPE_CHECKING:

letta/orm/enums.py

@@ -9,6 +9,7 @@ class ToolType(str, Enum):
     LETTA_SLEEPTIME_CORE = "letta_sleeptime_core"
     LETTA_VOICE_SLEEPTIME_CORE = "letta_voice_sleeptime_core"
     LETTA_BUILTIN = "letta_builtin"
+    LETTA_FILES_CORE = "letta_files_core"
     EXTERNAL_COMPOSIO = "external_composio"
     EXTERNAL_LANGCHAIN = "external_langchain"
     # TODO is "external" the right name here? Since as of now, MCP is local / doesn't support remote?

letta/orm/file.py

@@ -1,10 +1,13 @@
+import uuid
 from typing import TYPE_CHECKING, List, Optional
 
-from sqlalchemy import Integer, String
+from sqlalchemy import ForeignKey, Index, Integer, String, Text, UniqueConstraint, desc
+from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.mixins import OrganizationMixin, SourceMixin
 from letta.orm.sqlalchemy_base import SqlalchemyBase
+from letta.schemas.enums import FileProcessingStatus
 from letta.schemas.file import FileMetadata as PydanticFileMetadata
 
 if TYPE_CHECKING:
@@ -14,11 +17,36 @@ if TYPE_CHECKING:
     from letta.orm.source import Source
 
 
-class FileMetadata(SqlalchemyBase, OrganizationMixin, SourceMixin):
+# TODO: Note that this is NOT organization scoped, this is potentially dangerous if we misuse this
+# TODO: This should ONLY be manipulated internally in relation to FileMetadata.content
+# TODO: Leaving organization_id out of this for now for simplicity
+class FileContent(SqlalchemyBase):
+    """Holds the full text content of a file (potentially large)."""
+
+    __tablename__ = "file_contents"
+    __table_args__ = (UniqueConstraint("file_id", name="uq_file_contents_file_id"),)
+
+    # TODO: We want to migrate all the ORM models to do this, so we will need to move this to the SqlalchemyBase
+    # TODO: Some still rely on the Pydantic object to do this
+    id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"file_content-{uuid.uuid4()}")
+    file_id: Mapped[str] = mapped_column(ForeignKey("files.id", ondelete="CASCADE"), nullable=False, doc="Foreign key to files table.")
+
+    text: Mapped[str] = mapped_column(Text, nullable=False, doc="Full plain-text content of the file (e.g., extracted from a PDF).")
+
+    # back-reference to FileMetadata
+    file: Mapped["FileMetadata"] = relationship(back_populates="content", lazy="selectin")
+
+
+class FileMetadata(SqlalchemyBase, OrganizationMixin, SourceMixin, AsyncAttrs):
     """Represents an uploaded file."""
 
     __tablename__ = "files"
     __pydantic_model__ = PydanticFileMetadata
+    __table_args__ = (
+        Index("ix_files_org_created", "organization_id", desc("created_at")),
+        Index("ix_files_source_created", "source_id", desc("created_at")),
+        Index("ix_files_processing_status", "processing_status"),
+    )
 
     file_name: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The name of the file.")
     file_path: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The file path on the system.")
@@ -26,6 +54,11 @@ class FileMetadata(SqlalchemyBase, OrganizationMixin, SourceMixin):
     file_size: Mapped[Optional[int]] = mapped_column(Integer, nullable=True, doc="The size of the file in bytes.")
     file_creation_date: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The creation date of the file.")
     file_last_modified_date: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The last modified date of the file.")
+    processing_status: Mapped[FileProcessingStatus] = mapped_column(
+        String, default=FileProcessingStatus.PENDING, nullable=False, doc="The current processing status of the file."
+    )
+
+    error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Any error message encountered during processing.")
 
     # relationships
     organization: Mapped["Organization"] = relationship("Organization", back_populates="files", lazy="selectin")
@@ -33,4 +66,48 @@ class FileMetadata(SqlalchemyBase, OrganizationMixin, SourceMixin):
     source_passages: Mapped[List["SourcePassage"]] = relationship(
         "SourcePassage", back_populates="file", lazy="selectin", cascade="all, delete-orphan"
     )
-    file_agents: Mapped[List["FileAgent"]] = relationship("FileAgent", back_populates="file", lazy="selectin")
+    file_agents: Mapped[List["FileAgent"]] = relationship(
+        "FileAgent",
+        back_populates="file",
+        lazy="selectin",
+        cascade="all, delete-orphan",
+        passive_deletes=True,  # ← add this
+    )
+    content: Mapped[Optional["FileContent"]] = relationship(
+        "FileContent",
+        uselist=False,
+        back_populates="file",
+        lazy="raise",  # raises if you access without eager load
+        cascade="all, delete-orphan",
+    )
+
+    async def to_pydantic_async(self, include_content: bool = False) -> PydanticFileMetadata:
+        """
+        Async version of `to_pydantic` that supports optional relationship loading
+        without requiring `expire_on_commit=False`.
+        """
+
+        # Load content relationship if requested
+        if include_content:
+            content_obj = await self.awaitable_attrs.content
+            content_text = content_obj.text if content_obj else None
+        else:
+            content_text = None
+
+        return PydanticFileMetadata(
+            id=self.id,
+            organization_id=self.organization_id,
+            source_id=self.source_id,
+            file_name=self.file_name,
+            file_path=self.file_path,
+            file_type=self.file_type,
+            file_size=self.file_size,
+            file_creation_date=self.file_creation_date,
+            file_last_modified_date=self.file_last_modified_date,
+            processing_status=self.processing_status,
+            error_message=self.error_message,
+            created_at=self.created_at,
+            updated_at=self.updated_at,
+            is_deleted=self.is_deleted,
+            content=content_text,
+        )

letta/orm/files_agents.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional
 from sqlalchemy import Boolean, DateTime, ForeignKey, Index, String, Text, UniqueConstraint, func
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
+from letta.constants import CORE_MEMORY_SOURCE_CHAR_LIMIT, FILE_IS_TRUNCATED_WARNING
 from letta.orm.mixins import OrganizationMixin
 from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.block import Block as PydanticBlock
@@ -26,6 +27,8 @@ class FileAgent(SqlalchemyBase, OrganizationMixin):
     __table_args__ = (
         Index("ix_files_agents_file_id_agent_id", "file_id", "agent_id"),
         UniqueConstraint("file_id", "agent_id", name="uq_files_agents_file_agent"),
+        UniqueConstraint("agent_id", "file_name", name="uq_files_agents_agent_file_name"),
+        Index("ix_files_agents_agent_file_name", "agent_id", "file_name"),
     )
     __pydantic_model__ = PydanticFileAgent
 
@@ -33,6 +36,7 @@ class FileAgent(SqlalchemyBase, OrganizationMixin):
     # TODO: Some still rely on the Pydantic object to do this
     id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"file_agent-{uuid.uuid4()}")
     file_id: Mapped[str] = mapped_column(String, ForeignKey("files.id", ondelete="CASCADE"), primary_key=True, doc="ID of the file.")
+    file_name: Mapped[str] = mapped_column(String, nullable=False, doc="Denormalized copy of files.file_name; unique per agent.")
     agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True, doc="ID of the agent.")
 
     is_open: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True, doc="True if the agent currently has the file open.")
@@ -55,11 +59,20 @@ class FileAgent(SqlalchemyBase, OrganizationMixin):
         "FileMetadata",
         foreign_keys=[file_id],
         lazy="selectin",
+        back_populates="file_agents",
+        passive_deletes=True,  # ← add this
     )
 
     # TODO: This is temporary as we figure out if we want FileBlock as a first class citizen
     def to_pydantic_block(self) -> PydanticBlock:
         visible_content = self.visible_content if self.visible_content and self.is_open else ""
+
+        # Truncate content and add warnings here when converting from FileAgent to Block
+        if len(visible_content) > CORE_MEMORY_SOURCE_CHAR_LIMIT:
+            truncated_warning = f"...[TRUNCATED]\n{FILE_IS_TRUNCATED_WARNING}"
+            visible_content = visible_content[: CORE_MEMORY_SOURCE_CHAR_LIMIT - len(truncated_warning)]
+            visible_content += truncated_warning
+
         return PydanticBlock(
             organization_id=self.organization_id,
             value=visible_content,

letta/orm/passage.py

@@ -47,6 +47,8 @@ class SourcePassage(BasePassage, FileMixin, SourceMixin):
 
     __tablename__ = "source_passages"
 
+    file_name: Mapped[str] = mapped_column(doc="The name of the file that this passage was derived from")
+
     @declared_attr
     def file(cls) -> Mapped["FileMetadata"]:
         """Relationship to file"""

letta/orm/sqlalchemy_base.py

@@ -1,13 +1,15 @@
+import inspect
 from datetime import datetime
 from enum import Enum
 from functools import wraps
 from pprint import pformat
 from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Union
 
-from sqlalchemy import String, and_, delete, func, or_, select, text
+from sqlalchemy import Sequence, String, and_, delete, func, or_, select, text
 from sqlalchemy.exc import DBAPIError, IntegrityError, TimeoutError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Mapped, Session, mapped_column
+from sqlalchemy.orm.interfaces import ORMOption
 
 from letta.log import get_logger
 from letta.orm.base import Base, CommonSqlalchemyMetaMixins
@@ -23,16 +25,28 @@ logger = get_logger(__name__)
 
 def handle_db_timeout(func):
     """Decorator to handle SQLAlchemy TimeoutError and wrap it in a custom exception."""
+    if not inspect.iscoroutinefunction(func):
 
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except TimeoutError as e:
-            logger.error(f"Timeout while executing {func.__name__} with args {args} and kwargs {kwargs}: {e}")
-            raise DatabaseTimeoutError(message=f"Timeout occurred in {func.__name__}.", original_exception=e)
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except TimeoutError as e:
+                logger.error(f"Timeout while executing {func.__name__} with args {args} and kwargs {kwargs}: {e}")
+                raise DatabaseTimeoutError(message=f"Timeout occurred in {func.__name__}.", original_exception=e)
+
+        return wrapper
+    else:
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            try:
+                return await func(*args, **kwargs)
+            except TimeoutError as e:
+                logger.error(f"Timeout while executing {func.__name__} with args {args} and kwargs {kwargs}: {e}")
+                raise DatabaseTimeoutError(message=f"Timeout occurred in {func.__name__}.", original_exception=e)
 
-    return wrapper
+        return async_wrapper
 
 
 class AccessType(str, Enum):
@@ -163,6 +177,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         join_conditions: Optional[Union[Tuple, List]] = None,
         identifier_keys: Optional[List[str]] = None,
         identity_id: Optional[str] = None,
+        query_options: Sequence[ORMOption] | None = None,  # ← new
         **kwargs,
     ) -> List["SqlalchemyBase"]:
         """
@@ -224,6 +239,9 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             identity_id=identity_id,
             **kwargs,
         )
+        if query_options:
+            for opt in query_options:
+                query = query.options(opt)
 
         # Execute the query
         results = await db_session.execute(query)
@@ -472,14 +490,19 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         Raises:
             NoResultFound: if the object is not found
         """
+        from letta.settings import settings
+
         identifiers = [] if identifier is None else [identifier]
         query, query_conditions = cls._read_multiple_preprocess(identifiers, actor, access, access_type, check_is_deleted, **kwargs)
-        await db_session.execute(text("SET LOCAL enable_seqscan = OFF"))
+
+        if settings.letta_pg_uri_no_default:
+            await db_session.execute(text("SET LOCAL enable_seqscan = OFF"))
         try:
             result = await db_session.execute(query)
             item = result.scalar_one_or_none()
         finally:
-            await db_session.execute(text("SET LOCAL enable_seqscan = ON"))
+            if settings.letta_pg_uri_no_default:
+                await db_session.execute(text("SET LOCAL enable_seqscan = ON"))
 
         if item is None:
             raise NoResultFound(f"{cls.__name__} not found with {', '.join(query_conditions if query_conditions else ['no conditions'])}")

letta/otel/context.py

@@ -0,0 +1,25 @@
+from contextvars import ContextVar
+from typing import Any, Dict
+
+# Create context var at module level (outside middleware)
+request_attributes: ContextVar[Dict[str, Any]] = ContextVar("request_attributes", default={})
+
+
+# Helper functions
+def set_ctx_attributes(attrs: Dict[str, Any]):
+    """Set attributes in current context"""
+    current = request_attributes.get()
+    new_attrs = {**current, **attrs}
+    request_attributes.set(new_attrs)
+
+
+def add_ctx_attribute(key: str, value: Any):
+    """Add single attribute to current context"""
+    current = request_attributes.get()
+    new_attrs = {**current, key: value}
+    request_attributes.set(new_attrs)
+
+
+def get_ctx_attributes() -> Dict[str, Any]:
+    """Get all attributes from current context"""
+    return request_attributes.get()