letta-nightly 0.8.15.dev20250720104313__py3-none-any.whl → 0.8.16.dev20250721104533__py3-none-any.whl

letta/orm/message.py

@@ -11,7 +11,7 @@ from letta.schemas.letta_message_content import MessageContent
 from letta.schemas.letta_message_content import TextContent as PydanticTextContent
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.message import ToolReturn
-from letta.settings import settings
+from letta.settings import DatabaseChoice, settings
 
 
 class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
@@ -49,6 +49,9 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
         nullable=True,
         doc="The id of the LLMBatchItem that this message is associated with",
     )
+    is_err: Mapped[Optional[bool]] = mapped_column(
+        nullable=True, doc="Whether this message is part of an error step. Used only for debugging purposes."
+    )
 
     # Monotonically increasing sequence for efficient/correct listing
     sequence_id: Mapped[int] = mapped_column(
@@ -59,7 +62,7 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
     )
 
     # Relationships
-    organization: Mapped["Organization"] = relationship("Organization", back_populates="messages", lazy="selectin")
+    organization: Mapped["Organization"] = relationship("Organization", back_populates="messages", lazy="raise")
     step: Mapped["Step"] = relationship("Step", back_populates="messages", lazy="selectin")
 
     # Job relationship
@@ -78,7 +81,7 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
         if self.text and not model.content:
             model.content = [PydanticTextContent(text=self.text)]
         # If there are no tool calls, set tool_calls to None
-        if len(self.tool_calls) == 0:
+        if self.tool_calls is None or len(self.tool_calls) == 0:
             model.tool_calls = None
         return model
 
@@ -86,16 +89,139 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
 # listener
 
 
-@event.listens_for(Message, "before_insert")
-def set_sequence_id_for_sqlite(mapper, connection, target):
-    # TODO: Kind of hacky, used to detect if we are using sqlite or not
-    if not settings.letta_pg_uri_no_default:
-        session = Session.object_session(target)
+@event.listens_for(Session, "before_flush")
+def set_sequence_id_for_sqlite_bulk(session, flush_context, instances):
+    # Handle bulk inserts for SQLite
+    if settings.database_engine is DatabaseChoice.SQLITE:
+        # Find all new Message objects that need sequence IDs
+        new_messages = [obj for obj in session.new if isinstance(obj, Message) and obj.sequence_id is None]
+
+        if new_messages:
+            # Create a sequence table if it doesn't exist for atomic increments
+            session.execute(
+                text(
+                    """
+                CREATE TABLE IF NOT EXISTS message_sequence (
+                    id INTEGER PRIMARY KEY,
+                    next_val INTEGER NOT NULL DEFAULT 1
+                )
+            """
+                )
+            )
+
+            # Initialize the sequence table if empty
+            session.execute(
+                text(
+                    """
+                INSERT OR IGNORE INTO message_sequence (id, next_val)
+                SELECT 1, COALESCE(MAX(sequence_id), 0) + 1
+                FROM messages
+            """
+                )
+            )
+
+            # Get the number of records being inserted
+            records_count = len(new_messages)
+
+            # Atomically reserve a range of sequence values for this batch
+            result = session.execute(
+                text(
+                    """
+                UPDATE message_sequence
+                SET next_val = next_val + :count
+                WHERE id = 1
+                RETURNING next_val - :count
+            """
+                ),
+                {"count": records_count},
+            )
+
+            start_sequence_id = result.scalar()
+            if start_sequence_id is None:
+                # Fallback if RETURNING doesn't work (older SQLite versions)
+                session.execute(
+                    text(
+                        """
+                    UPDATE message_sequence
+                    SET next_val = next_val + :count
+                    WHERE id = 1
+                """
+                    ),
+                    {"count": records_count},
+                )
+                start_sequence_id = session.execute(
+                    text(
+                        """
+                    SELECT next_val - :count FROM message_sequence WHERE id = 1
+                """
+                    ),
+                    {"count": records_count},
+                ).scalar()
+
+            # Assign sequential IDs to each record
+            for i, obj in enumerate(new_messages):
+                obj.sequence_id = start_sequence_id + i
 
-        if not hasattr(session, "_sequence_id_counter"):
-            # Initialize counter for this flush
-            max_seq = connection.scalar(text("SELECT MAX(sequence_id) FROM messages"))
-            session._sequence_id_counter = max_seq or 0
 
-        session._sequence_id_counter += 1
-        target.sequence_id = session._sequence_id_counter
+@event.listens_for(Message, "before_insert")
+def set_sequence_id_for_sqlite(mapper, connection, target):
+    if settings.database_engine is DatabaseChoice.SQLITE:
+        # For SQLite, we need to generate sequence_id manually
+        # Use a database-level atomic operation to avoid race conditions
+
+        # Create a sequence table if it doesn't exist for atomic increments
+        connection.execute(
+            text(
+                """
+            CREATE TABLE IF NOT EXISTS message_sequence (
+                id INTEGER PRIMARY KEY,
+                next_val INTEGER NOT NULL DEFAULT 1
+            )
+        """
+            )
+        )
+
+        # Initialize the sequence table if empty
+        connection.execute(
+            text(
+                """
+            INSERT OR IGNORE INTO message_sequence (id, next_val)
+            SELECT 1, COALESCE(MAX(sequence_id), 0) + 1
+            FROM messages
+        """
+            )
+        )
+
+        # Atomically get the next sequence value
+        result = connection.execute(
+            text(
+                """
+            UPDATE message_sequence
+            SET next_val = next_val + 1
+            WHERE id = 1
+            RETURNING next_val - 1
+        """
+            )
+        )
+
+        sequence_id = result.scalar()
+        if sequence_id is None:
+            # Fallback if RETURNING doesn't work (older SQLite versions)
+            connection.execute(
+                text(
+                    """
+                UPDATE message_sequence
+                SET next_val = next_val + 1
+                WHERE id = 1
+            """
+                )
+            )
+            sequence_id = connection.execute(
+                text(
+                    """
+                SELECT next_val - 1 FROM message_sequence WHERE id = 1
+            """
+                )
+            ).scalar()
+
+        target.sequence_id = sequence_id

letta/orm/organization.py

@@ -6,18 +6,17 @@ from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.organization import Organization as PydanticOrganization
 
 if TYPE_CHECKING:
+    from letta.orm import Source
     from letta.orm.agent import Agent
-    from letta.orm.agent_passage import AgentPassage
     from letta.orm.block import Block
     from letta.orm.group import Group
     from letta.orm.identity import Identity
-    from letta.orm.llm_batch_item import LLMBatchItem
+    from letta.orm.llm_batch_items import LLMBatchItem
     from letta.orm.llm_batch_job import LLMBatchJob
     from letta.orm.message import Message
+    from letta.orm.passage import AgentPassage, SourcePassage
     from letta.orm.provider import Provider
-    from letta.orm.sandbox_config import AgentEnvironmentVariable, SandboxConfig
-    from letta.orm.sandbox_environment_variable import SandboxEnvironmentVariable
-    from letta.orm.source_passage import SourcePassage
+    from letta.orm.sandbox_config import AgentEnvironmentVariable, SandboxConfig, SandboxEnvironmentVariable
     from letta.orm.tool import Tool
     from letta.orm.user import User
 
@@ -48,6 +47,7 @@ class Organization(SqlalchemyBase):
 
     # relationships
     agents: Mapped[List["Agent"]] = relationship("Agent", back_populates="organization", cascade="all, delete-orphan")
