chore: add more indexes for block label and agent_id (#2827)

alembic/versions/18ff61fbc034_add_agent_id_index_to_mapping_tables.py

@@ -0,0 +1,39 @@
+"""add agent_id index to mapping tables
+
+Revision ID: 18ff61fbc034
+Revises: b888f21b151f
+Create Date: 2025-09-10 19:16:39.118760
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "18ff61fbc034"
+down_revision: Union[str, None] = "b888f21b151f"
+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_blocks_agents_block_id", "blocks_agents", ["block_id"], unique=False)
+    op.create_index("ix_block_label", "block", ["label"], unique=False)
+    op.create_index("ix_agents_organization_id", "agents", ["organization_id"], unique=False)
+    op.create_index("ix_tools_agents_tool_id", "tools_agents", ["tool_id"], unique=False)
+    op.create_index("ix_sources_agents_source_id", "sources_agents", ["source_id"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_sources_agents_source_id", table_name="sources_agents")
+    op.drop_index("ix_tools_agents_tool_id", table_name="tools_agents")
+    op.drop_index("ix_agents_organization_id", table_name="agents")
+    op.drop_index("ix_block_label", table_name="block")
+    op.drop_index("ix_blocks_agents_block_id", table_name="blocks_agents")
+    # ### end Alembic commands ###

letta/orm/agent.py

@@ -34,7 +34,10 @@ if TYPE_CHECKING:
 class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin, TemplateMixin, AsyncAttrs):
     __tablename__ = "agents"
     __pydantic_model__ = PydanticAgentState
-    __table_args__ = (Index("ix_agents_created_at", "created_at", "id"),)
+    __table_args__ = (
+        Index("ix_agents_created_at", "created_at", "id"),
+        Index("ix_agents_organization_id", "organization_id"),
+    )
 
     # 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/block.py

@@ -24,6 +24,7 @@ class Block(OrganizationMixin, SqlalchemyBase, ProjectMixin, TemplateEntityMixin
     __table_args__ = (
         UniqueConstraint("id", "label", name="unique_block_id_label"),
         Index("created_at_label_idx", "created_at", "label"),
+        Index("ix_block_label", "label"),
     )
 
     template_name: Mapped[Optional[str]] = mapped_column(

letta/orm/blocks_agents.py

@@ -20,6 +20,7 @@ class BlocksAgents(Base):
         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"),
+        Index("ix_blocks_agents_block_id", "block_id"),
     )
 
     # unique agent + block label

letta/orm/sources_agents.py

@@ -1,4 +1,4 @@
-from sqlalchemy import ForeignKey, String
+from sqlalchemy import ForeignKey, Index, String
 from sqlalchemy.orm import Mapped, mapped_column
 
 from letta.orm.base import Base
@@ -8,6 +8,7 @@ class SourcesAgents(Base):
     """Agents can have zero to many sources"""
 
     __tablename__ = "sources_agents"
+    __table_args__ = (Index("ix_sources_agents_source_id", "source_id"),)
 
     agent_id: Mapped[String] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True)
     source_id: Mapped[String] = mapped_column(String, ForeignKey("sources.id", ondelete="CASCADE"), primary_key=True)

letta/orm/tools_agents.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
 
 from letta.orm import Base
@@ -8,7 +8,10 @@ class ToolsAgents(Base):
     """Agents can have one or many tools associated with them."""
 
     __tablename__ = "tools_agents"
-    __table_args__ = (UniqueConstraint("agent_id", "tool_id", name="unique_agent_tool"),)
+    __table_args__ = (
+        UniqueConstraint("agent_id", "tool_id", name="unique_agent_tool"),
+        Index("ix_tools_agents_tool_id", "tool_id"),
+    )
 
     # Each agent must have unique tool names
     agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True)