diff --git a/letta/orm/agent.py b/letta/orm/agent.py index e6b14987..781ab383 100644 --- a/letta/orm/agent.py +++ b/letta/orm/agent.py @@ -70,7 +70,14 @@ class Agent(SqlalchemyBase, OrganizationMixin): ) tools: Mapped[List["Tool"]] = relationship("Tool", secondary="tools_agents", lazy="selectin", passive_deletes=True) sources: Mapped[List["Source"]] = relationship("Source", secondary="sources_agents", lazy="selectin") - core_memory: Mapped[List["Block"]] = relationship("Block", secondary="blocks_agents", lazy="selectin") + core_memory: Mapped[List["Block"]] = relationship( + "Block", + secondary="blocks_agents", + lazy="selectin", + passive_deletes=True, # Ensures SQLAlchemy doesn't fetch blocks_agents rows before deleting + back_populates="agents", + doc="Blocks forming the core memory of the agent.", + ) messages: Mapped[List["Message"]] = relationship( "Message", back_populates="agent", diff --git a/letta/orm/block.py b/letta/orm/block.py index ecfa6fe1..3e8c8006 100644 --- a/letta/orm/block.py +++ b/letta/orm/block.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional, Type +from typing import TYPE_CHECKING, List, Optional, Type from sqlalchemy import JSON, BigInteger, Index, Integer, UniqueConstraint, event from sqlalchemy.orm import Mapped, attributes, mapped_column, relationship @@ -39,6 +39,14 @@ class Block(OrganizationMixin, SqlalchemyBase): # relationships organization: Mapped[Optional["Organization"]] = relationship("Organization") + agents: Mapped[List["Agent"]] = relationship( + "Agent", + secondary="blocks_agents", + lazy="selectin", + passive_deletes=True, # Ensures SQLAlchemy doesn't fetch blocks_agents rows before deleting + back_populates="core_memory", + doc="Agents associated with this block.", + ) def to_pydantic(self) -> Type: match self.label: diff --git a/tests/test_managers.py b/tests/test_managers.py index 28ad057a..fe6280d5 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -1155,6 +1155,10 @@ def test_detach_block(server: SyncServer, sarah_agent, default_block, default_us agent = server.agent_manager.get_agent_by_id(sarah_agent.id, actor=default_user) assert len(agent.memory.blocks) == 0 + # Check that block still exists + block = server.block_manager.get_block_by_id(block_id=default_block.id, actor=default_user) + assert block + def test_detach_nonexistent_block(server: SyncServer, sarah_agent, default_user): """Test detaching a block that isn't attached.""" @@ -2080,6 +2084,26 @@ def test_delete_block(server: SyncServer, default_user): assert len(blocks) == 0 +def test_delete_block_detaches_from_agent(server: SyncServer, sarah_agent, default_user): + # Create and delete a block + block = server.block_manager.create_or_update_block(PydanticBlock(label="human", value="Sample content"), actor=default_user) + agent_state = server.agent_manager.attach_block(agent_id=sarah_agent.id, block_id=block.id, actor=default_user) + + # Check that block has been attached + assert block.id in [b.id for b in agent_state.memory.blocks] + + # Now attempt to delete the block + server.block_manager.delete_block(block_id=block.id, actor=default_user) + + # Verify that the block was deleted + blocks = server.block_manager.get_blocks(actor=default_user) + assert len(blocks) == 0 + + # Check that block has been detached too + agent_state = server.agent_manager.get_agent_by_id(agent_id=sarah_agent.id, actor=default_user) + assert not (block.id in [b.id for b in agent_state.memory.blocks]) + + # ====================================================================================================================== # SourceManager Tests - Sources # ======================================================================================================================