feat: add new modify approvals api (#4288)
* feat: add new modify approvals api * remove path params override
This commit is contained in:
@@ -481,6 +481,25 @@ async def detach_tool(
|
||||
return await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
|
||||
|
||||
|
||||
@router.patch("/{agent_id}/tools/approval/{tool_name}", response_model=AgentState, operation_id="modify_approval")
|
||||
async def modify_approval(
|
||||
agent_id: str,
|
||||
tool_name: str,
|
||||
requires_approval: bool,
|
||||
server: "SyncServer" = Depends(get_letta_server),
|
||||
actor_id: str | None = Header(None, alias="user_id"),
|
||||
):
|
||||
"""
|
||||
Attach a tool to an agent.
|
||||
"""
|
||||
actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id)
|
||||
await server.agent_manager.toggle_approvals_async(
|
||||
agent_id=agent_id, tool_name=tool_name, requires_approval=requires_approval, actor=actor
|
||||
)
|
||||
# TODO: Unfortunately we need this to preserve our current API behavior
|
||||
return await server.agent_manager.get_agent_by_id_async(agent_id=agent_id, actor=actor)
|
||||
|
||||
|
||||
@router.patch("/{agent_id}/sources/attach/{source_id}", response_model=AgentState, operation_id="attach_source_to_agent")
|
||||
async def attach_source(
|
||||
agent_id: str,
|
||||
|
||||
@@ -3170,6 +3170,33 @@ class AgentManager:
|
||||
|
||||
await session.commit()
|
||||
|
||||
@enforce_types
|
||||
@trace_method
|
||||
async def modify_approvals_async(self, agent_id: str, tool_name: str, requires_approval: bool, actor: PydanticUser) -> None:
|
||||
def is_target_rule(rule):
|
||||
return rule.tool_name == tool_name and rule.type == "requires_approval"
|
||||
|
||||
async with db_registry.async_session() as session:
|
||||
agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
|
||||
existing_rules = [rule for rule in agent.tool_rules if is_target_rule(rule)]
|
||||
|
||||
if len(existing_rules) == 1 and not requires_approval:
|
||||
tool_rules = [rule for rule in agent.tool_rules if not is_target_rule(rule)]
|
||||
elif len(existing_rules) == 0 and requires_approval:
|
||||
# Create a new list to ensure SQLAlchemy detects the change
|
||||
# This is critical for JSON columns - modifying in place doesn't trigger change detection
|
||||
tool_rules = list(agent.tool_rules) if agent.tool_rules else []
|
||||
tool_rules.append(RequiresApprovalToolRule(tool_name=tool_name))
|
||||
else:
|
||||
tool_rules = None
|
||||
|
||||
if tool_rules is None:
|
||||
return
|
||||
|
||||
agent.tool_rules = tool_rules
|
||||
session.add(agent)
|
||||
await session.commit()
|
||||
|
||||
@enforce_types
|
||||
@trace_method
|
||||
def list_attached_tools(self, agent_id: str, actor: PydanticUser) -> List[PydanticTool]:
|
||||
|
||||
@@ -1967,6 +1967,25 @@ async def test_attach_tool_with_default_requires_approval_on_creation(server: Sy
|
||||
assert len(tool_rules) == 1
|
||||
assert tool_rules[0].type == "requires_approval"
|
||||
|
||||
# Modify approval on tool after attach
|
||||
await server.agent_manager.modify_approvals_async(
|
||||
agent_id=agent.id, tool_name=bash_tool.name, requires_approval=False, actor=default_user
|
||||
)
|
||||
agent = await server.agent_manager.get_agent_by_id_async(agent_id=agent.id, actor=default_user)
|
||||
assert len([t for t in agent.tools if t.id == bash_tool.id]) == 1
|
||||
tool_rules = [rule for rule in agent.tool_rules if rule.tool_name == bash_tool.name]
|
||||
assert len(tool_rules) == 0
|
||||
|
||||
# Revert override
|
||||
await server.agent_manager.modify_approvals_async(
|
||||
agent_id=agent.id, tool_name=bash_tool.name, requires_approval=True, actor=default_user
|
||||
)
|
||||
agent = await server.agent_manager.get_agent_by_id_async(agent_id=agent.id, actor=default_user)
|
||||
assert len([t for t in agent.tools if t.id == bash_tool.id]) == 1
|
||||
tool_rules = [rule for rule in agent.tool_rules if rule.tool_name == bash_tool.name]
|
||||
assert len(tool_rules) == 1
|
||||
assert tool_rules[0].type == "requires_approval"
|
||||
|
||||
|
||||
# ======================================================================================================================
|
||||
# AgentManager Tests - Sources Relationship
|
||||
|
||||
Reference in New Issue
Block a user