From 4e1efa5c91a5307c265ffbb3b15931f0a8a8bfe9 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Mon, 12 May 2025 18:36:59 -0700
Subject: [PATCH] chore: add index to blocks and agents related tables  (#2127)

---
 .../1dc0fee72dea_add_block_related_indexes.py | 34 +++++++++++++++++++
 letta/orm/agents_tags.py                      |  7 ++--
 letta/orm/blocks_agents.py                    |  4 ++-
 3 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 alembic/versions/1dc0fee72dea_add_block_related_indexes.py

diff --git a/alembic/versions/1dc0fee72dea_add_block_related_indexes.py b/alembic/versions/1dc0fee72dea_add_block_related_indexes.py
new file mode 100644
index 00000000..89c7d739
--- /dev/null
+++ b/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")
diff --git a/letta/orm/agents_tags.py b/letta/orm/agents_tags.py
index 76ff9011..438910a3 100644
--- a/letta/orm/agents_tags.py
+++ b/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
diff --git a/letta/orm/blocks_agents.py b/letta/orm/blocks_agents.py
index 4774783b..15f2d015 100644
--- a/letta/orm/blocks_agents.py
+++ b/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