diff --git a/alembic/versions/348214cbc081_add_org_agent_id_indices.py b/alembic/versions/348214cbc081_add_org_agent_id_indices.py new file mode 100644 index 00000000..9bf7575b --- /dev/null +++ b/alembic/versions/348214cbc081_add_org_agent_id_indices.py @@ -0,0 +1,31 @@ +"""add org agent id indices + +Revision ID: 348214cbc081 +Revises: dd049fbec729 +Create Date: 2025-05-28 22:43:18.509397 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "348214cbc081" +down_revision: Union[str, None] = "dd049fbec729" +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_agent_passages_org_agent", "agent_passages", ["organization_id", "agent_id"], unique=False) + op.create_index("ix_messages_org_agent", "messages", ["organization_id", "agent_id"], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index("ix_messages_org_agent", table_name="messages") + op.drop_index("ix_agent_passages_org_agent", table_name="agent_passages") + # ### end Alembic commands ### diff --git a/letta/orm/message.py b/letta/orm/message.py index 0495da20..41551fb5 100644 --- a/letta/orm/message.py +++ b/letta/orm/message.py @@ -22,6 +22,7 @@ class Message(SqlalchemyBase, OrganizationMixin, AgentMixin): Index("ix_messages_agent_created_at", "agent_id", "created_at"), Index("ix_messages_created_at", "created_at", "id"), Index("ix_messages_agent_sequence", "agent_id", "sequence_id"), + Index("ix_messages_org_agent", "organization_id", "agent_id"), ) __pydantic_model__ = PydanticMessage diff --git a/letta/orm/passage.py b/letta/orm/passage.py index 4bec1717..76790fc2 100644 --- a/letta/orm/passage.py +++ b/letta/orm/passage.py @@ -41,16 +41,6 @@ class BasePassage(SqlalchemyBase, OrganizationMixin): """Relationship to organization""" return relationship("Organization", back_populates="passages", lazy="selectin") - @declared_attr - def __table_args__(cls): - if settings.letta_pg_uri_no_default: - return ( - Index(f"{cls.__tablename__}_org_idx", "organization_id"), - Index(f"{cls.__tablename__}_created_at_id_idx", "created_at", "id"), - {"extend_existing": True}, - ) - return (Index(f"{cls.__tablename__}_created_at_id_idx", "created_at", "id"), {"extend_existing": True}) - class SourcePassage(BasePassage, FileMixin, SourceMixin): """Passages derived from external files/sources""" @@ -66,6 +56,16 @@ class SourcePassage(BasePassage, FileMixin, SourceMixin): def organization(cls) -> Mapped["Organization"]: return relationship("Organization", back_populates="source_passages", lazy="selectin") + @declared_attr + def __table_args__(cls): + if settings.letta_pg_uri_no_default: + return ( + Index("source_passages_org_idx", "organization_id"), + Index("source_passages_created_at_id_idx", "created_at", "id"), + {"extend_existing": True}, + ) + return (Index("source_passages_created_at_id_idx", "created_at", "id"), {"extend_existing": True}) + @declared_attr def source(cls) -> Mapped["Source"]: """Relationship to source""" @@ -80,3 +80,14 @@ class AgentPassage(BasePassage, AgentMixin): @declared_attr def organization(cls) -> Mapped["Organization"]: return relationship("Organization", back_populates="agent_passages", lazy="selectin") + + @declared_attr + def __table_args__(cls): + if settings.letta_pg_uri_no_default: + return ( + Index("agent_passages_org_idx", "organization_id"), + Index("ix_agent_passages_org_agent", "organization_id", "agent_id"), + Index("agent_passages_created_at_id_idx", "created_at", "id"), + {"extend_existing": True}, + ) + return (Index("agent_passages_created_at_id_idx", "created_at", "id"), {"extend_existing": True}) diff --git a/letta/orm/sqlalchemy_base.py b/letta/orm/sqlalchemy_base.py index 5412c146..28bc5578 100644 --- a/letta/orm/sqlalchemy_base.py +++ b/letta/orm/sqlalchemy_base.py @@ -848,9 +848,9 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base): else: # Single value for equality filtering query = query.where(column == value) - # Handle soft deletes if the class has the 'is_deleted' attribute - if hasattr(cls, "is_deleted"): - query = query.where(cls.is_deleted == False) + # TODO: Handle soft deletes if the class has the 'is_deleted' attribute + # if hasattr(cls, "is_deleted"): + # query = query.where(cls.is_deleted == False) return query