b60439152b
- harness/es.py: append-only event store + проекции (образец — кооп-портал) - harness/levels.py: уровни 1с/1м/1ч/1д, контекст ∝ масштабу - harness/scheduler.py: уровневые потоки, идемпотентные тики, absorbed-стабы - harness/isolated_runner.py + _llm_worker.py: LLM-тик в сабпроцессе (глобальный asyncio-рантайм Strands не переживает 2+ create_harness) - cli.py: init/task/tasks/status/events + requeue (написан самим минутным агентом) - 9router: отдельный ключ sk-simbiot, модель auto (glm-5.3 первым)
448 lines
16 KiB
Python
448 lines
16 KiB
Python
"""Hook registry system for managing event callbacks in the Strands Agent SDK.
|
|
|
|
This module provides the core infrastructure for the typed hook system, enabling
|
|
composable extension of agent functionality through strongly-typed event callbacks.
|
|
The registry manages the mapping between event types and their associated callback
|
|
functions, supporting both individual callback registration and bulk registration
|
|
via hook provider objects.
|
|
"""
|
|
|
|
import bisect
|
|
import inspect
|
|
import logging
|
|
from collections.abc import Awaitable, Generator
|
|
from dataclasses import dataclass
|
|
from itertools import groupby
|
|
from typing import TYPE_CHECKING, Any, Generic, Protocol, runtime_checkable
|
|
|
|
from typing_extensions import TypeVar
|
|
|
|
from ..interrupt import Interrupt, InterruptException
|
|
from ._type_inference import infer_event_types
|
|
|
|
if TYPE_CHECKING:
|
|
from ..agent import Agent
|
|
from ..types.agent import LocalAgent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class HookOrder:
|
|
"""Named constants for hook execution priority.
|
|
|
|
Lower values execute first. Hooks with the same order preserve registration order.
|
|
"""
|
|
|
|
SDK_FIRST: int = -100
|
|
INTERVENTION_OUTPUT: int = -90
|
|
DEFAULT: int = 0
|
|
MODEL_ROUTING: int = 50
|
|
INTERVENTION_INPUT: int = 90
|
|
SDK_LAST: int = 100
|
|
|
|
|
|
@dataclass
|
|
class _CallbackEntry:
|
|
"""Internal entry pairing a callback with its execution order."""
|
|
|
|
callback: "HookCallback"
|
|
order: float
|
|
|
|
|
|
@dataclass
|
|
class BaseHookEvent:
|
|
"""Base class for all hook events."""
|
|
|
|
@property
|
|
def should_reverse_callbacks(self) -> bool:
|
|
"""Determine if callbacks for this event should be invoked in reverse order.
|
|
|
|
Returns:
|
|
False by default. Override to return True for events that should
|
|
invoke callbacks in reverse order (e.g., cleanup/teardown events).
|
|
"""
|
|
return False
|
|
|
|
def _can_write(self, name: str) -> bool:
|
|
"""Check if the given property can be written to.
|
|
|
|
Args:
|
|
name: The name of the property to check.
|
|
|
|
Returns:
|
|
True if the property can be written to, False otherwise.
|
|
"""
|
|
return False
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Disallow writes to non-approved properties."""
|
|
# This is needed as otherwise the class can't be initialized at all, so we trigger
|
|
# this after class initialization
|
|
super().__setattr__("_disallow_writes", True)
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
"""Prevent setting attributes on hook events.
|
|
|
|
Raises:
|
|
AttributeError: Always raised to prevent setting attributes on hook events.
|
|
"""
|
|
# Allow setting attributes:
|
|
# - during init (when __dict__) doesn't exist
|
|
# - if the subclass specifically said the property is writable
|
|
if not hasattr(self, "_disallow_writes") or self._can_write(name):
|
|
return super().__setattr__(name, value)
|
|
|
|
raise AttributeError(f"Property {name} is not writable")
|
|
|
|
|
|
_LocalAgentT = TypeVar("_LocalAgentT", bound="LocalAgent", default="Agent", covariant=True)
|
|
"""Agent type carried by hook events."""
|
|
|
|
|
|
@dataclass
|
|
class HookEvent(BaseHookEvent, Generic[_LocalAgentT]):
|
|
"""Base class for single agent hook events.
|
|
|
|
Attributes:
|
|
agent: The agent instance that triggered this event.
|
|
"""
|
|
|
|
agent: _LocalAgentT
|
|
|
|
|
|
TEvent = TypeVar("TEvent", bound=BaseHookEvent, contravariant=True)
|
|
"""Generic for adding callback handlers - contravariant to allow adding handlers which take in base classes."""
|
|
|
|
TInvokeEvent = TypeVar("TInvokeEvent", bound=BaseHookEvent)
|
|
"""Generic for invoking events - non-contravariant to enable returning events."""
|
|
|
|
|
|
@runtime_checkable
|
|
class HookProvider(Protocol):
|
|
"""Protocol for objects that provide hook callbacks to an agent.
|
|
|
|
Hook providers offer a composable way to extend agent functionality by
|
|
subscribing to various events in the agent lifecycle. This protocol enables
|
|
building reusable components that can hook into agent events.
|
|
|
|
Example:
|
|
```python
|
|
class MyHookProvider(HookProvider):
|
|
def register_hooks(self, registry: HookRegistry) -> None:
|
|
registry.add_callback(StartRequestEvent, self.on_request_start)
|
|
registry.add_callback(EndRequestEvent, self.on_request_end)
|
|
|
|
agent = Agent(hooks=[MyHookProvider()])
|
|
```
|
|
"""
|
|
|
|
def register_hooks(self, registry: "HookRegistry", **kwargs: Any) -> None:
|
|
"""Register callback functions for specific event types.
|
|
|
|
Args:
|
|
registry: The hook registry to register callbacks with.
|
|
**kwargs: Additional keyword arguments for future extensibility.
|
|
"""
|
|
...
|
|
|
|
|
|
class HookCallback(Protocol, Generic[TEvent]):
|
|
"""Protocol for callback functions that handle hook events.
|
|
|
|
Hook callbacks are functions that receive a single strongly-typed event
|
|
argument and perform some action in response. They should not return
|
|
values and any exceptions they raise will propagate to the caller.
|
|
|
|
Example:
|
|
```python
|
|
def my_callback(event: StartRequestEvent) -> None:
|
|
print(f"Request started for agent: {event.agent.name}")
|
|
|
|
# Or
|
|
|
|
async def my_callback(event: StartRequestEvent) -> None:
|
|
# await an async operation
|
|
```
|
|
"""
|
|
|
|
def __call__(self, event: TEvent) -> None | Awaitable[None]:
|
|
"""Handle a hook event.
|
|
|
|
Args:
|
|
event: The strongly-typed event to handle.
|
|
"""
|
|
...
|
|
|
|
|
|
class HookRegistry:
|
|
"""Registry for managing hook callbacks associated with event types.
|
|
|
|
The HookRegistry maintains a mapping of event types to callback functions
|
|
and provides methods for registering callbacks and invoking them when
|
|
events occur.
|
|
|
|
The registry handles callback ordering, including reverse ordering for
|
|
cleanup events, and provides type-safe event dispatching.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize an empty hook registry."""
|
|
self._registered_callbacks: dict[type, list[_CallbackEntry]] = {}
|
|
|
|
def add_callback(
|
|
self,
|
|
event_type: type[TEvent] | list[type[TEvent]] | None,
|
|
callback: HookCallback[TEvent],
|
|
*,
|
|
order: float = HookOrder.DEFAULT,
|
|
) -> None:
|
|
"""Register a callback function for a specific event type.
|
|
|
|
If ``event_type`` is None, then this will check the callback handler type hint
|
|
for the lifecycle event type. Union types (``A | B`` or ``Union[A, B]``) in
|
|
type hints will register the callback for each event type in the union.
|
|
|
|
If ``event_type`` is a list, the callback will be registered for each event
|
|
type in the list (duplicates are ignored).
|
|
|
|
Args:
|
|
event_type: The lifecycle event type(s) this callback should handle.
|
|
Can be a single type, a list of types, or None to infer from type hints.
|
|
callback: The callback function to invoke when events of this type occur.
|
|
order: Execution priority. Lower values execute first.
|
|
|
|
Raises:
|
|
ValueError: If event_type is not provided and cannot be inferred from
|
|
the callback's type hints, or if AgentInitializedEvent is registered
|
|
with an async callback, or if the event_type list is empty.
|
|
|
|
Example:
|
|
```python
|
|
def my_handler(event: StartRequestEvent):
|
|
print("Request started")
|
|
|
|
# With explicit event type
|
|
registry.add_callback(StartRequestEvent, my_handler)
|
|
|
|
# With event type inferred from type hint
|
|
registry.add_callback(None, my_handler)
|
|
|
|
# With union type hint (registers for both types)
|
|
def union_handler(event: BeforeModelCallEvent | AfterModelCallEvent):
|
|
print(f"Event: {type(event).__name__}")
|
|
registry.add_callback(None, union_handler)
|
|
|
|
# With list of event types
|
|
def multi_handler(event):
|
|
print(f"Event: {type(event).__name__}")
|
|
registry.add_callback([BeforeModelCallEvent, AfterModelCallEvent], multi_handler)
|
|
```
|
|
"""
|
|
resolved_event_types: list[type[TEvent]]
|
|
|
|
# Handle list of event types
|
|
if isinstance(event_type, list):
|
|
if not event_type:
|
|
raise ValueError("event_type list cannot be empty")
|
|
resolved_event_types = self._validate_event_type_list(event_type)
|
|
elif event_type is None:
|
|
# Infer event type(s) from callback type hints
|
|
resolved_event_types = infer_event_types(callback)
|
|
else:
|
|
# Single event type provided explicitly
|
|
resolved_event_types = [event_type]
|
|
|
|
# Deduplicate event types while preserving order
|
|
unique_event_types: set[type[TEvent]] = set(resolved_event_types)
|
|
|
|
# Register callback for each event type
|
|
for resolved_event_type in unique_event_types:
|
|
# Related issue: https://github.com/strands-agents/harness-sdk/issues/330
|
|
if resolved_event_type.__name__ == "AgentInitializedEvent" and inspect.iscoroutinefunction(callback):
|
|
raise ValueError("AgentInitializedEvent can only be registered with a synchronous callback")
|
|
|
|
entries = self._registered_callbacks.setdefault(resolved_event_type, [])
|
|
entry = _CallbackEntry(callback=callback, order=order)
|
|
bisect.insort(entries, entry, key=lambda e: e.order)
|
|
|
|
def _validate_event_type_list(self, event_types: list[type[TEvent]]) -> list[type[TEvent]]:
|
|
"""Validate that all types in a list are valid BaseHookEvent subclasses.
|
|
|
|
Args:
|
|
event_types: List of event types to validate.
|
|
|
|
Returns:
|
|
The validated list of event types.
|
|
|
|
Raises:
|
|
ValueError: If any type is not a valid BaseHookEvent subclass.
|
|
"""
|
|
validated: list[type[TEvent]] = []
|
|
for et in event_types:
|
|
if not (isinstance(et, type) and issubclass(et, BaseHookEvent)):
|
|
raise ValueError(f"Invalid event type: {et} | must be a subclass of BaseHookEvent")
|
|
validated.append(et)
|
|
return validated
|
|
|
|
def add_hook(self, hook: HookProvider) -> None:
|
|
"""Register all callbacks from a hook provider.
|
|
|
|
This method allows bulk registration of callbacks by delegating to
|
|
the hook provider's register_hooks method. This is the preferred
|
|
way to register multiple related callbacks.
|
|
|
|
Args:
|
|
hook: The hook provider containing callbacks to register.
|
|
|
|
Example:
|
|
```python
|
|
class MyHooks(HookProvider):
|
|
def register_hooks(self, registry: HookRegistry):
|
|
registry.add_callback(StartRequestEvent, self.on_start)
|
|
registry.add_callback(EndRequestEvent, self.on_end)
|
|
|
|
registry.add_hook(MyHooks())
|
|
```
|
|
"""
|
|
hook.register_hooks(self)
|
|
|
|
async def invoke_callbacks_async(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Interrupt]]:
|
|
"""Invoke all registered callbacks for the given event.
|
|
|
|
This method finds all callbacks registered for the event's type and
|
|
invokes them in the appropriate order. For events with should_reverse_callbacks=True,
|
|
callbacks are invoked in reverse registration order. Any exceptions raised by callback
|
|
functions will propagate to the caller.
|
|
|
|
Additionally, this method aggregates interrupts raised by the user to instantiate human-in-the-loop workflows.
|
|
|
|
Args:
|
|
event: The event to dispatch to registered callbacks.
|
|
|
|
Returns:
|
|
The event dispatched to registered callbacks and any interrupts raised by the user.
|
|
|
|
Raises:
|
|
ValueError: If interrupt name is used more than once.
|
|
|
|
Example:
|
|
```python
|
|
event = StartRequestEvent(agent=my_agent)
|
|
await registry.invoke_callbacks_async(event)
|
|
```
|
|
"""
|
|
interrupts: dict[str, Interrupt] = {}
|
|
|
|
for callback in self.get_callbacks_for(event):
|
|
try:
|
|
if inspect.iscoroutinefunction(callback):
|
|
await callback(event)
|
|
else:
|
|
callback(event)
|
|
|
|
except InterruptException as exception:
|
|
interrupt = exception.interrupt
|
|
if interrupt.name in interrupts:
|
|
message = f"interrupt_name=<{interrupt.name}> | interrupt name used more than once"
|
|
logger.error(message)
|
|
raise ValueError(message) from exception
|
|
|
|
# Each callback is allowed to raise their own interrupt.
|
|
interrupts[interrupt.name] = interrupt
|
|
|
|
return event, list(interrupts.values())
|
|
|
|
def invoke_callbacks(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Interrupt]]:
|
|
"""Invoke all registered callbacks for the given event.
|
|
|
|
This method finds all callbacks registered for the event's type and
|
|
invokes them in the appropriate order. For events with should_reverse_callbacks=True,
|
|
callbacks are invoked in reverse registration order. Any exceptions raised by callback
|
|
functions will propagate to the caller.
|
|
|
|
Additionally, this method aggregates interrupts raised by the user to instantiate human-in-the-loop workflows.
|
|
|
|
Args:
|
|
event: The event to dispatch to registered callbacks.
|
|
|
|
Returns:
|
|
The event dispatched to registered callbacks and any interrupts raised by the user.
|
|
|
|
Raises:
|
|
RuntimeError: If at least one callback is async.
|
|
ValueError: If interrupt name is used more than once.
|
|
|
|
Example:
|
|
```python
|
|
event = StartRequestEvent(agent=my_agent)
|
|
registry.invoke_callbacks(event)
|
|
```
|
|
"""
|
|
callbacks = list(self.get_callbacks_for(event))
|
|
interrupts: dict[str, Interrupt] = {}
|
|
|
|
if any(inspect.iscoroutinefunction(callback) for callback in callbacks):
|
|
raise RuntimeError(f"event=<{event}> | use invoke_callbacks_async to invoke async callback")
|
|
|
|
for callback in callbacks:
|
|
try:
|
|
callback(event)
|
|
except InterruptException as exception:
|
|
interrupt = exception.interrupt
|
|
if interrupt.name in interrupts:
|
|
message = f"interrupt_name=<{interrupt.name}> | interrupt name used more than once"
|
|
logger.error(message)
|
|
raise ValueError(message) from exception
|
|
|
|
# Each callback is allowed to raise their own interrupt.
|
|
interrupts[interrupt.name] = interrupt
|
|
|
|
return event, list(interrupts.values())
|
|
|
|
def has_callbacks(self) -> bool:
|
|
"""Check if the registry has any registered callbacks.
|
|
|
|
Returns:
|
|
True if there are any registered callbacks, False otherwise.
|
|
|
|
Example:
|
|
```python
|
|
if registry.has_callbacks():
|
|
print("Registry has callbacks registered")
|
|
```
|
|
"""
|
|
return bool(self._registered_callbacks)
|
|
|
|
def get_callbacks_for(self, event: TEvent) -> Generator[HookCallback[TEvent], None, None]:
|
|
"""Get callbacks registered for the given event in the appropriate order.
|
|
|
|
For normal events, callbacks are returned in order priority (lower first),
|
|
with registration order preserved within the same priority.
|
|
|
|
For reversed events (should_reverse_callbacks=True), order priority still
|
|
applies (lower first), but within the same priority group, registration
|
|
order is reversed.
|
|
|
|
Args:
|
|
event: The event to get callbacks for.
|
|
|
|
Yields:
|
|
Callback functions registered for this event type, in the appropriate order.
|
|
|
|
Example:
|
|
```python
|
|
event = EndRequestEvent(agent=my_agent)
|
|
for callback in registry.get_callbacks_for(event):
|
|
callback(event)
|
|
```
|
|
"""
|
|
event_type = type(event)
|
|
|
|
entries = self._registered_callbacks.get(event_type, [])
|
|
if event.should_reverse_callbacks:
|
|
for _order, group in groupby(entries, key=lambda e: e.order):
|
|
for entry in reversed(list(group)):
|
|
yield entry.callback
|
|
else:
|
|
for entry in entries:
|
|
yield entry.callback
|