-
Notifications
You must be signed in to change notification settings - Fork 676
feat(plugins): improve plugin creation devex with @hook and @tool decorators #1740
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
730e528
d2a680e
403999a
eaf6471
717cfec
c2a3296
0233286
3fc87d1
703cece
cb098d6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| """Utility for inferring event types from callback type hints.""" | ||
|
|
||
| import inspect | ||
| import logging | ||
| import types | ||
| from typing import TYPE_CHECKING, Union, cast, get_args, get_origin, get_type_hints | ||
|
|
||
| if TYPE_CHECKING: | ||
| from .registry import HookCallback, TEvent | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def infer_event_types(callback: "HookCallback[TEvent]") -> "list[type[TEvent]]": | ||
| """Infer the event type(s) from a callback's type hints. | ||
|
|
||
| Supports both single types and union types (A | B or Union[A, B]). | ||
|
|
||
| Args: | ||
| callback: The callback function to inspect. | ||
|
|
||
| Returns: | ||
| A list of event types inferred from the callback's first parameter type hint. | ||
|
|
||
| Raises: | ||
| ValueError: If the event type cannot be inferred from the callback's type hints, | ||
| or if a union contains None or non-BaseHookEvent types. | ||
| """ | ||
| # Import here to avoid circular dependency | ||
| from .registry import BaseHookEvent | ||
|
|
||
| try: | ||
| hints = get_type_hints(callback) | ||
| except Exception as e: | ||
| logger.debug("callback=<%s>, error=<%s> | failed to get type hints", callback, e) | ||
| raise ValueError( | ||
| "failed to get type hints for callback | cannot infer event type, please provide event_type explicitly" | ||
| ) from e | ||
|
|
||
| # Get the first parameter's type hint | ||
| sig = inspect.signature(callback) | ||
| params = list(sig.parameters.values()) | ||
|
|
||
| if not params: | ||
| raise ValueError("callback has no parameters | cannot infer event type, please provide event_type explicitly") | ||
|
|
||
| # Skip 'self' parameter for methods | ||
| first_param = params[0] | ||
| if first_param.name == "self" and len(params) > 1: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Question: do we have validation check else where that hanldes edge cases like |
||
| first_param = params[1] | ||
|
|
||
| type_hint = hints.get(first_param.name) | ||
|
|
||
| if type_hint is None: | ||
| raise ValueError( | ||
| f"parameter=<{first_param.name}> has no type hint | " | ||
| "cannot infer event type, please provide event_type explicitly" | ||
| ) | ||
|
|
||
| # Check if it's a Union type (Union[A, B] or A | B) | ||
| origin = get_origin(type_hint) | ||
| if origin is Union or origin is types.UnionType: | ||
| event_types: list[type[TEvent]] = [] | ||
| for arg in get_args(type_hint): | ||
| if arg is type(None): | ||
| raise ValueError("None is not a valid event type in union") | ||
| if not (isinstance(arg, type) and issubclass(arg, BaseHookEvent)): | ||
| raise ValueError(f"Invalid type in union: {arg} | must be a subclass of BaseHookEvent") | ||
| event_types.append(cast("type[TEvent]", arg)) | ||
| return event_types | ||
|
|
||
| # Handle single type | ||
| if isinstance(type_hint, type) and issubclass(type_hint, BaseHookEvent): | ||
| return [cast("type[TEvent]", type_hint)] | ||
|
|
||
| raise ValueError( | ||
| f"parameter=<{first_param.name}>, type=<{type_hint}> | type hint must be a subclass of BaseHookEvent" | ||
| ) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,25 +1,13 @@ | ||
| """Plugin system for extending agent functionality. | ||
|
|
||
| This module provides a composable mechanism for building objects that can | ||
| extend agent behavior through a standardized initialization pattern. | ||
|
|
||
| Example Usage: | ||
| ```python | ||
| from strands.plugins import Plugin | ||
|
|
||
| class LoggingPlugin(Plugin): | ||
| name = "logging" | ||
|
|
||
| def init_agent(self, agent: Agent) -> None: | ||
| agent.add_hook(self.on_model_call, BeforeModelCallEvent) | ||
|
|
||
| def on_model_call(self, event: BeforeModelCallEvent) -> None: | ||
| print(f"Model called for {event.agent.name}") | ||
| ``` | ||
| extend agent behavior through automatic hook and tool registration. | ||
| """ | ||
|
|
||
| from .decorator import hook | ||
| from .plugin import Plugin | ||
|
|
||
| __all__ = [ | ||
| "Plugin", | ||
| "hook", | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: would it also make sense to export tool from here? I'm thinking the devx of: from strands.plugins import Plugin, hook, tool
class MyPlugin(Plugin):
...I don't have a strong preference here though |
||
| ] | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| """Hook decorator for Plugin methods. | ||
|
|
||
| Marks methods as hook callbacks for automatic registration when the plugin | ||
| is attached to an agent. Infers event types from type hints and supports | ||
| union types for multiple events. | ||
|
|
||
| Example: | ||
| ```python | ||
| class MyPlugin(Plugin): | ||
| @hook | ||
| def on_model_call(self, event: BeforeModelCallEvent): | ||
| print(event) | ||
| ``` | ||
| """ | ||
|
|
||
| from collections.abc import Callable | ||
| from typing import Generic, cast, overload | ||
|
|
||
| from ..hooks._type_inference import infer_event_types | ||
| from ..hooks.registry import HookCallback, TEvent | ||
|
|
||
|
|
||
| class _WrappedHookCallable(HookCallback, Generic[TEvent]): | ||
| """Wrapped version of HookCallback that includes a `_hook_event_types` attribute.""" | ||
|
|
||
| _hook_event_types: list[type[TEvent]] | ||
|
|
||
|
|
||
| # Handle @hook | ||
| @overload | ||
| def hook(__func: HookCallback) -> _WrappedHookCallable: ... | ||
|
|
||
|
|
||
| # Handle @hook() | ||
| @overload | ||
| def hook() -> Callable[[HookCallback], _WrappedHookCallable]: ... | ||
|
|
||
|
|
||
| def hook( | ||
| func: HookCallback | None = None, | ||
| ) -> _WrappedHookCallable | Callable[[HookCallback], _WrappedHookCallable]: | ||
| """Mark a method as a hook callback for automatic registration. | ||
|
|
||
| Infers event type from the callback's type hint. Supports union types | ||
| for multiple events. Can be used as @hook or @hook(). | ||
|
|
||
| Args: | ||
| func: The function to decorate. | ||
|
|
||
| Returns: | ||
| The decorated function with hook metadata. | ||
|
|
||
| Raises: | ||
| ValueError: If event type cannot be inferred from type hints. | ||
| """ | ||
|
|
||
| def decorator(f: HookCallback[TEvent]) -> _WrappedHookCallable[TEvent]: | ||
| # Infer event types from type hints | ||
| event_types: list[type[TEvent]] = infer_event_types(f) | ||
|
|
||
| # Store hook metadata on the function | ||
| f_wrapped = cast(_WrappedHookCallable, f) | ||
| f_wrapped._hook_event_types = event_types | ||
|
|
||
| return f_wrapped | ||
|
|
||
| if func is None: | ||
| return decorator | ||
| return decorator(func) |
Uh oh!
There was an error while loading. Please reload this page.