| 1 | """Shared lifecycle hook primitives for agent runs.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import Any |
| 7 | |
| 8 | from loguru import logger |
| 9 | |
| 10 | from nanobot.providers.base import LLMResponse, ToolCallRequest |
| 11 | |
| 12 | |
| 13 | @dataclass(slots=True) |
| 14 | class AgentHookContext: |
| 15 | """Mutable per-iteration state exposed to runner hooks.""" |
| 16 | |
| 17 | iteration: int |
| 18 | messages: list[dict[str, Any]] |
| 19 | response: LLMResponse | None = None |
| 20 | usage: dict[str, int] = field(default_factory=dict) |
| 21 | tool_calls: list[ToolCallRequest] = field(default_factory=list) |
| 22 | tool_results: list[Any] = field(default_factory=list) |
| 23 | tool_events: list[dict[str, str]] = field(default_factory=list) |
| 24 | final_content: str | None = None |
| 25 | stop_reason: str | None = None |
| 26 | error: str | None = None |
| 27 | |
| 28 | |
| 29 | class AgentHook: |
| 30 | """Minimal lifecycle surface for shared runner customization.""" |
| 31 | |
| 32 | def __init__(self, reraise: bool = False) -> None: |
| 33 | self._reraise = reraise |
| 34 | |
| 35 | def wants_streaming(self) -> bool: |
| 36 | return False |
| 37 | |
| 38 | async def before_iteration(self, context: AgentHookContext) -> None: |
| 39 | pass |
| 40 | |
| 41 | async def on_stream(self, context: AgentHookContext, delta: str) -> None: |
| 42 | pass |
| 43 | |
| 44 | async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: |
| 45 | pass |
| 46 | |
| 47 | async def before_execute_tools(self, context: AgentHookContext) -> None: |
| 48 | pass |
| 49 | |
| 50 | async def after_iteration(self, context: AgentHookContext) -> None: |
| 51 | pass |
| 52 | |
| 53 | def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None: |
| 54 | return content |
| 55 | |
| 56 | |
| 57 | class CompositeHook(AgentHook): |
| 58 | """Fan-out hook that delegates to an ordered list of hooks. |
| 59 | |
| 60 | Error isolation: async methods catch and log per-hook exceptions |
| 61 | so a faulty custom hook cannot crash the agent loop. |
| 62 | ``finalize_content`` is a pipeline (no isolation — bugs should surface). |
| 63 | """ |
| 64 | |
| 65 | __slots__ = ("_hooks",) |
| 66 | |
| 67 | def __init__(self, hooks: list[AgentHook]) -> None: |
| 68 | super().__init__() |
| 69 | self._hooks = list(hooks) |
| 70 | |
| 71 | def wants_streaming(self) -> bool: |
| 72 | return any(h.wants_streaming() for h in self._hooks) |
| 73 | |
| 74 | async def _for_each_hook_safe(self, method_name: str, *args: Any, **kwargs: Any) -> None: |
| 75 | for h in self._hooks: |
| 76 | if getattr(h, "_reraise", False): |
| 77 | await getattr(h, method_name)(*args, **kwargs) |
| 78 | continue |
| 79 | |
| 80 | try: |
| 81 | await getattr(h, method_name)(*args, **kwargs) |
| 82 | except Exception: |
| 83 | logger.exception("AgentHook.{} error in {}", method_name, type(h).__name__) |
| 84 | |
| 85 | async def before_iteration(self, context: AgentHookContext) -> None: |
| 86 | await self._for_each_hook_safe("before_iteration", context) |
| 87 | |
| 88 | async def on_stream(self, context: AgentHookContext, delta: str) -> None: |
| 89 | await self._for_each_hook_safe("on_stream", context, delta) |
| 90 | |
| 91 | async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None: |
| 92 | await self._for_each_hook_safe("on_stream_end", context, resuming=resuming) |
| 93 | |
| 94 | async def before_execute_tools(self, context: AgentHookContext) -> None: |
| 95 | await self._for_each_hook_safe("before_execute_tools", context) |
| 96 | |
| 97 | async def after_iteration(self, context: AgentHookContext) -> None: |
| 98 | await self._for_each_hook_safe("after_iteration", context) |
| 99 | |
| 100 | def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None: |
| 101 | for h in self._hooks: |
| 102 | content = h.finalize_content(context, content) |
| 103 | return content |
| 104 |