feat: remove agents SELECTIN from blocks model (#3427)
This commit is contained in:
@@ -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.",
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user