feat: optimize size async queries (#2512)

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 ###

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
 

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})

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
+        # TODO: Handle soft deletes if the class has the 'is_deleted' attribute
-        if hasattr(cls, "is_deleted"):
+        # if hasattr(cls, "is_deleted"):
-            query = query.where(cls.is_deleted == False)
+        #    query = query.where(cls.is_deleted == False)
 
         return query