diff --git a/alembic/versions/18ff61fbc034_add_agent_id_index_to_mapping_tables.py b/alembic/versions/18ff61fbc034_add_agent_id_index_to_mapping_tables.py new file mode 100644 index 00000000..825ead37 --- /dev/null +++ b/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 ### diff --git a/letta/orm/agent.py b/letta/orm/agent.py index 5c50017f..1a14adec 100644 --- a/letta/orm/agent.py +++ b/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 diff --git a/letta/orm/block.py b/letta/orm/block.py index 4fe3c78b..bec65f4e 100644 --- a/letta/orm/block.py +++ b/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( diff --git a/letta/orm/blocks_agents.py b/letta/orm/blocks_agents.py index 15f2d015..918363d8 100644 --- a/letta/orm/blocks_agents.py +++ b/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 diff --git a/letta/orm/sources_agents.py b/letta/orm/sources_agents.py index ffe8a9d0..b22c61c6 100644 --- a/letta/orm/sources_agents.py +++ b/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) diff --git a/letta/orm/tools_agents.py b/letta/orm/tools_agents.py index 52c1e0a1..ffc07c09 100644 --- a/letta/orm/tools_agents.py +++ b/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)