chore: add index to blocks and agents related tables (#2127)

alembic/versions/1dc0fee72dea_add_block_related_indexes.py

@@ -0,0 +1,34 @@
+"""add block-related indexes
+
+Revision ID: 1dc0fee72dea
+Revises: 18e300709530
+Create Date: 2025-05-12 17:06:32.055091
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "1dc0fee72dea"
+down_revision: Union[str, None] = "18e300709530"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade():
+    # add index for blocks_agents table
+    op.create_index("ix_blocks_agents_block_label_agent_id", "blocks_agents", ["block_label", "agent_id"], unique=False)
+
+    # add index for just block_label
+    op.create_index("ix_blocks_block_label", "blocks_agents", ["block_label"], unique=False)
+
+    # add index for agent_tags for agent_id and tag
+    op.create_index("ix_agents_tags_agent_id_tag", "agents_tags", ["agent_id", "tag"], unique=False)
+
+
+def downgrade():
+    op.drop_index("ix_blocks_agents_block_label_agent_id", table_name="blocks_agents")
+    op.drop_index("ix_blocks_block_label", table_name="blocks_agents")
+    op.drop_index("ix_agents_tags_agent_id_tag", table_name="agents_tags")

letta/orm/agents_tags.py

@@ -1,4 +1,4 @@
-from sqlalchemy import ForeignKey, String, UniqueConstraint
+from sqlalchemy import ForeignKey, Index, String, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.orm.base import Base
@@ -6,7 +6,10 @@ from letta.orm.base import Base
 
 class AgentsTags(Base):
     __tablename__ = "agents_tags"
-    __table_args__ = (UniqueConstraint("agent_id", "tag", name="unique_agent_tag"),)
+    __table_args__ = (
+        UniqueConstraint("agent_id", "tag", name="unique_agent_tag"),
+        Index("ix_agents_tags_agent_id_tag", "agent_id", "tag"),
+    )
 
     # # agent generates its own id
     # # TODO: We want to migrate all the ORM models to do this, so we will need to move this to the SqlalchemyBase

letta/orm/blocks_agents.py

@@ -1,4 +1,4 @@
-from sqlalchemy import ForeignKey, ForeignKeyConstraint, String, UniqueConstraint
+from sqlalchemy import ForeignKey, ForeignKeyConstraint, Index, String, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column
 
 from letta.orm.base import Base
@@ -18,6 +18,8 @@ class BlocksAgents(Base):
             ["block_id", "block_label"], ["block.id", "block.label"], name="fk_block_id_label", deferrable=True, initially="DEFERRED"
         ),
         UniqueConstraint("agent_id", "block_id", name="unique_agent_block"),
+        Index("ix_blocks_agents_block_label_agent_id", "block_label", "agent_id"),
+        Index("ix_blocks_block_label", "block_label"),
     )
 
     # unique agent + block label