letta-nightly 0.7.13.dev20250511104036__py3-none-any.whl → 0.7.14.dev20250513020711__py3-none-any.whl

letta/orm/blocks_agents.py

@@ -1,4 +1,4 @@
-from sqlalchemy import ForeignKey, ForeignKeyConstraint, String, UniqueConstraint
+from sqlalchemy import ForeignKey, ForeignKeyConstraint, Index, String, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column
 
 from letta.orm.base import Base
@@ -18,6 +18,8 @@ class BlocksAgents(Base):
             ["block_id", "block_label"], ["block.id", "block.label"], name="fk_block_id_label", deferrable=True, initially="DEFERRED"
         ),
         UniqueConstraint("agent_id", "block_id", name="unique_agent_block"),
+        Index("ix_blocks_agents_block_label_agent_id", "block_label", "agent_id"),
+        Index("ix_blocks_block_label", "block_label"),
     )
 
     # unique agent + block label

letta/orm/sqlalchemy_base.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Union
 
 from sqlalchemy import String, and_, func, or_, select
 from sqlalchemy.exc import DBAPIError, IntegrityError, TimeoutError
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Mapped, Session, mapped_column
 
 from letta.log import get_logger
@@ -300,6 +301,44 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             raise NoResultFound(f"{cls.__name__} not found with {', '.join(conditions if conditions else ['no conditions'])}")
         return found[0]
 