+    sources: Mapped[List["Source"]] = relationship("Source", cascade="all, delete-orphan")
     messages: Mapped[List["Message"]] = relationship("Message", back_populates="organization", cascade="all, delete-orphan")
     source_passages: Mapped[List["SourcePassage"]] = relationship(
         "SourcePassage", back_populates="organization", cascade="all, delete-orphan"

letta/orm/passage.py

@@ -9,7 +9,7 @@ from letta.orm.custom_columns import CommonVector, EmbeddingConfigColumn
 from letta.orm.mixins import AgentMixin, FileMixin, OrganizationMixin, SourceMixin
 from letta.orm.sqlalchemy_base import SqlalchemyBase
 from letta.schemas.passage import Passage as PydanticPassage
-from letta.settings import settings
+from letta.settings import DatabaseChoice, settings
 
 config = LettaConfig()
 
@@ -29,7 +29,7 @@ class BasePassage(SqlalchemyBase, OrganizationMixin):
     metadata_: Mapped[dict] = mapped_column(JSON, doc="Additional metadata")
 
     # Vector embedding field based on database type
-    if settings.letta_pg_uri_no_default:
+    if settings.database_engine is DatabaseChoice.POSTGRES:
         from pgvector.sqlalchemy import Vector
 
         embedding = mapped_column(Vector(MAX_EMBEDDING_DIM))
@@ -56,7 +56,7 @@ class SourcePassage(BasePassage, FileMixin, SourceMixin):
     @declared_attr
     def __table_args__(cls):
         # TODO (cliandy): investigate if this is necessary, may be for SQLite compatability or do we need to add as well?
-        if settings.letta_pg_uri_no_default:
+        if settings.database_engine is DatabaseChoice.POSTGRES:
             return (
                 Index("source_passages_org_idx", "organization_id"),
                 Index("source_passages_created_at_id_idx", "created_at", "id"),
@@ -81,7 +81,7 @@ class AgentPassage(BasePassage, AgentMixin):
 
     @declared_attr
     def __table_args__(cls):
-        if settings.letta_pg_uri_no_default:
+        if settings.database_engine is DatabaseChoice.POSTGRES:
             return (
                 Index("agent_passages_org_idx", "organization_id"),
                 Index("ix_agent_passages_org_agent", "organization_id", "agent_id"),

letta/orm/source.py

@@ -20,7 +20,7 @@ class Source(SqlalchemyBase, OrganizationMixin):
     __pydantic_model__ = PydanticSource
 
     __table_args__ = (
-        Index(f"source_created_at_id_idx", "created_at", "id"),
+        Index("source_created_at_id_idx", "created_at", "id"),
         UniqueConstraint("name", "organization_id", name="uq_source_name_organization"),
         {"extend_existing": True},
     )

letta/orm/sqlalchemy_base.py

@@ -5,7 +5,7 @@ from functools import wraps
 from pprint import pformat
 from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Union
 
-from sqlalchemy import Sequence, String, and_, delete, func, or_, select, text
+from sqlalchemy import Sequence, String, and_, delete, func, or_, select
 from sqlalchemy.exc import DBAPIError, IntegrityError, TimeoutError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Mapped, Session, mapped_column
@@ -15,6 +15,7 @@ from letta.log import get_logger
 from letta.orm.base import Base, CommonSqlalchemyMetaMixins
 from letta.orm.errors import DatabaseTimeoutError, ForeignKeyConstraintViolationError, NoResultFound, UniqueConstraintViolationError
 from letta.orm.sqlite_functions import adapt_array
+from letta.settings import DatabaseChoice
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
@@ -353,10 +354,12 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
 
             if before_obj and after_obj:
                 # Window-based query - get records between before and after
-                conditions = [
-                    or_(cls.created_at < before_obj.created_at, and_(cls.created_at == before_obj.created_at, cls.id < before_obj.id)),
-                    or_(cls.created_at > after_obj.created_at, and_(cls.created_at == after_obj.created_at, cls.id > after_obj.id)),
-                ]
+                conditions.append(
+                    or_(cls.created_at < before_obj.created_at, and_(cls.created_at == before_obj.created_at, cls.id < before_obj.id))
+                )
+                conditions.append(
+                    or_(cls.created_at > after_obj.created_at, and_(cls.created_at == after_obj.created_at, cls.id > after_obj.id))
+                )
             else:
                 # Pure pagination query
                 if before_obj:
@@ -393,7 +396,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
 
             from letta.settings import settings
 
-            if settings.letta_pg_uri_no_default:
+            if settings.database_engine is DatabaseChoice.POSTGRES:
                 # PostgreSQL with pgvector
                 query = query.order_by(cls.embedding.cosine_distance(query_embedding).asc())
             else:
@@ -509,14 +512,9 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         query, query_conditions = cls._read_multiple_preprocess(identifiers, actor, access, access_type, check_is_deleted, **kwargs)
         if query is None:
             raise NoResultFound(f"{cls.__name__} not found with identifier {identifier}")
-        if is_postgresql_session(db_session):
-            await db_session.execute(text("SET LOCAL enable_seqscan = OFF"))
-        try:
-            result = await db_session.execute(query)
-            item = result.scalar_one_or_none()
-        finally:
-            if is_postgresql_session(db_session):
-                await db_session.execute(text("SET LOCAL enable_seqscan = ON"))
+
+        result = await db_session.execute(query)
+        item = result.scalar_one_or_none()
 
         if item is None:
             raise NoResultFound(f"{cls.__name__} not found with {', '.join(query_conditions if query_conditions else ['no conditions'])}")
@@ -656,7 +654,13 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             self._handle_dbapi_error(e)
 
     @handle_db_timeout
-    async def create_async(self, db_session: "AsyncSession", actor: Optional["User"] = None, no_commit: bool = False) -> "SqlalchemyBase":
+    async def create_async(
+        self,
+        db_session: "AsyncSession",
+        actor: Optional["User"] = None,
+        no_commit: bool = False,
+        no_refresh: bool = False,
+    ) -> "SqlalchemyBase":
         """Async version of create function"""
         logger.debug(f"Creating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
 
@@ -668,7 +672,9 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
                 await db_session.flush()  # no commit, just flush to get PK
             else:
                 await db_session.commit()
-            await db_session.refresh(self)
+
+            if not no_refresh:
+                await db_session.refresh(self)
             return self
         except (DBAPIError, IntegrityError) as e:
             self._handle_dbapi_error(e)
@@ -717,7 +723,12 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
     @classmethod
     @handle_db_timeout
     async def batch_create_async(
-        cls, items: List["SqlalchemyBase"], db_session: "AsyncSession", actor: Optional["User"] = None
+        cls,
+        items: List["SqlalchemyBase"],
+        db_session: "AsyncSession",
+        actor: Optional["User"] = None,
+        no_commit: bool = False,
+        no_refresh: bool = False,
     ) -> List["SqlalchemyBase"]:
         """
         Async version of batch_create method.
@@ -726,10 +737,13 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             items: List of model instances to create
             db_session: AsyncSession session
             actor: Optional user performing the action
+            no_commit: Whether to commit the transaction
+            no_refresh: Whether to refresh the created objects
         Returns:
             List of created model instances
         """
         logger.debug(f"Async batch creating {len(items)} {cls.__name__} items with actor={actor}")
+
         if not items:
             return []
 
@@ -740,21 +754,22 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
 
         try:
             db_session.add_all(items)
-            await db_session.flush()  # Flush to generate IDs but don't commit yet
-
-            # Collect IDs to fetch the complete objects after commit
-            item_ids = [item.id for item in items]
-
-            await db_session.commit()
-
-            # Re-query the objects to get them with relationships loaded
-            query = select(cls).where(cls.id.in_(item_ids))
-            if hasattr(cls, "created_at"):
-                query = query.order_by(cls.created_at)
+            if no_commit:
+                await db_session.flush()
+            else:
+                await db_session.commit()
 
-            result = await db_session.execute(query)
-            return list(result.scalars())
+            if no_refresh:
+                return items
+            else:
+                # Re-query the objects to get them with relationships loaded
+                item_ids = [item.id for item in items]
+                query = select(cls).where(cls.id.in_(item_ids))
+                if hasattr(cls, "created_at"):
+                    query = query.order_by(cls.created_at)
 
+                result = await db_session.execute(query)
+                return list(result.scalars())
         except (DBAPIError, IntegrityError) as e:
             cls._handle_dbapi_error(e)
 
@@ -854,20 +869,27 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         return self
 
     @handle_db_timeout
-    async def update_async(self, db_session: AsyncSession, actor: "User | None" = None, no_commit: bool = False) -> "SqlalchemyBase":
+    async def update_async(
+        self, db_session: "AsyncSession", actor: Optional["User"] = None, no_commit: bool = False, no_refresh: bool = False
+    ) -> "SqlalchemyBase":
         """Async version of update function"""
-        logger.debug(...)
+        logger.debug(f"Updating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
+
         if actor:
             self._set_created_and_updated_by_fields(actor.id)
         self.set_updated_at()
+        try:
+            db_session.add(self)
+            if no_commit:
+                await db_session.flush()
+            else:
+                await db_session.commit()
 
-        db_session.add(self)
-        if no_commit:
-            await db_session.flush()
-        else:
-            await db_session.commit()
-        await db_session.refresh(self)
-        return self
+            if not no_refresh:
+                await db_session.refresh(self)
+            return self
+        except (DBAPIError, IntegrityError) as e:
+            self._handle_dbapi_error(e)
 
     @classmethod
     def _size_preprocess(

letta/orm/step.py

@@ -5,6 +5,7 @@ from sqlalchemy import JSON, ForeignKey, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.sqlalchemy_base import SqlalchemyBase
+from letta.schemas.letta_stop_reason import StopReasonType
 from letta.schemas.step import Step as PydanticStep
 
 if TYPE_CHECKING:
@@ -45,6 +46,7 @@ class Step(SqlalchemyBase):
     prompt_tokens: Mapped[int] = mapped_column(default=0, doc="Number of tokens in the prompt")
     total_tokens: Mapped[int] = mapped_column(default=0, doc="Total number of tokens processed by the agent")
     completion_tokens_details: Mapped[Optional[Dict]] = mapped_column(JSON, nullable=True, doc="metadata for the agent.")
+    stop_reason: Mapped[Optional[StopReasonType]] = mapped_column(None, nullable=True, doc="The stop reason associated with this step.")
     tags: Mapped[Optional[List]] = mapped_column(JSON, doc="Metadata tags.")
     tid: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="Transaction ID that processed the step.")
     trace_id: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The trace id of the agent step.")