fix: move all attach detach to be under agents (#723)

Co-authored-by: Mindy Long <[email protected]>

letta/server/rest_api/routers/v1/agents.py

@@ -126,43 +126,63 @@ def get_tools_from_agent(
 ):
     """Get tools from an existing agent"""
     actor = server.user_manager.get_user_or_default(user_id=user_id)
-    return server.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor).tools
+    return server.agent_manager.list_attached_tools(agent_id=agent_id, actor=actor)
 
 
-@router.patch("/{agent_id}/add-tool/{tool_id}", response_model=AgentState, operation_id="add_tool_to_agent")
-def add_tool_to_agent(
+@router.patch("/{agent_id}/tools/attach/{tool_id}", response_model=AgentState, operation_id="attach_tool_to_agent")
+def attach_tool(
     agent_id: str,
     tool_id: str,
     server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    user_id: Optional[str] = Header(None, alias="user_id"),
 ):
-    """Add tools to an existing agent"""
+    """
+    Attach a tool to an agent.
+    """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
     return server.agent_manager.attach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
 
 
-@router.patch("/{agent_id}/remove-tool/{tool_id}", response_model=AgentState, operation_id="remove_tool_from_agent")
-def remove_tool_from_agent(
+@router.patch("/{agent_id}/tools/detach/{tool_id}", response_model=AgentState, operation_id="detach_tool_from_agent")
+def detach_tool(
     agent_id: str,
     tool_id: str,
     server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    user_id: Optional[str] = Header(None, alias="user_id"),
 ):
-    """Add tools to an existing agent"""
+    """
+    Detach a tool from an agent.
+    """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
     return server.agent_manager.detach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
 
 
-@router.patch("/{agent_id}/reset-messages", response_model=AgentState, operation_id="reset_messages")
-def reset_messages(
+@router.patch("/{agent_id}/sources/attach/{source_id}", response_model=AgentState, operation_id="attach_source_to_agent")
+def attach_source(
     agent_id: str,
-    add_default_initial_messages: bool = Query(default=False, description="If true, adds the default initial messages after resetting."),
+    source_id: str,
     server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    user_id: Optional[str] = Header(None, alias="user_id"),
 ):
-    """Resets the messages for an agent"""
+    """
+    Attach a source to an agent.
+    """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
-    return server.agent_manager.reset_messages(agent_id=agent_id, actor=actor, add_default_initial_messages=add_default_initial_messages)
+    return server.agent_manager.attach_source(agent_id=agent_id, source_id=source_id, actor=actor)
+
+
+@router.patch("/{agent_id}/sources/detach/{source_id}", response_model=AgentState, operation_id="detach_source_from_agent")
+def detach_source(
+    agent_id: str,
+    source_id: str,
+    server: "SyncServer" = Depends(get_letta_server),
+    user_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Detach a source from an agent.
+    """
+    actor = server.user_manager.get_user_or_default(user_id=user_id)
+    return server.agent_manager.detach_source(agent_id=agent_id, source_id=source_id, actor=actor)
 
 
 @router.get("/{agent_id}", response_model=AgentState, operation_id="get_agent")
@@ -263,69 +283,54 @@ def list_agent_memory_blocks(
         raise HTTPException(status_code=404, detail=str(e))
 
 
-@router.post("/{agent_id}/core_memory/blocks", response_model=Memory, operation_id="add_agent_memory_block")
-def add_agent_memory_block(
+@router.patch("/{agent_id}/core_memory/blocks/{block_label}", response_model=Block, operation_id="update_agent_memory_block_by_label")
+def update_agent_memory_block(
     agent_id: str,
-    create_block: CreateBlock = Body(...),
+    block_label: str,
+    block_update: BlockUpdate = Body(...),
     server: "SyncServer" = Depends(get_letta_server),
     user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
 ):
     """
-    Creates a memory block and links it to the agent.
+    Updates a memory block of an agent.
     """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
 
-    # Copied from POST /blocks
-    # TODO: Should have block_manager accept only CreateBlock
-    # TODO: This will be possible once we move ID creation to the ORM
-    block_req = Block(**create_block.model_dump())
-    block = server.block_manager.create_or_update_block(actor=actor, block=block_req)
+    block = server.agent_manager.get_block_with_label(agent_id=agent_id, block_label=block_label, actor=actor)
+    block = server.block_manager.update_block(block.id, block_update=block_update, actor=actor)
 
-    # Link the block to the agent
-    agent = server.agent_manager.attach_block(agent_id=agent_id, block_id=block.id, actor=actor)
-    return agent.memory
+    # This should also trigger a system prompt change in the agent
+    server.agent_manager.rebuild_system_prompt(agent_id=agent_id, actor=actor, force=True, update_timestamp=False)
+
+    return block
 
 
-@router.delete("/{agent_id}/core_memory/blocks/{block_label}", response_model=Memory, operation_id="remove_agent_memory_block_by_label")
-def remove_agent_memory_block(
+@router.patch("/{agent_id}/core_memory/blocks/attach/{block_id}", response_model=AgentState, operation_id="attach_block_to_agent")
+def attach_block(
     agent_id: str,
-    # TODO should this be block_id, or the label?
-    # I think label is OK since it's user-friendly + guaranteed to be unique within a Memory object
-    block_label: str,
+    block_id: str,
     server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    user_id: Optional[str] = Header(None, alias="user_id"),
 ):
     """
-    Removes a memory block from an agent by unlnking it. If the block is not linked to any other agent, it is deleted.
+    Attach a block to an agent.
     """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
+    return server.agent_manager.attach_block(agent_id=agent_id, block_id=block_id, actor=actor)
 
-    # Unlink the block from the agent
-    agent = server.agent_manager.detach_block_with_label(agent_id=agent_id, block_label=block_label, actor=actor)
 
-    return agent.memory
-
-
-@router.patch("/{agent_id}/core_memory/blocks/{block_label}", response_model=Block, operation_id="update_agent_memory_block_by_label")
-def update_agent_memory_block(
+@router.patch("/{agent_id}/core_memory/blocks/detach/{block_id}", response_model=AgentState, operation_id="detach_block_from_agent")
+def detach_block(
     agent_id: str,
-    block_label: str,
-    block_update: BlockUpdate = Body(...),
+    block_id: str,
     server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+    user_id: Optional[str] = Header(None, alias="user_id"),
 ):
     """
-    Removes a memory block from an agent by unlnking it. If the block is not linked to any other agent, it is deleted.
+    Detach a block from an agent.
     """
     actor = server.user_manager.get_user_or_default(user_id=user_id)
-
-    block = server.agent_manager.get_block_with_label(agent_id=agent_id, block_label=block_label, actor=actor)
-    block = server.block_manager.update_block(block.id, block_update=block_update, actor=actor)
-
-    # This should also trigger a system prompt change in the agent
-    server.agent_manager.rebuild_system_prompt(agent_id=agent_id, actor=actor, force=True, update_timestamp=False)
-
-    return block
+    return server.agent_manager.detach_block(agent_id=agent_id, block_id=block_id, actor=actor)
 
 
 @router.get("/{agent_id}/archival_memory", response_model=List[Passage], operation_id="list_agent_archival_memory")
@@ -610,3 +615,15 @@ async def send_message_async(
     )
 
     return run
+
+
+@router.patch("/{agent_id}/reset-messages", response_model=AgentState, operation_id="reset_messages")
+def reset_messages(
+    agent_id: str,
+    add_default_initial_messages: bool = Query(default=False, description="If true, adds the default initial messages after resetting."),
+    server: "SyncServer" = Depends(get_letta_server),
+    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+):
+    """Resets the messages for an agent"""
+    actor = server.user_manager.get_user_or_default(user_id=user_id)
+    return server.agent_manager.reset_messages(agent_id=agent_id, actor=actor, add_default_initial_messages=add_default_initial_messages)

letta/server/rest_api/routers/v1/blocks.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, List, Optional
 
-from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query, Response
+from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
 
 from letta.orm.errors import NoResultFound
 from letta.schemas.block import Block, BlockUpdate, CreateBlock
@@ -73,41 +73,3 @@ def get_block(
         return block
     except NoResultFound:
         raise HTTPException(status_code=404, detail="Block not found")
-
-
-@router.patch("/{block_id}/attach", response_model=None, status_code=204, operation_id="link_agent_memory_block")
-def link_agent_memory_block(
-    block_id: str,
-    agent_id: str = Query(..., description="The unique identifier of the agent to attach the source to."),
-    server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
-):
-    """
-    Link a memory block to an agent.
-    """
-    actor = server.user_manager.get_user_or_default(user_id=user_id)
-
-    try:
-        server.agent_manager.attach_block(agent_id=agent_id, block_id=block_id, actor=actor)
-        return Response(status_code=204)
-    except NoResultFound as e:
-        raise HTTPException(status_code=404, detail=str(e))
-
-
-@router.patch("/{block_id}/detach", response_model=None, status_code=204, operation_id="unlink_agent_memory_block")
-def unlink_agent_memory_block(
-    block_id: str,
-    agent_id: str = Query(..., description="The unique identifier of the agent to attach the source to."),
-    server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
-):
-    """
-    Unlink a memory block from an agent
-    """
-    actor = server.user_manager.get_user_or_default(user_id=user_id)
-
-    try:
-        server.agent_manager.detach_block(agent_id=agent_id, block_id=block_id, actor=actor)
-        return Response(status_code=204)
-    except NoResultFound as e:
-        raise HTTPException(status_code=404, detail=str(e))

letta/server/rest_api/routers/v1/sources.py

@@ -111,36 +111,6 @@ def delete_source(
     server.delete_source(source_id=source_id, actor=actor)
 
 
-@router.post("/{source_id}/attach", response_model=Source, operation_id="attach_agent_to_source")
-def attach_source_to_agent(
-    source_id: str,
-    agent_id: str = Query(..., description="The unique identifier of the agent to attach the source to."),
-    server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
-):
-    """
-    Attach a data source to an existing agent.
-    """
-    actor = server.user_manager.get_user_or_default(user_id=user_id)
-    server.agent_manager.attach_source(source_id=source_id, agent_id=agent_id, actor=actor)
-    return server.source_manager.get_source_by_id(source_id=source_id, actor=actor)
-
-
-@router.post("/{source_id}/detach", response_model=Source, operation_id="detach_agent_from_source")
-def detach_source_from_agent(
-    source_id: str,
-    agent_id: str = Query(..., description="The unique identifier of the agent to detach the source from."),
-    server: "SyncServer" = Depends(get_letta_server),
-    user_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
-) -> None:
-    """
-    Detach a data source from an existing agent.
-    """
-    actor = server.user_manager.get_user_or_default(user_id=user_id)
-    server.agent_manager.detach_source(agent_id=agent_id, source_id=source_id, actor=actor)
-    return server.source_manager.get_source_by_id(source_id=source_id, actor=actor)
-
-
 @router.post("/{source_id}/upload", response_model=Job, operation_id="upload_file_to_source")
 def upload_file_to_source(
     file: UploadFile,

letta/services/agent_manager.py

@@ -25,6 +25,7 @@
 from letta.schemas.message import MessageCreate
 from letta.schemas.passage import Passage as PydanticPassage
 from letta.schemas.source import Source as PydanticSource
+from letta.schemas.tool import Tool as PydanticTool
 from letta.schemas.tool_rule import ToolRule as PydanticToolRule
 from letta.schemas.user import User as PydanticUser
 from letta.services.block_manager import BlockManager
@@ -537,7 +538,7 @@ def reset_messages(self, agent_id: str, actor: PydanticUser, add_default_initial
     # Source Management
     # ======================================================================================================================
     @enforce_types
-    def attach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> None:
+    def attach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> PydanticAgentState:
         """
         Attaches a source to an agent.
 
@@ -567,6 +568,7 @@ def attach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> N
 
             # Commit the changes
             agent.update(session, actor=actor)
+            return agent.to_pydantic()
 
     @enforce_types
     def list_attached_sources(self, agent_id: str, actor: PydanticUser) -> List[PydanticSource]:
@@ -588,7 +590,7 @@ def list_attached_sources(self, agent_id: str, actor: PydanticUser) -> List[Pyda
             return [source.to_pydantic() for source in agent.sources]
 
     @enforce_types
-    def detach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> None:
+    def detach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> PydanticAgentState:
         """
         Detaches a source from an agent.
 
@@ -602,10 +604,17 @@ def detach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> N
             agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
 
             # Remove the source from the relationship
-            agent.sources = [s for s in agent.sources if s.id != source_id]
+            remaining_sources = [s for s in agent.sources if s.id != source_id]
+
+            if len(remaining_sources) == len(agent.sources):  # Source ID was not in the relationship
+                logger.warning(f"Attempted to remove unattached source id={source_id} from agent id={agent_id} by actor={actor}")
+
+            # Update the sources relationship
+            agent.sources = remaining_sources
 
             # Commit the changes
             agent.update(session, actor=actor)
+            return agent.to_pydantic()
 
     # ======================================================================================================================
     # Block management
@@ -1011,6 +1020,22 @@ def detach_tool(self, agent_id: str, tool_id: str, actor: PydanticUser) -> Pydan
             agent.update(session, actor=actor)
             return agent.to_pydantic()
 
+    @enforce_types
+    def list_attached_tools(self, agent_id: str, actor: PydanticUser) -> List[PydanticTool]:
+        """
+        List all tools attached to an agent.
+
+        Args:
+            agent_id: ID of the agent to list tools for.
+            actor: User performing the action.
+
+        Returns:
+            List[PydanticTool]: List of tools attached to the agent.
+        """
+        with self.session_maker() as session:
+            agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+            return [tool.to_pydantic() for tool in agent.tools]
+
     # ======================================================================================================================
     # Tag Management
     # ======================================================================================================================