feat: remove agents SELECTIN from blocks model (#3427)

This commit is contained in:
cthomas
2025-07-19 23:04:33 -07:00
committed by GitHub
parent 7a37557bb9
commit 4f37a39741
3 changed files with 26 additions and 10 deletions

View File

@@ -59,7 +59,7 @@ class Block(OrganizationMixin, SqlalchemyBase):
agents: Mapped[List["Agent"]] = relationship(
"Agent",
secondary="blocks_agents",
lazy="selectin",
lazy="raise",
passive_deletes=True, # Ensures SQLAlchemy doesn't fetch blocks_agents rows before deleting
back_populates="core_memory",
doc="Agents associated with this block.",

View File

@@ -72,14 +72,14 @@ async def modify_block(
return await server.block_manager.update_block_async(block_id=block_id, block_update=block_update, actor=actor)
@router.delete("/{block_id}", response_model=Block, operation_id="delete_block")
@router.delete("/{block_id}", operation_id="delete_block")
async def delete_block(
block_id: str,
server: SyncServer = Depends(get_letta_server),
actor_id: Optional[str] = Header(None, alias="user_id"),
):
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
return await server.block_manager.delete_block_async(block_id=block_id, actor=actor)
await server.block_manager.delete_block_async(block_id=block_id, actor=actor)
@router.get("/{block_id}", response_model=Block, operation_id="retrieve_block")

View File

@@ -1,12 +1,14 @@
import asyncio
from typing import Dict, List, Optional
from sqlalchemy import select
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from letta.log import get_logger
from letta.orm.agent import Agent as AgentModel
from letta.orm.block import Block as BlockModel
from letta.orm.block_history import BlockHistory
from letta.orm.blocks_agents import BlocksAgents
from letta.orm.errors import NoResultFound
from letta.otel.tracing import trace_method
from letta.schemas.agent import AgentState as PydanticAgentState
@@ -135,21 +137,29 @@ class BlockManager:
@enforce_types
@trace_method
def delete_block(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
def delete_block(self, block_id: str, actor: PydanticUser) -> None:
"""Delete a block by its ID."""
with db_registry.session() as session:
# First, delete all references in blocks_agents table
session.execute(delete(BlocksAgents).where(BlocksAgents.block_id == block_id))
session.flush()
# Then delete the block itself
block = BlockModel.read(db_session=session, identifier=block_id)
block.hard_delete(db_session=session, actor=actor)
return block.to_pydantic()
@enforce_types
@trace_method
async def delete_block_async(self, block_id: str, actor: PydanticUser) -> PydanticBlock:
async def delete_block_async(self, block_id: str, actor: PydanticUser) -> None:
"""Delete a block by its ID."""
async with db_registry.async_session() as session:
# First, delete all references in blocks_agents table
await session.execute(delete(BlocksAgents).where(BlocksAgents.block_id == block_id))
await session.flush()
# Then delete the block itself
block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
await block.hard_delete_async(db_session=session, actor=actor)
return block.to_pydantic()
@enforce_types
@trace_method
@@ -296,8 +306,14 @@ class BlockManager:
Retrieve all agents associated with a given block.
"""
async with db_registry.async_session() as session:
block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
agents_orm = block.agents
query = (
select(AgentModel)
.where(AgentModel.id.in_(select(BlocksAgents.agent_id).where(BlocksAgents.block_id == block_id)))
.where(AgentModel.organization_id == actor.organization_id)
)
result = await session.execute(query)
agents_orm = result.scalars().all()
agents = await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents_orm])
return agents