+    @classmethod
+    @handle_db_timeout
+    async def read_async(
+        cls,
+        db_session: "Session",
+        identifier: Optional[str] = None,
+        actor: Optional["User"] = None,
+        access: Optional[List[Literal["read", "write", "admin"]]] = ["read"],
+        access_type: AccessType = AccessType.ORGANIZATION,
+        **kwargs,
+    ) -> "SqlalchemyBase":
+        """The primary accessor for an ORM record. Async version of read method.
+        Args:
+            db_session: the database session to use when retrieving the record
+            identifier: the identifier of the record to read, can be the id string or the UUID object for backwards compatibility
+            actor: if specified, results will be scoped only to records the user is able to access
+            access: if actor is specified, records will be filtered to the minimum permission level for the actor
+            kwargs: additional arguments to pass to the read, used for more complex objects
+        Returns:
+            The matching object
+        Raises:
+            NoResultFound: if the object is not found
+        """
+        # this is ok because read_multiple will check if the
+        identifiers = [] if identifier is None else [identifier]
+        found = await cls.read_multiple_async(db_session, identifiers, actor, access, access_type, **kwargs)
+        if len(found) == 0:
+            # for backwards compatibility.
+            conditions = []
+            if identifier:
+                conditions.append(f"id={identifier}")
+            if actor:
+                conditions.append(f"access level in {access} for {actor}")
+            if hasattr(cls, "is_deleted"):
+                conditions.append("is_deleted=False")
+            raise NoResultFound(f"{cls.__name__} not found with {', '.join(conditions if conditions else ['no conditions'])}")
+        return found[0]
+
     @classmethod
     @handle_db_timeout
     def read_multiple(
@@ -323,6 +362,38 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         Raises:
             NoResultFound: if the object is not found
         """
+        query, query_conditions = cls._read_multiple_preprocess(identifiers, actor, access, access_type, **kwargs)
+        results = db_session.execute(query).scalars().all()
+        return cls._read_multiple_postprocess(results, identifiers, query_conditions)
+
+    @classmethod
+    @handle_db_timeout
+    async def read_multiple_async(
+        cls,
+        db_session: "AsyncSession",
+        identifiers: List[str] = [],
+        actor: Optional["User"] = None,
+        access: Optional[List[Literal["read", "write", "admin"]]] = ["read"],
+        access_type: AccessType = AccessType.ORGANIZATION,
+        **kwargs,
+    ) -> List["SqlalchemyBase"]:
+        """
+        Async version of read_multiple(...)
+        The primary accessor for ORM record(s)
+        """
+        query, query_conditions = cls._read_multiple_preprocess(identifiers, actor, access, access_type, **kwargs)
+        results = await db_session.execute(query)
+        return cls._read_multiple_postprocess(results.scalars().all(), identifiers, query_conditions)
+
+    @classmethod
+    def _read_multiple_preprocess(
+        cls,
+        identifiers: List[str],
+        actor: Optional["User"],
+        access: Optional[List[Literal["read", "write", "admin"]]],
+        access_type: AccessType,
+        **kwargs,
+    ):
         logger.debug(f"Reading {cls.__name__} with ID(s): {identifiers} with actor={actor}")
 
         # Start the query
@@ -350,7 +421,10 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             query = query.where(cls.is_deleted == False)
             query_conditions.append("is_deleted=False")
 
-        results = db_session.execute(query).scalars().all()
+        return query, query_conditions
+
+    @classmethod
+    def _read_multiple_postprocess(cls, results, identifiers: List[str], query_conditions) -> List["SqlalchemyBase"]:
         if results:  # if empty list a.k.a. no results
             if len(identifiers) > 0:
                 # find which identifiers were not found
@@ -471,6 +545,22 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         db_session.refresh(self)
         return self
 
+    @handle_db_timeout
+    async def update_async(self, db_session: AsyncSession, actor: "User | None" = None, no_commit: bool = False) -> "SqlalchemyBase":
+        """Async version of update function"""
+        logger.debug(...)
+        if actor:
+            self._set_created_and_updated_by_fields(actor.id)
+        self.set_updated_at()
+
+        db_session.add(self)
+        if no_commit:
+            await db_session.flush()
+        else:
+            await db_session.commit()
+        await db_session.refresh(self)
+        return self
+
     @classmethod
     @handle_db_timeout
     def size(

letta/schemas/message.py

@@ -404,7 +404,7 @@ class Message(BaseMessage):
                     stdout=self.tool_returns[0].stdout if self.tool_returns else None,
                     stderr=self.tool_returns[0].stderr if self.tool_returns else None,
                     name=self.name,
-                    otid=self.id.replace("message-", ""),
+                    otid=Message.generate_otid_from_id(self.id, len(messages)),
                     sender_id=self.sender_id,
                     step_id=self.step_id,
                 )

letta/serialize_schemas/marshmallow_agent.py

@@ -1,6 +1,7 @@
 from typing import Dict
 
 from marshmallow import fields, post_dump, pre_load
+from sqlalchemy.orm import sessionmaker
 
 import letta
 from letta.orm import Agent
@@ -14,7 +15,6 @@ from letta.serialize_schemas.marshmallow_custom_fields import EmbeddingConfigFie
 from letta.serialize_schemas.marshmallow_message import SerializedMessageSchema
 from letta.serialize_schemas.marshmallow_tag import SerializedAgentTagSchema
 from letta.serialize_schemas.marshmallow_tool import SerializedToolSchema
-from letta.server.db import SessionLocal
 
 
 class MarshmallowAgentSchema(BaseSchema):
@@ -41,7 +41,7 @@ class MarshmallowAgentSchema(BaseSchema):
     tool_exec_environment_variables = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema))
     tags = fields.List(fields.Nested(SerializedAgentTagSchema))
 
-    def __init__(self, *args, session: SessionLocal, actor: User, **kwargs):
+    def __init__(self, *args, session: sessionmaker, actor: User, **kwargs):
         super().__init__(*args, actor=actor, **kwargs)
         self.session = session
 
@@ -60,9 +60,9 @@ class MarshmallowAgentSchema(BaseSchema):
         After dumping the agent, load all its Message rows and serialize them here.
         """
         # TODO: This is hacky, but want to move fast, please refactor moving forward
-        from letta.server.db import db_context as session_maker
+        from letta.server.db import db_registry
 
-        with session_maker() as session:
+        with db_registry.session() as session:
             agent_id = data.get("id")
             msgs = (
                 session.query(MessageModel)

letta/server/db.py

@@ -1,28 +1,19 @@
 import os
 import threading
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, AsyncGenerator, Generator
 
 from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
-from sqlalchemy import create_engine
+from sqlalchemy import Engine, create_engine
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 from sqlalchemy.orm import sessionmaker
 
 from letta.config import LettaConfig
 from letta.log import get_logger
-from letta.orm import Base
 from letta.settings import settings
 
-# Use globals for the lock and initialization flag
-_engine_lock = threading.Lock()
-_engine_initialized = False
-
-# Create variables in global scope but don't initialize them yet
-config = LettaConfig.load()
-logger = get_logger(__name__)
-engine = None
-SessionLocal = None
-
 
 def print_sqlite_schema_error():
     """Print a formatted error message for SQLite schema issues"""
@@ -54,86 +45,187 @@ def db_error_handler():
         exit(1)
 
 
-def initialize_engine():
-    """Initialize the database engine only when needed."""
-    global engine, SessionLocal, _engine_initialized
-
-    with _engine_lock:
-        # Check again inside the lock to prevent race conditions
-        if _engine_initialized:
-            return
-
-        if settings.letta_pg_uri_no_default:
-            logger.info("Creating postgres engine")
-            config.recall_storage_type = "postgres"
-            config.recall_storage_uri = settings.letta_pg_uri_no_default
-            config.archival_storage_type = "postgres"
-            config.archival_storage_uri = settings.letta_pg_uri_no_default
-
-            # create engine
-            engine = create_engine(
-                settings.letta_pg_uri,
-                # f"{settings.letta_pg_uri}?options=-c%20client_encoding=UTF8",
-                pool_size=settings.pg_pool_size,
-                max_overflow=settings.pg_max_overflow,
-                pool_timeout=settings.pg_pool_timeout,
-                pool_recycle=settings.pg_pool_recycle,
-                echo=settings.pg_echo,
-                # connect_args={"client_encoding": "utf8"},
-            )
-        else:
-            # TODO: don't rely on config storage
-            engine_path = "sqlite:///" + os.path.join(config.recall_storage_path, "sqlite.db")
-            logger.info("Creating sqlite engine " + engine_path)
-
-            engine = create_engine(engine_path)
-
-            # Store the original connect method
-            original_connect = engine.connect
-
-            def wrapped_connect(*args, **kwargs):
-                with db_error_handler():
-                    # Get the connection
-                    connection = original_connect(*args, **kwargs)
-
-                    # Store the original execution method
-                    original_execute = connection.execute
-
-                    # Wrap the execute method of the connection
-                    def wrapped_execute(*args, **kwargs):
-                        with db_error_handler():
-                            return original_execute(*args, **kwargs)
-
-                    # Replace the connection's execute method
-                    connection.execute = wrapped_execute
-
-                    return connection
-
-            # Replace the engine's connect method
-            engine.connect = wrapped_connect
-
-            Base.metadata.create_all(bind=engine)
-
-        # Create the session factory
-        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-        _engine_initialized = True
+class DatabaseRegistry:
+    """Registry for database connections and sessions.
+
+    This class manages both synchronous and asynchronous database connections
+    and provides context managers for session handling.
+    """
+
+    def __init__(self):
+        self._engines: dict[str, Engine] = {}
+        self._async_engines: dict[str, AsyncEngine] = {}
+        self._session_factories: dict[str, sessionmaker] = {}
+        self._async_session_factories: dict[str, async_sessionmaker] = {}
+        self._initialized: dict[str, bool] = {"sync": False, "async": False}
+        self._lock = threading.Lock()
+        self.config = LettaConfig.load()
+        self.logger = get_logger(__name__)
+
+    def initialize_sync(self, force: bool = False) -> None:
+        """Initialize the synchronous database engine if not already initialized."""
+        with self._lock:
+            if self._initialized.get("sync") and not force:
+                return
+
+            # Postgres engine
+            if settings.letta_pg_uri_no_default:
+                self.logger.info("Creating postgres engine")
+                self.config.recall_storage_type = "postgres"
+                self.config.recall_storage_uri = settings.letta_pg_uri_no_default
+                self.config.archival_storage_type = "postgres"
+                self.config.archival_storage_uri = settings.letta_pg_uri_no_default
+
+                engine = create_engine(
+                    settings.letta_pg_uri,
+                    # f"{settings.letta_pg_uri}?options=-c%20client_encoding=UTF8",
+                    pool_size=settings.pg_pool_size,
+                    max_overflow=settings.pg_max_overflow,
+                    pool_timeout=settings.pg_pool_timeout,
+                    pool_recycle=settings.pg_pool_recycle,
+                    echo=settings.pg_echo,
+                    # connect_args={"client_encoding": "utf8"},
+                )
+
+                self._engines["default"] = engine
+            # SQLite engine
+            else:
+                from letta.orm import Base
+
+                # TODO: don't rely on config storage
+                engine_path = "sqlite:///" + os.path.join(self.config.recall_storage_path, "sqlite.db")
+                self.logger.info("Creating sqlite engine " + engine_path)
+
+                engine = create_engine(engine_path)
+
+                # Wrap the engine with error handling
+                self._wrap_sqlite_engine(engine)
+
+                Base.metadata.create_all(bind=engine)
+                self._engines["default"] = engine
+
+            # Create session factory
+            self._session_factories["default"] = sessionmaker(autocommit=False, autoflush=False, bind=self._engines["default"])
+            self._initialized["sync"] = True
+
+    def initialize_async(self, force: bool = False) -> None:
+        """Initialize the asynchronous database engine if not already initialized."""
+        with self._lock:
+            if self._initialized.get("async") and not force:
+                return
+
+            if settings.letta_pg_uri_no_default:
+                self.logger.info("Creating async postgres engine")
+
+                # Create async engine - convert URI to async format
+                pg_uri = settings.letta_pg_uri
+                if pg_uri.startswith("postgresql://"):
+                    async_pg_uri = pg_uri.replace("postgresql://", "postgresql+asyncpg://")
+                else:
+                    async_pg_uri = f"postgresql+asyncpg://{pg_uri.split('://', 1)[1]}" if "://" in pg_uri else pg_uri
+
+                async_engine = create_async_engine(
+                    async_pg_uri,
+                    pool_size=settings.pg_pool_size,
+                    max_overflow=settings.pg_max_overflow,
+                    pool_timeout=settings.pg_pool_timeout,
+                    pool_recycle=settings.pg_pool_recycle,
+                    echo=settings.pg_echo,
+                )
+
+                self._async_engines["default"] = async_engine
+
+                # Create async session factory
+                self._async_session_factories["default"] = async_sessionmaker(
+                    autocommit=False, autoflush=False, bind=self._async_engines["default"], class_=AsyncSession
+                )
+                self._initialized["async"] = True
+            else:
+                self.logger.warning("Async SQLite is currently not supported. Please use PostgreSQL for async database operations.")
+                # TODO (cliandy): unclear around async sqlite support in sqlalchemy, we will not currently support this
+                self._initialized["async"] = False
+
+    def _wrap_sqlite_engine(self, engine: Engine) -> None:
+        """Wrap SQLite engine with error handling."""
+        original_connect = engine.connect
+
+        def wrapped_connect(*args, **kwargs):
+            with db_error_handler():
+                connection = original_connect(*args, **kwargs)
+                original_execute = connection.execute
+
+                def wrapped_execute(*args, **kwargs):
+                    with db_error_handler():
+                        return original_execute(*args, **kwargs)
+
+                connection.execute = wrapped_execute
+                return connection
+
+        engine.connect = wrapped_connect
+
+    def get_engine(self, name: str = "default") -> Engine:
+        """Get a database engine by name."""
+        self.initialize_sync()
+        return self._engines.get(name)
+
+    def get_async_engine(self, name: str = "default") -> AsyncEngine:
+        """Get an async database engine by name."""
+        self.initialize_async()
+        return self._async_engines.get(name)
+
+    def get_session_factory(self, name: str = "default") -> sessionmaker:
+        """Get a session factory by name."""
+        self.initialize_sync()
+        return self._session_factories.get(name)
+
+    def get_async_session_factory(self, name: str = "default") -> async_sessionmaker:
+        """Get an async session factory by name."""
+        self.initialize_async()
+        return self._async_session_factories.get(name)
+
+    @contextmanager
+    def session(self, name: str = "default") -> Generator[Any, None, None]:
+        """Context manager for database sessions."""
+        session_factory = self.get_session_factory(name)
+        if not session_factory:
+            raise ValueError(f"No session factory found for '{name}'")
+
+        session = session_factory()
+        try:
+            yield session
+        finally:
+            session.close()
+
+    @asynccontextmanager
+    async def async_session(self, name: str = "default") -> AsyncGenerator[AsyncSession, None]:
+        """Async context manager for database sessions."""
+        session_factory = self.get_async_session_factory(name)
+        if not session_factory:
+            raise ValueError(f"No async session factory found for '{name}' or async database is not configured")
+
+        session = session_factory()
+        try:
+            yield session
+        finally:
+            await session.close()
+
+
+# Create a singleton instance
+db_registry = DatabaseRegistry()
 
 
 def get_db():
-    """Get a database session, initializing the engine if needed."""
-    global engine, SessionLocal
+    """Get a database session."""
+    with db_registry.session() as session:
+        yield session
 
-    # Make sure engine is initialized
-    if not _engine_initialized:
-        initialize_engine()
 
-    # Now SessionLocal should be defined and callable
-    db = SessionLocal()
-    try:
-        yield db
-    finally:
-        db.close()
+async def get_db_async():
+    """Get an async database session."""
+    async with db_registry.async_session() as session:
+        yield session
 
 
-# Define db_context as a context manager that uses get_db
+# Prefer calling db_registry.session() or db_registry.async_session() directly
+# This is for backwards compatibility
 db_context = contextmanager(get_db)

letta/server/rest_api/app.py

@@ -184,11 +184,14 @@ def create_application() -> "FastAPI":
     @app.exception_handler(Exception)
     async def generic_error_handler(request: Request, exc: Exception):
         # Log the actual error for debugging
-        log.error(f"Unhandled error: {exc}", exc_info=True)
-        print(f"Unhandled error: {exc}")
+        log.error(f"Unhandled error: {str(exc)}", exc_info=True)
+        print(f"Unhandled error: {str(exc)}")
+
+        import traceback
 
         # Print the stack trace
-        print(f"Stack trace: {exc}")
+        print(f"Stack trace: {traceback.format_exc()}")
+
         if (os.getenv("SENTRY_DSN") is not None) and (os.getenv("SENTRY_DSN") != ""):
             import sentry_sdk
 

letta/server/rest_api/chat_completions_interface.py

@@ -162,6 +162,7 @@ class ChatCompletionsStreamingInterface(AgentChunkStreamingInterface):
         expect_reasoning_content: bool = False,
         name: Optional[str] = None,
         message_index: int = 0,
+        prev_message_type: Optional[str] = None,
     ) -> None:
         """
         Called externally with a ChatCompletionChunkResponse. Transforms