diff --git a/mellea/plugins/base.py b/mellea/plugins/base.py index 357e2a515..a68244835 100644 --- a/mellea/plugins/base.py +++ b/mellea/plugins/base.py @@ -16,7 +16,12 @@ @dataclass(frozen=True) class PluginMeta: - """Metadata attached to Plugin subclasses.""" + """Metadata attached to Plugin subclasses. + + Args: + name: Unique identifier for the plugin. + priority: Execution priority — lower numbers execute first. + """ name: str priority: int = 50 @@ -121,7 +126,14 @@ async def __aexit__( class PluginViolationError(Exception): - """Raised when a plugin blocks execution in enforce mode.""" + """Raised when a plugin blocks execution in enforce mode. + + Args: + hook_type: The hook point where the violation occurred. + reason: Human-readable reason for the violation. + code: Machine-readable violation code. + plugin_name: Name of the plugin that raised the violation. + """ def __init__( # noqa: D107 self, hook_type: str, reason: str, code: str = "", plugin_name: str = "" diff --git a/mellea/plugins/context.py b/mellea/plugins/context.py index bed622794..2243083a8 100644 --- a/mellea/plugins/context.py +++ b/mellea/plugins/context.py @@ -23,7 +23,12 @@ def build_global_context(*, backend: Backend | None = None, **extra_fields: Any) Hook-specific data (context, session, action, etc.) belongs on the typed payload, not here. - Returns ``None`` if ContextForge is not installed. + Args: + backend: Optional backend whose ``model_id`` is added to the context. + **extra_fields: Additional key-value pairs merged into the context state. + + Returns: + A ``GlobalContext`` instance, or ``None`` if ContextForge is not installed. """ if not _HAS_PLUGIN_FRAMEWORK: return None diff --git a/mellea/plugins/decorators.py b/mellea/plugins/decorators.py index 27789c647..43bf61af3 100644 --- a/mellea/plugins/decorators.py +++ b/mellea/plugins/decorators.py @@ -10,7 +10,13 @@ @dataclass(frozen=True) class HookMeta: - """Metadata attached by the @hook decorator.""" + """Metadata attached by the @hook decorator. + + Args: + hook_type: The hook point name (e.g., ``"generation_pre_call"``). + mode: Execution mode for the hook handler. + priority: Execution priority — lower numbers execute first. + """ hook_type: str mode: PluginMode = PluginMode.SEQUENTIAL @@ -31,6 +37,12 @@ def hook( ``PluginMode.AUDIT``, or ``PluginMode.FIRE_AND_FORGET``. priority: Lower numbers execute first. For methods on a ``Plugin`` subclass, falls back to the class-level priority, then 50. For standalone functions, defaults to 50. + + Returns: + A decorator that attaches ``HookMeta`` to the decorated function. + + Raises: + TypeError: If the decorated function is not ``async def``. """ def decorator(fn: Callable) -> Callable: diff --git a/mellea/plugins/manager.py b/mellea/plugins/manager.py index 25598f520..978b7ae99 100644 --- a/mellea/plugins/manager.py +++ b/mellea/plugins/manager.py @@ -37,6 +37,13 @@ def has_plugins(hook_type: HookType | None = None) -> bool: When ``hook_type`` is provided, also checks whether any plugin has registered a handler for that specific hook, enabling callers to skip payload construction entirely when no plugin subscribes. + + Args: + hook_type: Optional hook type to check for registered handlers. + + Returns: + ``True`` if plugins are enabled and (when ``hook_type`` is given) + at least one plugin subscribes to that hook. """ if not _plugins_enabled or _plugin_manager is None: return False @@ -46,12 +53,23 @@ def has_plugins(hook_type: HookType | None = None) -> bool: def get_plugin_manager() -> Any | None: - """Returns the initialized PluginManager, or ``None`` if plugins are not configured.""" + """Return the initialized PluginManager, or ``None`` if plugins are not configured. + + Returns: + The singleton ``PluginManager`` instance, or ``None``. + """ return _plugin_manager def ensure_plugin_manager() -> Any: - """Lazily initialize the PluginManager if not already created.""" + """Lazily initialize the PluginManager if not already created. + + Returns: + The singleton ``PluginManager`` instance. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. + """ global _plugin_manager, _plugins_enabled if not _HAS_PLUGIN_FRAMEWORK: @@ -87,6 +105,12 @@ async def initialize_plugins( Args: config_path: Optional path to a YAML plugin configuration file. timeout: Maximum execution time per plugin in seconds. + + Returns: + The initialized ``PluginManager`` instance. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. """ global _plugin_manager, _plugins_enabled @@ -122,12 +146,21 @@ async def shutdown_plugins() -> None: def track_session_plugin(session_id: str, plugin_name: str) -> None: - """Track a plugin as belonging to a session for later deregistration.""" + """Track a plugin as belonging to a session for later deregistration. + + Args: + session_id: Identifier for the session that owns the plugin. + plugin_name: Registered name of the plugin. + """ _session_tags.setdefault(session_id, set()).add(plugin_name) def deregister_session_plugins(session_id: str) -> None: - """Deregister all plugins scoped to the given session.""" + """Deregister all plugins scoped to the given session. + + Args: + session_id: Identifier for the session whose plugins should be removed. + """ if not _plugins_enabled or _plugin_manager is None: return @@ -158,6 +191,20 @@ async def invoke_hook( 1. ``_plugins_enabled`` boolean — single pointer dereference 2. ``has_hooks_for(hook_type)`` — skips when no plugin subscribes 3. Returns immediately when either guard fails + + Args: + hook_type: The hook point to invoke. + payload: The immutable payload to pass to plugin handlers. + backend: Optional backend for building the global context. + **context_fields: Additional fields passed to ``build_global_context``. + + Returns: + A ``(result, payload)`` tuple where *result* is the ``PluginResult`` + (or ``None`` when no plugins ran) and *payload* is the + possibly-modified payload. + + Raises: + PluginViolationError: If a plugin blocks execution. """ if not _plugins_enabled or _plugin_manager is None: return None, payload diff --git a/mellea/plugins/pluginset.py b/mellea/plugins/pluginset.py index d485ee65a..099bd5d74 100644 --- a/mellea/plugins/pluginset.py +++ b/mellea/plugins/pluginset.py @@ -23,6 +23,11 @@ class PluginSet: # or async async with PluginSet("guards", [safety_hook]): result, ctx = await ainstruct("Generate code", ctx, backend) + + Args: + name: Human-readable label for this group. + items: Hook functions, plugin instances, or nested ``PluginSet`` instances. + priority: Optional priority override applied to all items in this set. """ def __init__( # noqa: D107 diff --git a/mellea/plugins/registry.py b/mellea/plugins/registry.py index 181044ec6..a43ea17ca 100644 --- a/mellea/plugins/registry.py +++ b/mellea/plugins/registry.py @@ -52,7 +52,7 @@ def modify(payload: Any, **field_updates: Any) -> Any: listed in the hook's ``HookPayloadPolicy.writable_fields`` will be accepted by the framework; changes to read-only fields are silently discarded. - Mirrors :func:`block` for the modification case:: + Mirrors ``block()`` for the modification case:: # instead of: modified = payload.model_copy(update={"model_output": new_mot}) @@ -64,6 +64,12 @@ def modify(payload: Any, **field_updates: Any) -> Any: Args: payload: The original (frozen) payload received by the hook. **field_updates: Fields to update on the payload copy. + + Returns: + A ``PluginResult`` with ``continue_processing=True`` and the modified payload. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. """ if not _HAS_PLUGIN_FRAMEWORK: raise ImportError( @@ -90,6 +96,12 @@ def block( code: Machine-readable violation code. description: Longer description (defaults to ``reason``). details: Additional structured details. + + Returns: + A ``PluginResult`` with ``continue_processing=False`` and a violation. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. """ if not _HAS_PLUGIN_FRAMEWORK: raise ImportError( @@ -119,6 +131,16 @@ def register( Accepts standalone ``@hook`` functions, ``@plugin``-decorated class instances, ``MelleaPlugin`` instances, ``PluginSet`` instances, or lists thereof. + + Args: + items: One or more plugins to register — standalone ``@hook`` functions, + ``@plugin``-decorated class instances, ``MelleaPlugin`` instances, + ``PluginSet`` instances, or a list of any combination. + session_id: Optional session identifier. When provided, the plugins are + scoped to that session and automatically deregistered on session cleanup. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. """ if not _HAS_PLUGIN_FRAMEWORK: raise ImportError( @@ -366,7 +388,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: class _PluginScope: - """Context manager returned by :func:`plugin_scope`. + """Context manager returned by ``plugin_scope()``. Supports both synchronous and asynchronous ``with`` statements. """ @@ -453,11 +475,18 @@ def unregister( ) -> None: """Unregister globally-registered plugins. - Accepts the same items as :func:`register`: standalone ``@hook``-decorated + Accepts the same items as ``register()``: standalone ``@hook``-decorated functions, ``Plugin`` subclass instances, ``MelleaPlugin`` instances, ``PluginSet`` instances, or lists thereof. Silently ignores items that are not currently registered. + + Args: + items: One or more plugins to unregister — same types accepted by + ``register()``. + + Raises: + ImportError: If the ContextForge plugin framework is not installed. """ if not _HAS_PLUGIN_FRAMEWORK: raise ImportError( @@ -485,9 +514,9 @@ def unregister( def plugin_scope(*items: Callable | Any | PluginSet) -> _PluginScope: """Return a context manager that temporarily registers plugins for a block of code. - Accepts the same items as :func:`register`: standalone ``@hook``-decorated + Accepts the same items as ``register()``: standalone ``@hook``-decorated functions, ``@plugin``-decorated class instances, ``MelleaPlugin`` instances, - and :class:`~mellea.plugins.PluginSet` instances — or any mix thereof. + and ``PluginSet`` instances — or any mix thereof. Supports both synchronous and asynchronous ``with`` statements:: diff --git a/mellea/plugins/types.py b/mellea/plugins/types.py index 107c0a97d..d64eb4876 100644 --- a/mellea/plugins/types.py +++ b/mellea/plugins/types.py @@ -28,7 +28,11 @@ class PluginMode(str, Enum): class HookType(str, Enum): - """All Mellea hook types.""" + """All Mellea hook types. + + Hook types are organized by lifecycle stage: session, component, + generation pipeline, validation, sampling pipeline, and tool execution. + """ # Session Lifecycle SESSION_PRE_INIT = "session_pre_init" diff --git a/mellea/telemetry/tracing.py b/mellea/telemetry/tracing.py index cab2c2bc6..f5cc03c7e 100644 --- a/mellea/telemetry/tracing.py +++ b/mellea/telemetry/tracing.py @@ -116,11 +116,11 @@ def trace_application(name: str, **attributes: Any) -> Generator[Any, None, None """Create an application trace span if application tracing is enabled. Args: - name: Name of the span - **attributes: Additional attributes to add to the span + name: Name of the span. + **attributes: Additional attributes to add to the span. Yields: - The span object if tracing is enabled, otherwise a no-op context manager + The span object if tracing is enabled, otherwise ``None``. """ if _TRACE_APPLICATION_ENABLED and _application_tracer is not None: with _application_tracer.start_as_current_span(name) as span: # type: ignore @@ -139,11 +139,11 @@ def trace_backend(name: str, **attributes: Any) -> Generator[Any, None, None]: Follows Gen-AI semantic conventions for LLM operations. Args: - name: Name of the span - **attributes: Additional attributes to add to the span + name: Name of the span. + **attributes: Additional attributes to add to the span. Yields: - The span object if tracing is enabled, otherwise a no-op context manager + The span object if tracing is enabled, otherwise ``None``. """ if _TRACE_BACKEND_ENABLED and _backend_tracer is not None: with _backend_tracer.start_as_current_span(name) as span: # type: ignore