feat(asyncify): agent blocks routes (#2363)
This commit is contained in:
@@ -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")
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user