refactor: unify agent runner lifecycle hooks
This commit is contained in:
Xubin Ren
Xubin Ren
+34
-23
@@ -3,12 +3,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
||||
from nanobot.providers.base import LLMProvider, ToolCallRequest
|
||||
from nanobot.utils.helpers import build_assistant_message
|
||||
|
||||
_DEFAULT_MAX_ITERATIONS_MESSAGE = (
|
||||
@@ -29,11 +29,7 @@ class AgentRunSpec:
|
||||
temperature: float | None = None
|
||||
max_tokens: int | None = None
|
||||
reasoning_effort: str | None = None
|
||||
on_stream: Callable[[str], Awaitable[None]] | None = None
|
||||
on_stream_end: Callable[..., Awaitable[None]] | None = None
|
||||
on_tool_calls: Callable[[LLMResponse], Awaitable[None] | None] | None = None
|
||||
before_execute_tools: Callable[[list[ToolCallRequest]], Awaitable[None] | None] | None = None
|
||||
finalize_content: Callable[[str | None], str | None] | None = None
|
||||
hook: AgentHook | None = None
|
||||
error_message: str | None = _DEFAULT_ERROR_MESSAGE
|
||||
max_iterations_message: str | None = None
|
||||
concurrent_tools: bool = False
|
||||
@@ -60,6 +56,7 @@ class AgentRunner:
|
||||
self.provider = provider
|
||||
|
||||
async def run(self, spec: AgentRunSpec) -> AgentRunResult:
|
||||
hook = spec.hook or AgentHook()
|
||||
messages = list(spec.initial_messages)
|
||||
final_content: str | None = None
|
||||
tools_used: list[str] = []
|
||||
@@ -68,7 +65,9 @@ class AgentRunner:
|
||||
stop_reason = "completed"
|
||||
tool_events: list[dict[str, str]] = []
|
||||
|
||||
for _ in range(spec.max_iterations):
|
||||
for iteration in range(spec.max_iterations):
|
||||
context = AgentHookContext(iteration=iteration, messages=messages)
|
||||
await hook.before_iteration(context)
|
||||
kwargs: dict[str, Any] = {
|
||||
"messages": messages,
|
||||
"tools": spec.tools.get_definitions(),
|
||||
@@ -81,10 +80,13 @@ class AgentRunner:
|
||||
if spec.reasoning_effort is not None:
|
||||
kwargs["reasoning_effort"] = spec.reasoning_effort
|
||||
|
||||
if spec.on_stream:
|
||||
if hook.wants_streaming():
|
||||
async def _stream(delta: str) -> None:
|
||||
await hook.on_stream(context, delta)
|
||||
|
||||
response = await self.provider.chat_stream_with_retry(
|
||||
**kwargs,
|
||||
on_content_delta=spec.on_stream,
|
||||
on_content_delta=_stream,
|
||||
)
|
||||
else:
|
||||
response = await self.provider.chat_with_retry(**kwargs)
|
||||
@@ -94,14 +96,13 @@ class AgentRunner:
|
||||
"prompt_tokens": int(raw_usage.get("prompt_tokens", 0) or 0),
|
||||
"completion_tokens": int(raw_usage.get("completion_tokens", 0) or 0),
|
||||
}
|
||||
context.response = response
|
||||
context.usage = usage
|
||||
context.tool_calls = list(response.tool_calls)
|
||||
|
||||
if response.has_tool_calls:
|
||||
if spec.on_stream_end:
|
||||
await spec.on_stream_end(resuming=True)
|
||||
if spec.on_tool_calls:
|
||||
maybe = spec.on_tool_calls(response)
|
||||
if maybe is not None:
|
||||
await maybe
|
||||
if hook.wants_streaming():
|
||||
await hook.on_stream_end(context, resuming=True)
|
||||
|
||||
messages.append(build_assistant_message(
|
||||
response.content or "",
|
||||
@@ -111,16 +112,18 @@ class AgentRunner:
|
||||
))
|
||||
tools_used.extend(tc.name for tc in response.tool_calls)
|
||||
|
||||
if spec.before_execute_tools:
|
||||
maybe = spec.before_execute_tools(response.tool_calls)
|
||||
if maybe is not None:
|
||||
await maybe
|
||||
await hook.before_execute_tools(context)
|
||||
|
||||
results, new_events, fatal_error = await self._execute_tools(spec, response.tool_calls)
|
||||
tool_events.extend(new_events)
|
||||
context.tool_results = list(results)
|
||||
context.tool_events = list(new_events)
|
||||
if fatal_error is not None:
|
||||
error = f"Error: {type(fatal_error).__name__}: {fatal_error}"
|
||||
stop_reason = "tool_error"
|
||||
context.error = error
|
||||
context.stop_reason = stop_reason
|
||||
await hook.after_iteration(context)
|
||||
break
|
||||
for tool_call, result in zip(response.tool_calls, results):
|
||||
messages.append({
|
||||
@@ -129,16 +132,21 @@ class AgentRunner:
|
||||
"name": tool_call.name,
|
||||
"content": result,
|
||||
})
|
||||
await hook.after_iteration(context)
|
||||
continue
|
||||
|
||||
if spec.on_stream_end:
|
||||
await spec.on_stream_end(resuming=False)
|
||||
if hook.wants_streaming():
|
||||
await hook.on_stream_end(context, resuming=False)
|
||||
|
||||
clean = spec.finalize_content(response.content) if spec.finalize_content else response.content
|
||||
clean = hook.finalize_content(context, response.content)
|
||||
if response.finish_reason == "error":
|
||||
final_content = clean or spec.error_message or _DEFAULT_ERROR_MESSAGE
|
||||
stop_reason = "error"
|
||||
error = final_content
|
||||
context.final_content = final_content
|
||||
context.error = error
|
||||
context.stop_reason = stop_reason
|
||||
await hook.after_iteration(context)
|
||||
break
|
||||
|
||||
messages.append(build_assistant_message(
|
||||
@@ -147,6 +155,9 @@ class AgentRunner:
|
||||
thinking_blocks=response.thinking_blocks,
|
||||
))
|
||||
final_content = clean
|
||||
context.final_content = final_content
|
||||
context.stop_reason = stop_reason
|
||||
await hook.after_iteration(context)
|
||||
break
|
||||
else:
|
||||
stop_reason = "max_iterations"
|
||||
|
||||
Reference in New Issue
Block a user