diff --git a/alembic/versions/d6632deac81d_add_composite_index_to_messages_table.py b/alembic/versions/d6632deac81d_add_composite_index_to_messages_table.py new file mode 100644 index 0000000000..68ceec88f5 --- /dev/null +++ b/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 ### diff --git a/letta/orm/message.py b/letta/orm/message.py index e1d6da7889..a8bbb90017 100644 --- a/letta/orm/message.py +++ b/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") diff --git a/letta/orm/passage.py b/letta/orm/passage.py index d38878418b..492c602174 100644 --- a/letta/orm/passage.py +++ b/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")