fix: Add index to messages table (#2280)

alembic/versions/d6632deac81d_add_composite_index_to_messages_table.py

@@ -0,0 +1,29 @@
+"""Add composite index to messages table
+
+Revision ID: d6632deac81d
+Revises: 54dec07619c4
+Create Date: 2024-12-18 13:38:56.511701
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "d6632deac81d"
+down_revision: Union[str, None] = "54dec07619c4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_index("ix_messages_agent_created_at", "messages", ["agent_id", "created_at"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_messages_agent_created_at", table_name="messages")
+    # ### end Alembic commands ###

letta/orm/message.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from sqlalchemy import Index
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.custom_columns import ToolCallColumn
@@ -13,7 +14,7 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
     """Defines data model for storing Message objects"""
 
     __tablename__ = "messages"
-    __table_args__ = {"extend_existing": True}
+    __table_args__ = (Index("ix_messages_agent_created_at", "agent_id", "created_at"),)
     __pydantic_model__ = PydanticMessage
 
     id: Mapped[str] = mapped_column(primary_key=True, doc="Unique message identifier")

letta/orm/passage.py

@@ -1,26 +1,26 @@
 from typing import TYPE_CHECKING
-from sqlalchemy import Column, JSON, Index
-from sqlalchemy.orm import Mapped, mapped_column, relationship, declared_attr
 
-from letta.orm.mixins import FileMixin, OrganizationMixin
+from sqlalchemy import JSON, Column, Index
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
+
+from letta.config import LettaConfig
+from letta.constants import MAX_EMBEDDING_DIM
 from letta.orm.custom_columns import CommonVector, EmbeddingConfigColumn
-from letta.orm.sqlalchemy_base import SqlalchemyBase
 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.config import LettaConfig
-from letta.constants import MAX_EMBEDDING_DIM
-
 config = LettaConfig()
 
 if TYPE_CHECKING:
-    from letta.orm.organization import Organization
     from letta.orm.agent import Agent
+    from letta.orm.organization import Organization
 
 
 class BasePassage(SqlalchemyBase, OrganizationMixin):
     """Base class for all passage types with common fields"""
+
     __abstract__ = True
     __pydantic_model__ = PydanticPassage
 
@@ -45,17 +45,15 @@ def organization(cls) -> Mapped["Organization"]:
     @declared_attr
     def __table_args__(cls):
         if settings.letta_pg_uri_no_default:
-            return (
-                Index(f'{cls.__tablename__}_org_idx', 'organization_id'),
-                {"extend_existing": True}
-            )
+            return (Index(f"{cls.__tablename__}_org_idx", "organization_id"), {"extend_existing": True})
         return ({"extend_existing": True},)
 
 
 class SourcePassage(BasePassage, FileMixin, SourceMixin):
     """Passages derived from external files/sources"""
+
     __tablename__ = "source_passages"
-    
+
     @declared_attr
     def file(cls) -> Mapped["FileMetadata"]:
         """Relationship to file"""
@@ -64,7 +62,7 @@ def file(cls) -> Mapped["FileMetadata"]:
     @declared_attr
     def organization(cls) -> Mapped["Organization"]:
         return relationship("Organization", back_populates="source_passages", lazy="selectin")
-        
+
     @declared_attr
     def source(cls) -> Mapped["Source"]:
         """Relationship to source"""
@@ -73,8 +71,9 @@ def source(cls) -> Mapped["Source"]:
 
 class AgentPassage(BasePassage, AgentMixin):
     """Passages created by agents as archival memories"""
+
     __tablename__ = "agent_passages"
-    
+
     @declared_attr
     def organization(cls) -> Mapped["Organization"]:
         return relationship("Organization", back_populates="agent_passages", lazy="selectin")