diff --git a/examples/fastapi-vite/backend/agent.py b/examples/fastapi-vite/backend/agent.py index ff296293..4f2dc7aa 100644 --- a/examples/fastapi-vite/backend/agent.py +++ b/examples/fastapi-vite/backend/agent.py @@ -1,11 +1,11 @@ """Agent logic for the chat demo. Demonstrates human-in-the-loop tool approval using ToolApproval hooks. -Every tool call is gated behind user approval before execution. +``talk_to_mothership`` is gated behind user approval via the +``require_approval=True`` flag on the tool. """ import asyncio -from collections.abc import AsyncGenerator import ai @@ -42,7 +42,7 @@ async def get_population(city: str) -> int: return {"new york": 8_336_817, "tokyo": 13_960_000}.get(city.lower(), 1_000_000) -@ai.tool +@ai.tool(require_approval=True) async def talk_to_mothership(question: str) -> ai.SubAgentTool: """Contact the mothership for important decisions.""" mothership = ai.agent() @@ -58,65 +58,4 @@ async def talk_to_mothership(question: str) -> ai.SubAgentTool: TOOLS: list[ai.AgentTool] = [get_weather, get_population, talk_to_mothership] -class ChatAgent(ai.Agent): - """Agent graph with human-in-the-loop tool approval. - - Loops: stream LLM -> request approval -> execute tools -> repeat. - The ToolApproval hook suspends execution and emits an approval- - request event on the SSE stream. The frontend displays Approve / - Reject buttons and sends the decision back on the next request. - """ - - async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]: - while context.keep_running(): - async with ( - ai.stream(context=context) as s, - ai.ToolRunner() as tr, - ): - async for event in ai.util.merge(s, tr.events()): - yield event - if isinstance(event, ai.events.ToolEnd): - tc = _resolve(context, event.tool_call) - tr.schedule(tc) - - context.add(s.message) - context.add(tr.get_tool_message()) - - -chat_agent = ChatAgent(tools=TOOLS) - - -def _resolve( - context: ai.Context, tool_call: ai.messages.ToolCallPart -) -> ai.ToolCallLike: - tc = context.resolve(tool_call) - if tc.name == "talk_to_mothership": - return lambda: _execute_with_approval(tc) - else: - return tc - - -async def _execute_with_approval(tc: ai.ToolCall) -> ai.events.ToolCallResult: - """Execute a tool call only after the user grants approval. - - Creates a ToolApproval hook that suspends execution until the - frontend responds with an approve/reject decision. - """ - try: - approval = await ai.hook( - f"approve_{tc.id}", - payload=ai.tools.ToolApproval, - metadata={"tool_name": tc.name, "tool_kwargs": tc.kwargs}, - ) - except ai.agents.hooks.HookPendingError as e: - return ai.pending_tool_result(e.hook, tool_call_id=tc.id, tool_name=tc.name) - - if approval.granted: - return await tc() - - return ai.tool_result( - tool_call_id=tc.id, - tool_name=tc.name, - result="Tool call was denied by the user.", - is_error=True, - ) +chat_agent = ai.Agent(tools=TOOLS) diff --git a/examples/multiagent-textual/server.py b/examples/multiagent-textual/server.py index a57ce784..66d389cd 100644 --- a/examples/multiagent-textual/server.py +++ b/examples/multiagent-textual/server.py @@ -78,7 +78,7 @@ class GatedCall: path before the underlying tool runs. """ - def __init__(self, tc: ai.ToolCall, label: str) -> None: + def __init__(self, tc: ai.ToolCallLike, label: str) -> None: self._tc = tc self._label = label diff --git a/examples/samples/agent_hooks.py b/examples/samples/agent_hooks.py index 8200c3ed..e54c0335 100644 --- a/examples/samples/agent_hooks.py +++ b/examples/samples/agent_hooks.py @@ -1,81 +1,27 @@ """Human-in-the-loop approval hooks. Demonstrates the function-based hook API: - - await hook("label", payload=Model) to suspend inside the loop + - mark a tool with ``require_approval=True`` to gate its execution + behind an approval hook - resolve_hook("label", data) to unblock from outside - Hook signals arrive as HookEvent events - -The custom loop uses the concurrent ``ToolRunner`` flow: tools are -scheduled and run concurrently as the model emits them. The approval -hook is awaited inside a ``ToolCall``-shaped wrapper that is scheduled -in place of the bare tool call, so gating composes naturally with the -runner's merge-and-iterate behaviour. """ import asyncio -from collections.abc import AsyncGenerator import ai -@ai.tool +@ai.tool(require_approval=True) async def contact_mothership(query: str) -> str: """Contact the mothership for important decisions.""" return "Soon." -class GatedCall: - """ToolCall-shaped wrapper that awaits an approval hook before executing. - - ``ToolRunner.schedule`` only consumes the ``__call__`` shape of - ``ToolCall``; this wrapper supplies the same shape while inserting - the hook await + denial path before the underlying tool runs. - """ - - def __init__(self, tc: ai.ToolCall) -> None: - self._tc = tc - - async def __call__(self) -> ai.events.ToolCallResult: - tc = self._tc - approval = await ai.hook( - f"approve_{tc.id}", - payload=ai.tools.ToolApproval, - metadata={"tool": tc.name, "kwargs": tc.kwargs}, - ) - if approval.granted: - return await tc() - return ai.tool_result( - tool_call_id=tc.id, - tool_name=tc.name, - result=f"Rejected: {approval.reason}", - is_error=True, - ) - - -class ApprovalAgent(ai.Agent): - async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]: - while context.keep_running(): - async with ( - ai.stream(context=context) as s, - ai.ToolRunner() as tr, - ): - async for event in ai.util.merge(s, tr.events()): - yield event - if isinstance(event, ai.events.ToolEnd): - tc = context.resolve(event.tool_call) - if tc.name == "contact_mothership": - tr.schedule(GatedCall(tc)) - else: - tr.schedule(tc) - - context.add(s.message) - context.add(tr.get_tool_message()) - - async def main() -> None: model = ai.ai_gateway("anthropic/claude-sonnet-4") - my_agent = ApprovalAgent(tools=[contact_mothership]) + my_agent = ai.Agent(tools=[contact_mothership]) messages = [ ai.system_message( diff --git a/examples/samples/agent_hooks_serverless.py b/examples/samples/agent_hooks_serverless.py index 7412304f..33978c9d 100644 --- a/examples/samples/agent_hooks_serverless.py +++ b/examples/samples/agent_hooks_serverless.py @@ -13,14 +13,13 @@ """ import asyncio -from collections.abc import AsyncGenerator import ai FILES_DELETED = set() -@ai.tool +@ai.tool(require_approval=True) async def delete_file(path: str) -> str: """Delete a file at the given path.""" print("FILE DELETED:", path) @@ -39,61 +38,10 @@ async def audit_log(message: str) -> str: return f"Logged {message!r}" -class GatedCall: - """ToolCall-shaped wrapper that awaits an approval hook before executing. - - ``ToolRunner.schedule`` only consumes the ``__call__`` shape of - ``ToolCall``; this wrapper supplies the same shape while inserting - the hook await + denial path before the underlying tool runs. - """ - - def __init__(self, tc: ai.ToolCall) -> None: - self._tc = tc - - async def __call__(self) -> ai.events.ToolCallResult: - tc = self._tc - try: - approval = await ai.hook( - f"approve_{tc.id}", - payload=ai.tools.ToolApproval, - metadata={"tool": tc.name, "kwargs": tc.kwargs}, - ) - except ai.agents.hooks.HookPendingError as e: - return ai.pending_tool_result(e.hook, tool_call_id=tc.id, tool_name=tc.name) - if approval.granted: - return await tc() - return ai.tool_result( - tool_call_id=tc.id, - tool_name=tc.name, - result=f"Rejected: {approval.reason}", - is_error=True, - ) - - -class ConfirmAgent(ai.Agent): - async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]: - while context.keep_running(): - async with ( - ai.stream(context=context) as s, - ai.ToolRunner() as tr, - ): - async for event in ai.util.merge(s, tr.events()): - yield event - if isinstance(event, ai.events.ToolEnd): - tc = context.resolve(event.tool_call) - if tc.name == "delete_file": - tr.schedule(GatedCall(tc)) - else: - tr.schedule(tc) - - context.add(s.message) - context.add(tr.get_tool_message()) - - async def main() -> None: model = ai.ai_gateway("anthropic/claude-sonnet-4") - my_agent = ConfirmAgent(tools=[delete_file, audit_log]) + my_agent = ai.Agent(tools=[delete_file, audit_log]) messages = [ ai.system_message(""" diff --git a/examples/temporal-direct/main.py b/examples/temporal-direct/main.py index bfb8b3ef..60d1a106 100644 --- a/examples/temporal-direct/main.py +++ b/examples/temporal-direct/main.py @@ -174,8 +174,8 @@ async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent def _activity_tool_call( tc: ai.messages.ToolCallPart, -) -> ai.ToolCallLike: - """Build a ``ToolCallLike`` that runs the tool as a Temporal activity. +) -> ai.agents.ToolCallCallable: + """Build a ``ToolCallCallable`` that runs the tool as a Temporal activity. ``ToolRunner.schedule`` accepts any zero-arg callable that returns a coroutine resolving to a ``ToolCallResult``. This lets us route diff --git a/skills/ai/SKILL.md b/skills/ai/SKILL.md index fdb01327..cadd44a8 100644 --- a/skills/ai/SKILL.md +++ b/skills/ai/SKILL.md @@ -181,17 +181,28 @@ Each forwarded event is wrapped in `ai.events.PartialToolCallResult` carrying th ## Hooks -Typed suspension points for human-in-the-loop: +Typed suspension points for human-in-the-loop. + +**Tool approval (built-in shortcut).** Flag a tool with `require_approval=True` and the default loop gates each call behind an `ai.tools.ToolApproval` hook (label `approve_{tool_call_id}`, payload carries `granted` + `reason`): ```python -class Approval(pydantic.BaseModel): - granted: bool - reason: str +@ai.tool(require_approval=True) +async def delete_file(path: str) -> str: + ... + +# consumer-side resolve: +ai.resolve_hook(hook.hook_id, ai.tools.ToolApproval(granted=True)) ``` -Inside agent code (blocks until resolved): +Denial returns an error `ToolCallResult` with `result=f"Rejected: {reason}"`. For custom payloads, label schemes, or per-call gating, write a custom loop using the primitives below. + +**Manual hooks.** Inside agent code (blocks until resolved): ```python +class Approval(pydantic.BaseModel): + granted: bool + reason: str + approval = await ai.hook( "approve_send_email", payload=Approval, diff --git a/src/ai/agents/__init__.py b/src/ai/agents/__init__.py index 8095e7c6..8bd77a2f 100644 --- a/src/ai/agents/__init__.py +++ b/src/ai/agents/__init__.py @@ -14,6 +14,7 @@ SubAgentTool, Tool, ToolCall, + ToolCallCallable, ToolCallLike, ToolRunner, agent, @@ -44,6 +45,7 @@ "SubAgentTool", "Tool", "ToolCall", + "ToolCallCallable", "ToolCallLike", "ToolRunner", "StreamingStatusTool", diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index e2b62736..e7065feb 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -37,6 +37,7 @@ from .. import models, types, util from ..types import builders from ..types import events as events_ +from . import hooks as hooks_ from . import middleware as middleware_ from . import runtime @@ -356,6 +357,10 @@ class AgentTool: def name(self) -> str: return self.tool.name + @property + def require_approval(self) -> bool: + return self.tool.require_approval + @property def _aggregator( self, @@ -382,9 +387,15 @@ def tool[**P, R](fn: Callable[P, Awaitable[R]], /) -> AgentTool: ... def tool[**P, T](fn: Callable[P, AsyncGenerator[T]], /) -> AgentTool: ... +@overload +def tool[**P](*, require_approval: bool) -> Callable[[Callable[P, Any]], AgentTool]: ... + + @overload def tool[**P, T, R]( - *, aggregator: Callable[[], events_.Aggregator[T, Any, R]] + *, + aggregator: Callable[[], events_.Aggregator[T, Any, R]], + require_approval: bool = False, ) -> Callable[[Callable[P, AsyncGenerator[T]]], AgentTool]: ... @@ -393,7 +404,12 @@ def tool[**P, T, R]( /, *, aggregator: Callable[[], events_.Aggregator[T, Any, R]] | None = None, -) -> Callable[[Callable[P, AsyncGenerator[T]]], AgentTool] | AgentTool: + require_approval: bool = False, +) -> ( + Callable[[Callable[P, AsyncGenerator[T]]], AgentTool] + | Callable[[Callable[P, Awaitable[R]]], AgentTool] + | AgentTool +): """Decorator: turn an async function into a :class:`Tool`. For async-generator tools, declare the aggregator either via the @@ -433,6 +449,7 @@ def wrap(fn: Any) -> AgentTool: description=inspect.getdoc(fn) or "", params=validator.model_json_schema(), ), + require_approval=require_approval, ) return AgentTool( @@ -565,7 +582,54 @@ async def _real(call: middleware_.ToolContext) -> events_.ToolCallResult: return await chain(call) -class ToolCallLike(Protocol): +class GatedCall: + """ToolCall-shaped wrapper that awaits an approval hook before executing. + + ``ToolRunner.schedule`` only consumes the ``__call__`` shape of + ``ToolCall``; this wrapper supplies the same shape while inserting + the hook await + denial path before the underlying tool runs. + """ + + def __init__(self, tc: ToolCallLike) -> None: + self._tc = tc + + @property + def id(self) -> str: + return self._tc.id + + @property + def name(self) -> str: + return self._tc.name + + @property + def fn(self) -> Callable[..., Awaitable[Any]]: + return self._tc.fn + + @property + def kwargs(self) -> dict[str, Any]: + return self._tc.kwargs + + async def __call__(self) -> events_.ToolCallResult: + tc = self._tc + try: + approval = await hooks_.hook( + f"approve_{tc.id}", + payload=types.tools.ToolApproval, + metadata={"tool": tc.name, "kwargs": tc.kwargs}, + ) + except hooks_.HookPendingError as e: + return pending_tool_result(e.hook, tool_call_id=tc.id, tool_name=tc.name) + if approval.granted: + return await tc() + return tool_result( + tool_call_id=tc.id, + tool_name=tc.name, + result=f"Rejected: {approval.reason}", + is_error=True, + ) + + +class ToolCallCallable(Protocol): """Anything ``ToolRunner.schedule`` can accept. Satisfied by :class:`ToolCall` and by any zero-arg callable returning @@ -576,6 +640,22 @@ class ToolCallLike(Protocol): def __call__(self) -> Coroutine[Any, Any, events_.ToolCallResult]: ... +class ToolCallLike(ToolCallCallable, Protocol): + """Something with all the key information for a tool call.""" + + @property + def id(self) -> str: ... + + @property + def name(self) -> str: ... + + @property + def fn(self) -> Callable[..., Awaitable[Any]]: ... + + @property + def kwargs(self) -> dict[str, Any]: ... + + class _RestartableToolStream: def __init__(self, tr: ToolRunner) -> None: self._tr = tr @@ -611,10 +691,10 @@ async def __aexit__(self, *args: Any) -> None: def events(self) -> _RestartableToolStream: return _RestartableToolStream(self) - def schedule(self, tc: ToolCallLike) -> None: + def schedule(self, tc: ToolCallCallable) -> None: """Schedule a tool call (or any callable producing a ToolCallResult). - See :class:`ToolCallLike` — accepts both :class:`ToolCall` and + See :class:`ToolCallCallable` — accepts both :class:`ToolCall` and any zero-arg callable returning a coroutine that resolves to a :class:`ToolCallResult`. The latter lets you wrap a ``ToolCall`` in custom logic (e.g. an approval hook await) and still ride the @@ -699,16 +779,16 @@ def keep_running(self) -> bool: return last_message.replay or last_message.role not in ("assistant", "internal") @overload - def resolve(self, tool_part: types.messages.ToolCallPart) -> ToolCall: ... + def resolve(self, tool_part: types.messages.ToolCallPart) -> ToolCallLike: ... @overload def resolve( self, tool_part: Sequence[types.messages.ToolCallPart] - ) -> list[ToolCall]: ... + ) -> list[ToolCallLike]: ... def resolve( self, tool_part: types.messages.ToolCallPart | Sequence[types.messages.ToolCallPart], - ) -> ToolCall | list[ToolCall]: + ) -> ToolCallLike | list[ToolCallLike]: """Resolve ToolCallPart(s) into callable ToolCall object(s).""" if isinstance(tool_part, types.messages.ToolCallPart): tool = self._agent_tools_by_name.get(tool_part.tool_name) @@ -716,7 +796,10 @@ def resolve( raise KeyError( f"No agent executor registered for tool {tool_part.tool_name!r}" ) - return ToolCall(part=tool_part, tool=tool) + tc = ToolCall(part=tool_part, tool=tool) + if tool.require_approval: + return GatedCall(tc) + return tc return [self.resolve(tp) for tp in tool_part] def add( diff --git a/src/ai/types/tools.py b/src/ai/types/tools.py index f6a96d12..b8193bbf 100644 --- a/src/ai/types/tools.py +++ b/src/ai/types/tools.py @@ -21,6 +21,7 @@ class Tool(pydantic.BaseModel): kind: Literal["function", "provider"] name: str args: pydantic.BaseModel + require_approval: bool = False @pydantic.model_validator(mode="after") def validate_args_shape(self) -> Self: