feat: Make blocks agents mapping table

alembic/versions/1c8880d671ee_make_an_blocks_agents_mapping_table.py

@@ -0,0 +1,52 @@
+"""Make an blocks agents mapping table
+
+Revision ID: 1c8880d671ee
+Revises: f81ceea2c08d
+Create Date: 2024-11-22 15:42:47.209229
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "1c8880d671ee"
+down_revision: Union[str, None] = "f81ceea2c08d"
+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_unique_constraint("unique_block_id_label", "block", ["id", "label"])
+
+    op.create_table(
+        "blocks_agents",
+        sa.Column("agent_id", sa.String(), nullable=False),
+        sa.Column("block_id", sa.String(), nullable=False),
+        sa.Column("block_label", sa.String(), nullable=False),
+        sa.Column("id", sa.String(), nullable=False),
+        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True),
+        sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True),
+        sa.Column("is_deleted", sa.Boolean(), server_default=sa.text("FALSE"), nullable=False),
+        sa.Column("_created_by_id", sa.String(), nullable=True),
+        sa.Column("_last_updated_by_id", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["agent_id"],
+            ["agents.id"],
+        ),
+        sa.ForeignKeyConstraint(["block_id", "block_label"], ["block.id", "block.label"], name="fk_block_id_label"),
+        sa.PrimaryKeyConstraint("agent_id", "block_id", "block_label", "id"),
+        sa.UniqueConstraint("agent_id", "block_label", name="unique_label_per_agent"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint("unique_block_id_label", "block", type_="unique")
+    op.drop_table("blocks_agents")
+    # ### end Alembic commands ###

letta/orm/__init__.py

@@ -1,5 +1,6 @@
 from letta.orm.base import Base
 from letta.orm.block import Block
+from letta.orm.blocks_agents import BlocksAgents
 from letta.orm.file import FileMetadata
 from letta.orm.organization import Organization
 from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable

letta/orm/block.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Optional, Type
 
-from sqlalchemy import JSON, BigInteger, Integer
+from sqlalchemy import JSON, BigInteger, Integer, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from letta.constants import CORE_MEMORY_BLOCK_CHAR_LIMIT
@@ -18,6 +18,8 @@ class Block(OrganizationMixin, SqlalchemyBase):
 
     __tablename__ = "block"
     __pydantic_model__ = PydanticBlock
+    # This may seem redundant, but is necessary for the BlocksAgents composite FK relationship
+    __table_args__ = (UniqueConstraint("id", "label", name="unique_block_id_label"),)
 
     template_name: Mapped[Optional[str]] = mapped_column(
         nullable=True, doc="the unique name that identifies a block in a human-readable way"

letta/orm/blocks_agents.py

@@ -0,0 +1,29 @@
+from sqlalchemy import ForeignKey, ForeignKeyConstraint, String, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column
+
+from letta.orm.sqlalchemy_base import SqlalchemyBase
+from letta.schemas.blocks_agents import BlocksAgents as PydanticBlocksAgents
+
+
+class BlocksAgents(SqlalchemyBase):
+    """Agents must have one or many blocks to make up their core memory."""
+
+    __tablename__ = "blocks_agents"
+    __pydantic_model__ = PydanticBlocksAgents
+    __table_args__ = (
+        UniqueConstraint(
+            "agent_id",
+            "block_label",
+            name="unique_label_per_agent",
+        ),
+        ForeignKeyConstraint(
+            ["block_id", "block_label"],
+            ["block.id", "block.label"],
+            name="fk_block_id_label",
+        ),
+    )
+
+    # unique agent + block label
+    agent_id: Mapped[str] = mapped_column(String, ForeignKey("agents.id"), primary_key=True)
+    block_id: Mapped[str] = mapped_column(String, primary_key=True)
+    block_label: Mapped[str] = mapped_column(String, primary_key=True)

letta/schemas/blocks_agents.py

@@ -0,0 +1,32 @@
+from datetime import datetime
+from typing import Optional
+
+from pydantic import Field
+
+from letta.schemas.letta_base import LettaBase
+
+
+class BlocksAgentsBase(LettaBase):
+    __id_prefix__ = "blocks_agents"
+
+
+class BlocksAgents(BlocksAgentsBase):
+    """
+    Schema representing the relationship between blocks and agents.
+
+    Parameters:
+        agent_id (str): The ID of the associated agent.
+        block_id (str): The ID of the associated block.
+        block_label (str): The label of the block.
+        created_at (datetime): The date this relationship was created.
+        updated_at (datetime): The date this relationship was last updated.
+        is_deleted (bool): Whether this block-agent relationship is deleted or not.
+    """
+
+    id: str = BlocksAgentsBase.generate_id_field()
+    agent_id: str = Field(..., description="The ID of the associated agent.")
+    block_id: str = Field(..., description="The ID of the associated block.")
+    block_label: str = Field(..., description="The label of the block.")
+    created_at: Optional[datetime] = Field(None, description="The creation date of the association.")
+    updated_at: Optional[datetime] = Field(None, description="The update date of the association.")
+    is_deleted: bool = Field(False, description="Whether this block-agent relationship is deleted or not.")

letta/server/server.py

@@ -77,6 +77,7 @@
 from letta.schemas.user import User
 from letta.services.agents_tags_manager import AgentsTagsManager
 from letta.services.block_manager import BlockManager
+from letta.services.blocks_agents_manager import BlocksAgentsManager
 from letta.services.organization_manager import OrganizationManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
 from letta.services.source_manager import SourceManager
@@ -248,6 +249,7 @@ def __init__(
         self.block_manager = BlockManager()
         self.source_manager = SourceManager()
         self.agents_tags_manager = AgentsTagsManager()
+        self.blocks_agents_manager = BlocksAgentsManager()
         self.sandbox_config_manager = SandboxConfigManager(tool_settings)
 
         # Make default user and org

letta/services/block_manager.py

@@ -36,7 +36,7 @@ def create_or_update_block(self, block: Block, actor: PydanticUser) -> PydanticB
             return block.to_pydantic()
 
     @enforce_types
-    def update_block(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser, limit: Optional[int] = None) -> PydanticBlock:
+    def update_block(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser) -> PydanticBlock:
         """Update a block by its ID with the given BlockUpdate object."""
         with self.session_maker() as session:
             block = BlockModel.read(db_session=session, identifier=block_id, actor=actor)

letta/services/blocks_agents_manager.py

@@ -0,0 +1,84 @@
+import warnings
+from typing import List
+
+from letta.orm.blocks_agents import BlocksAgents as BlocksAgentsModel
+from letta.orm.errors import NoResultFound
+from letta.schemas.blocks_agents import BlocksAgents as PydanticBlocksAgents
+from letta.utils import enforce_types
+
+
+# TODO: DELETE THIS ASAP
+# TODO: So we have a patch where we manually specify CRUD operations
+# TODO: This is because Agent is NOT migrated to the ORM yet
+# TODO: Once we migrate Agent to the ORM, we should deprecate any agents relationship table managers
+class BlocksAgentsManager:
+    """Manager class to handle business logic related to Blocks and Agents."""
+
+    def __init__(self):
+        from letta.server.server import db_context
+
+        self.session_maker = db_context
+
+    @enforce_types
+    def add_block_to_agent(self, agent_id: str, block_id: str, block_label: str) -> PydanticBlocksAgents:
+        """Add a block to an agent. If the label already exists on that agent, this will error."""
+        with self.session_maker() as session:
+            try:
+                # Check if the block-label combination already exists for this agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                warnings.warn(f"Block label '{block_label}' already exists for agent '{agent_id}'.")
+            except NoResultFound:
+                blocks_agents_record = PydanticBlocksAgents(agent_id=agent_id, block_id=block_id, block_label=block_label)
+                blocks_agents_record = BlocksAgentsModel(**blocks_agents_record.model_dump(exclude_none=True))
+                blocks_agents_record.create(session)
+
+            return blocks_agents_record.to_pydantic()
+
+    @enforce_types
+    def remove_block_with_label_from_agent(self, agent_id: str, block_label: str) -> PydanticBlocksAgents:
+        """Remove a block with a label from an agent."""
+        with self.session_maker() as session:
+            try:
+                # Find and delete the block-label association for the agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                blocks_agents_record.hard_delete(session)
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block label '{block_label}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def remove_block_with_id_from_agent(self, agent_id: str, block_id: str) -> PydanticBlocksAgents:
+        """Remove a block with a label from an agent."""
+        with self.session_maker() as session:
+            try:
+                # Find and delete the block-label association for the agent
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_id=block_id)
+                blocks_agents_record.hard_delete(session)
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block id '{block_id}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def update_block_id_for_agent(self, agent_id: str, block_label: str, new_block_id: str) -> PydanticBlocksAgents:
+        """Update the block ID for a specific block label for an agent."""
+        with self.session_maker() as session:
+            try:
+                blocks_agents_record = BlocksAgentsModel.read(db_session=session, agent_id=agent_id, block_label=block_label)
+                blocks_agents_record.block_id = new_block_id
+                return blocks_agents_record.to_pydantic()
+            except NoResultFound:
+                raise ValueError(f"Block label '{block_label}' not found for agent '{agent_id}'.")
+
+    @enforce_types
+    def list_block_ids_for_agent(self, agent_id: str) -> List[str]:
+        """List all blocks associated with a specific agent."""
+        with self.session_maker() as session:
+            blocks_agents_record = BlocksAgentsModel.list(db_session=session, agent_id=agent_id)
+            return [record.block_id for record in blocks_agents_record]
+
+    @enforce_types
+    def list_agent_ids_with_block(self, block_id: str) -> List[str]:
+        """List all agents associated with a specific block."""
+        with self.session_maker() as session:
+            blocks_agents_record = BlocksAgentsModel.list(db_session=session, block_id=block_id)
+            return [record.agent_id for record in blocks_agents_record]