feat(asyncify): agent blocks routes (#2363)

This commit is contained in:
cthomas
2025-05-23 00:42:44 -07:00
committed by GitHub
parent 7467bd5b14
commit 7499a09030
3 changed files with 120 additions and 16 deletions

View File

@@ -356,7 +356,7 @@ async def retrieve_agent(
@router.delete("/{agent_id}", response_model=None, operation_id="delete_agent")
def delete_agent(
async def delete_agent(
agent_id: str,
server: "SyncServer" = Depends(get_letta_server),
actor_id: Optional[str] = Header(None, alias="user_id"), # Extract user_id from header, default to None if not present
@@ -364,9 +364,9 @@ def delete_agent(
"""
Delete an agent.
"""
actor = server.user_manager.get_user_or_default(user_id=actor_id)
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
try:
server.agent_manager.delete_agent(agent_id=agent_id, actor=actor)
await server.agent_manager.delete_agent_async(agent_id=agent_id, actor=actor)
return JSONResponse(status_code=status.HTTP_200_OK, content={"message": f"Agent id={agent_id} successfully deleted"})
except NoResultFound:
raise HTTPException(status_code=404, detail=f"Agent agent_id={agent_id} not found for user_id={actor.id}.")
@@ -387,7 +387,7 @@ async def list_agent_sources(
# TODO: remove? can also get with agent blocks
@router.get("/{agent_id}/core-memory", response_model=Memory, operation_id="retrieve_agent_memory")
def retrieve_agent_memory(
async def retrieve_agent_memory(
agent_id: str,
server: "SyncServer" = Depends(get_letta_server),
actor_id: Optional[str] = Header(None, alias="user_id"), # Extract user_id from header, default to None if not present
@@ -396,13 +396,13 @@ def retrieve_agent_memory(
Retrieve the memory state of a specific agent.
This endpoint fetches the current memory state of the agent identified by the user ID and agent ID.
"""
actor = server.user_manager.get_user_or_default(user_id=actor_id)
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
return server.get_agent_memory(agent_id=agent_id, actor=actor)
return await server.get_agent_memory_async(agent_id=agent_id, actor=actor)
@router.get("/{agent_id}/core-memory/blocks/{block_label}", response_model=Block, operation_id="retrieve_core_memory_block")
def retrieve_block(
async def retrieve_block(
agent_id: str,
block_label: str,
server: "SyncServer" = Depends(get_letta_server),
@@ -411,10 +411,10 @@ def retrieve_block(
"""
Retrieve a core memory block from an agent.
"""
actor = server.user_manager.get_user_or_default(user_id=actor_id)
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
try:
return server.agent_manager.get_block_with_label(agent_id=agent_id, block_label=block_label, actor=actor)
return await server.agent_manager.get_block_with_label_async(agent_id=agent_id, block_label=block_label, actor=actor)
except NoResultFound as e:
raise HTTPException(status_code=404, detail=str(e))
@@ -454,13 +454,13 @@ async def modify_block(
)
# 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)
await server.agent_manager.rebuild_system_prompt_async(agent_id=agent_id, actor=actor, force=True, update_timestamp=False)
return block
@router.patch("/{agent_id}/core-memory/blocks/attach/{block_id}", response_model=AgentState, operation_id="attach_core_memory_block")
def attach_block(
async def attach_block(
agent_id: str,
block_id: str,
server: "SyncServer" = Depends(get_letta_server),
@@ -469,12 +469,12 @@ def attach_block(
"""
Attach a core memoryblock to an agent.
"""
actor = server.user_manager.get_user_or_default(user_id=actor_id)
return server.agent_manager.attach_block(agent_id=agent_id, block_id=block_id, actor=actor)
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
return await server.agent_manager.attach_block_async(agent_id=agent_id, block_id=block_id, actor=actor)
@router.patch("/{agent_id}/core-memory/blocks/detach/{block_id}", response_model=AgentState, operation_id="detach_core_memory_block")
def detach_block(
async def detach_block(
agent_id: str,
block_id: str,
server: "SyncServer" = Depends(get_letta_server),
@@ -483,8 +483,8 @@ def detach_block(
"""
Detach a core memory block from an agent.
"""
actor = server.user_manager.get_user_or_default(user_id=actor_id)
return server.agent_manager.detach_block(agent_id=agent_id, block_id=block_id, actor=actor)
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
return await server.agent_manager.detach_block_async(agent_id=agent_id, block_id=block_id, actor=actor)
@router.get("/{agent_id}/archival-memory", response_model=List[Passage], operation_id="list_passages")

View File

@@ -969,6 +969,11 @@ class SyncServer(Server):
"""Return the memory of an agent (core memory)"""
return self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor).memory
async def get_agent_memory_async(self, agent_id: str, actor: User) -> Memory:
"""Return the memory of an agent (core memory)"""
agent = await self.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
return agent.memory
def get_archival_memory_summary(self, agent_id: str, actor: User) -> ArchivalMemorySummary:
return ArchivalMemorySummary(size=self.agent_manager.passage_size(actor=actor, agent_id=agent_id))

View File

@@ -1078,6 +1078,56 @@ class AgentManager:
else:
logger.debug(f"Agent with ID {agent_id} successfully hard deleted")
@trace_method
@enforce_types
async def delete_agent_async(self, agent_id: str, actor: PydanticUser) -> None:
"""
Deletes an agent and its associated relationships.
Ensures proper permission checks and cascades where applicable.
Args:
agent_id: ID of the agent to be deleted.
actor: User performing the action.
Raises:
NoResultFound: If agent doesn't exist
"""
async with db_registry.async_session() as session:
# Retrieve the agent
logger.debug(f"Hard deleting Agent with ID: {agent_id} with actor={actor}")
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
agents_to_delete = [agent]
sleeptime_group_to_delete = None
# Delete sleeptime agent and group (TODO this is flimsy pls fix)
if agent.multi_agent_group:
participant_agent_ids = agent.multi_agent_group.agent_ids
if agent.multi_agent_group.manager_type in {ManagerType.sleeptime, ManagerType.voice_sleeptime} and participant_agent_ids:
for participant_agent_id in participant_agent_ids:
try:
sleeptime_agent = await AgentModel.read_async(db_session=session, identifier=participant_agent_id, actor=actor)
agents_to_delete.append(sleeptime_agent)
except NoResultFound:
pass # agent already deleted
sleeptime_agent_group = await GroupModel.read_async(
db_session=session, identifier=agent.multi_agent_group.id, actor=actor
)
sleeptime_group_to_delete = sleeptime_agent_group
try:
if sleeptime_group_to_delete is not None:
await session.delete(sleeptime_group_to_delete)
await session.commit()
for agent in agents_to_delete:
await session.delete(agent)
await session.commit()
except Exception as e:
await session.rollback()
logger.exception(f"Failed to hard delete Agent with ID {agent_id}")
raise ValueError(f"Failed to hard delete Agent with ID {agent_id}: {e}")
else:
logger.debug(f"Agent with ID {agent_id} successfully hard deleted")
@trace_method
@enforce_types
def serialize(self, agent_id: str, actor: PydanticUser) -> AgentSchema:
@@ -1678,6 +1728,22 @@ class AgentManager:
return block.to_pydantic()
raise NoResultFound(f"No block with label '{block_label}' found for agent '{agent_id}'")
@trace_method
@enforce_types
async def get_block_with_label_async(
self,
agent_id: str,
block_label: str,
actor: PydanticUser,
) -> PydanticBlock:
"""Gets a block attached to an agent by its label."""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
for block in agent.core_memory:
if block.label == block_label:
return block.to_pydantic()
raise NoResultFound(f"No block with label '{block_label}' found for agent '{agent_id}'")
@trace_method
@enforce_types
async def modify_block_by_label_async(
@@ -1743,6 +1809,18 @@ class AgentManager:
agent.update(session, actor=actor)
return agent.to_pydantic()
@trace_method
@enforce_types
async def attach_block_async(self, agent_id: str, block_id: str, actor: PydanticUser) -> PydanticAgentState:
"""Attaches a block to an agent."""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
agent.core_memory.append(block)
await agent.update_async(session, actor=actor)
return await agent.to_pydantic_async()
@trace_method
@enforce_types
def detach_block(
@@ -1764,6 +1842,27 @@ class AgentManager:
agent.update(session, actor=actor)
return agent.to_pydantic()
@trace_method
@enforce_types
async def detach_block_async(
self,
agent_id: str,
block_id: str,
actor: PydanticUser,
) -> PydanticAgentState:
"""Detaches a block from an agent."""
async with db_registry.async_session() as session:
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
original_length = len(agent.core_memory)
agent.core_memory = [b for b in agent.core_memory if b.id != block_id]
if len(agent.core_memory) == original_length:
raise NoResultFound(f"No block with id '{block_id}' found for agent '{agent_id}' with actor id: '{actor.id}'")
await agent.update_async(session, actor=actor)
return await agent.to_pydantic_async()
@trace_method
@enforce_types
def detach_block_with_label(