feat: Add FK cascades, remove app-side label propagation [LET-4690] (#5219)

add FK cascades, remove app-side label propagation
This commit is contained in:
Matthew Zhou
2025-10-07 13:26:05 -07:00
committed by Caren Thomas
parent ade992d5ca
commit 932b418681
3 changed files with 64 additions and 18 deletions

View File

@@ -0,0 +1,55 @@
"""add cascades to blocks_agents FKs; set initially immediate
Revision ID: 038e68cdf0df
Revises: b6061da886ee
Create Date: 2025-10-07 13:01:17.872405
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "038e68cdf0df"
down_revision: Union[str, None] = "b6061da886ee"
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.drop_constraint(op.f("blocks_agents_agent_id_fkey"), "blocks_agents", type_="foreignkey")
op.drop_constraint(op.f("fk_block_id_label"), "blocks_agents", type_="foreignkey")
op.create_foreign_key(
"fk_block_id_label",
"blocks_agents",
"block",
["block_id", "block_label"],
["id", "label"],
onupdate="CASCADE",
ondelete="CASCADE",
initially="IMMEDIATE",
deferrable=True,
)
op.create_foreign_key(None, "blocks_agents", "agents", ["agent_id"], ["id"], ondelete="CASCADE")
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(None, "blocks_agents", type_="foreignkey")
op.drop_constraint("fk_block_id_label", "blocks_agents", type_="foreignkey")
op.create_foreign_key(
op.f("fk_block_id_label"),
"blocks_agents",
"block",
["block_id", "block_label"],
["id", "label"],
initially="DEFERRED",
deferrable=True,
)
op.create_foreign_key(op.f("blocks_agents_agent_id_fkey"), "blocks_agents", "agents", ["agent_id"], ["id"])
# ### end Alembic commands ###

View File

@@ -1,7 +1,7 @@
from typing import TYPE_CHECKING, List, Optional, Type
from sqlalchemy import JSON, BigInteger, ForeignKey, Index, Integer, String, UniqueConstraint, event
from sqlalchemy.orm import Mapped, attributes, declared_attr, mapped_column, relationship
from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
from letta.constants import CORE_MEMORY_BLOCK_CHAR_LIMIT
from letta.orm.block_history import BlockHistory
@@ -110,21 +110,6 @@ class Block(OrganizationMixin, SqlalchemyBase, ProjectMixin, TemplateEntityMixin
) # Helps manage potential FK cycles
@event.listens_for(Block, "after_update") # Changed from 'before_update'
def block_before_update(mapper, connection, target):
"""Handle updating BlocksAgents when a block's label changes."""
label_history = attributes.get_history(target, "label")
if not label_history.has_changes():
return
blocks_agents = BlocksAgents.__table__
connection.execute(
blocks_agents.update()
.where(blocks_agents.c.block_id == target.id, blocks_agents.c.block_label == label_history.deleted[0])
.values(block_label=label_history.added[0])
)
@event.listens_for(Block, "before_insert")
@event.listens_for(Block, "before_update")
def validate_value_length(mapper, connection, target):

View File

@@ -15,7 +15,13 @@ class BlocksAgents(Base):
name="unique_label_per_agent",
),
ForeignKeyConstraint(
["block_id", "block_label"], ["block.id", "block.label"], name="fk_block_id_label", deferrable=True, initially="DEFERRED"
["block_id", "block_label"],
["block.id", "block.label"],
name="fk_block_id_label",
onupdate="CASCADE",
ondelete="CASCADE",
deferrable=True,
initially="IMMEDIATE",
),
UniqueConstraint("agent_id", "block_id", name="unique_agent_block"),
Index("ix_blocks_agents_block_label_agent_id", "block_label", "agent_id"),
@@ -24,6 +30,6 @@ class BlocksAgents(Base):
)
# unique agent + block label
agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id"), primary_key=True)
agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id", ondelete="CASCADE"), primary_key=True)
block_id: Mapped[str] = mapped_column(String, primary_key=True)
block_label: Mapped[str] = mapped_column(String, primary_key=True)