feat: hook system and plugin support for Mellea

docs/docs/docs.json

@@ -54,7 +54,8 @@
               "core-concept/tuning",
               "core-concept/adapters",
               "core-concept/alora",
-              "core-concept/interoperability"
+              "core-concept/interoperability",
+              "core-concept/plugins"
             ]
           }
         ]

docs/examples/plugins/class_plugin.py

@@ -0,0 +1,128 @@
+# pytest: ollama, llm
+#
+# Class-based plugin — group related hooks in a single Plugin subclass.
+#
+# This example creates a PII protection plugin that:
+#   1. Blocks input containing SSN patterns before component execution
+#   2. Scans LLM output for SSN patterns after generation (observe-only)
+#
+# Run:
+#   uv run python docs/examples/plugins/class_plugin.py
+
+import logging
+import re
+import sys
+
+from mellea import start_session
+from mellea.plugins import HookType, Plugin, PluginViolationError, block, hook, register
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%S",
+)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("fancy_logger").setLevel(logging.ERROR)
+log = logging.getLogger("class_plugin")
+
+
+class PIIRedactor(Plugin, name="pii-redactor", priority=5):
+    """Redacts PII patterns from both input and output.
+
+    .. warning:: Shared mutable state
+        ``redaction_count`` is shared across all hook invocations.  This is
+        safe today because all hooks run on the same ``asyncio`` event loop,
+        but would require a lock or ``contextvars`` if hooks ever execute in
+        parallel threads.
+    """
+
+    def __init__(self, patterns: list[str] | None = None):
+        self.patterns = patterns or [
+            r"\d{3}-\d{2}-\d{4}",  # SSN
+            r"\b\d{16}\b",  # credit card (simplified)
+        ]
+        self.redaction_count = 0
+
+    @hook(HookType.COMPONENT_PRE_EXECUTE)
+    async def reject_pii_input(self, payload, ctx):
+        """Block component execution if the action contains PII patterns."""
+        if payload.component_type != "Instruction":
+            return
+        original = (
+            str(payload.action._description) if payload.action._description else ""
+        )
+        if self._contains_pii(original):
+            log.warning("[pii-redactor] PII detected in component action — blocking")
+            self.redaction_count += 1
+            return block(
+                "Input contains PII patterns that must be removed before processing",
+                code="PII_INPUT_DETECTED",
+            )
+        log.info("[pii-redactor] no PII found in input")
+
+    @hook(HookType.GENERATION_POST_CALL)
+    async def scan_output(self, payload, ctx):
+        """Scan LLM output for PII and log a warning if detected.
+
+        ``generation_post_call`` is observe-only — plugins cannot modify the
+        ``model_output``.  This hook therefore only inspects the output and
+        records a warning for downstream monitoring/alerting.
+        """
+        mot_value = getattr(payload.model_output, "value", None)
+        if mot_value is None:
+            log.info("[pii-redactor] output not yet computed — skipping output scan")
+            return
+        original = str(mot_value)
+        if self._contains_pii(original):
+            log.warning("[pii-redactor] PII detected in LLM output (observe-only)")
+            self.redaction_count += 1
+        else:
+            log.info("[pii-redactor] no PII found in output")
+
+    def _contains_pii(self, text: str) -> bool:
+        return any(re.search(p, text) for p in self.patterns)
+
+
+# Create an instance and register it globally
+redactor = PIIRedactor()
+register(redactor)
+
+if __name__ == "__main__":
+    log.info("--- Class-based Plugin example (PII Redactor) ---")
+    log.info("")
+
+    with start_session() as m:
+        log.info("Session started (id=%s)", m.id)
+        log.info("")
+
+        # Request 1: contains an SSN — the input hook blocks execution.
+        log.info("Request 1: input with PII (should be blocked)")
+        try:
+            m.instruct(
+                "Summarize this customer record: "
+                "Name: Jane Doe, SSN: 123-45-6789, Status: Active"
+            )
+        except PluginViolationError as e:
+            log.info(
+                "Blocked as expected on %s: [%s] %s", e.hook_type, e.code, e.reason
+            )
+        log.info("")
+
+        # Request 2: clean input — no PII, so it reaches the LLM.
+        # If the LLM output contains PII, scan_output logs a warning (observe-only).
+        log.info("Request 2: clean input (should succeed)")
+        try:
+            result = m.instruct("Name the three primary colors.")
+            log.info("Result: %s", result)
+        except PluginViolationError as e:
+            log.warning(
+                "Execution blocked on %s: [%s] %s (plugin=%s)",
+                e.hook_type,
+                e.code,
+                e.reason,
+                e.plugin_name,
+            )
+            sys.exit(1)
+
+        log.info("")
+        log.info("Total PII detections: %d", redactor.redaction_count)

docs/examples/plugins/execution_modes.py

@@ -0,0 +1,130 @@
+# pytest: ollama, llm
+#
+# Execution modes — all five PluginMode values side by side.
+#
+# This example registers five hooks on the same hook type
+# (COMPONENT_PRE_EXECUTE), each using a different execution mode.
+# It demonstrates:
+#
+#   1. SEQUENTIAL      — serial, can block + modify
+#   2. TRANSFORM       — serial, can modify only (blocks suppressed)
+#   3. AUDIT           — serial, observe-only (modifications discarded, blocks logged)
+#   4. CONCURRENT      — parallel, can block only (modifications discarded)
+#   5. FIRE_AND_FORGET — background, observe-only (result ignored)
+#
+# Execution order: SEQUENTIAL → TRANSFORM → AUDIT → CONCURRENT → FIRE_AND_FORGET
+#
+# Run:
+#   uv run python docs/examples/plugins/execution_modes.py
+
+import logging
+
+from mellea import start_session
+from mellea.plugins import (
+    HookType,
+    PluginMode,
+    PluginViolationError,
+    block,
+    hook,
+    modify,
+    plugin_scope,
+)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%S",
+)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("fancy_logger").setLevel(logging.ERROR)
+log = logging.getLogger("execution_modes")
+
+
+# --- Hook 1: SEQUENTIAL (priority=10) ---
+# Serial, chained execution. Can block the pipeline and modify writable
+# payload fields. Each hook receives the payload from the prior one.
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.SEQUENTIAL, priority=10)
+async def sequential_hook(payload, ctx):
+    """Sequential hook — can block + modify, runs inline in priority order."""
+    log.info("[SEQUENTIAL      p=10] component=%s", payload.component_type)
+
+
+# --- Hook 2: TRANSFORM (priority=20) ---
+# Serial, chained execution after all SEQUENTIAL hooks. Can modify writable
+# payload fields but CANNOT block — blocking results are suppressed with a
+# warning. Ideal for data transformation (PII redaction, prompt rewriting).
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.TRANSFORM, priority=20)
+async def transform_hook(payload, ctx):
+    """Transform hook — can modify but cannot block."""
+    log.info("[TRANSFORM       p=20] enriching model_options")
+    opts = dict(payload.model_options or {})
+    opts.setdefault("temperature", 0.7)
+    return modify(payload, model_options=opts)
+
+
+# --- Hook 3: AUDIT (priority=30) ---
+# Serial execution after TRANSFORM. Observe-only: payload modifications are
+# discarded and violations are logged but do NOT block. Use for monitoring,
+# metrics, and gradual policy rollout.
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.AUDIT, priority=30)
+async def audit_hook(payload, ctx):
+    """Audit hook — observe-only; violations logged but not enforced."""
+    log.info("[AUDIT           p=30] would block, but audit mode only logs")
+    return block("Audit-mode violation: for monitoring only", code="AUDIT_001")
+
+
+# --- Hook 4: CONCURRENT (priority=40) ---
+# Dispatched in parallel after AUDIT. Can block the pipeline (fail-fast on
+# first blocking result) but payload modifications are discarded to avoid
+# non-deterministic last-writer-wins races.
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.CONCURRENT, priority=40)
+async def concurrent_hook(payload, ctx):
+    """Concurrent hook — can block but cannot modify, runs in parallel."""
+    log.info("[CONCURRENT      p=40] component=%s", payload.component_type)
+
+
+# --- Hook 5: FIRE_AND_FORGET (priority=50) ---
+# Dispatched via asyncio.create_task() after all other phases. Receives a
+# copy-on-write snapshot of the payload. Cannot modify payloads or block
+# execution. Any exceptions are logged but do not propagate.
+# The log line may appear *after* the main result is printed.
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.FIRE_AND_FORGET, priority=50)
+async def fire_and_forget_hook(payload, ctx):
+    """Fire-and-forget hook — runs in background, never blocks."""
+    log.info("[FIRE_AND_FORGET p=50] logging in the background")
+
+
+if __name__ == "__main__":
+    log.info("--- Execution modes example ---")
+    log.info("")
+
+    with start_session() as m:
+        with plugin_scope(
+            sequential_hook,
+            transform_hook,
+            audit_hook,
+            concurrent_hook,
+            fire_and_forget_hook,
+        ):
+            try:
+                result = m.instruct("Name the four seasons.")
+                log.info("")
+                log.info("Result: %s", result)
+            except PluginViolationError as e:
+                log.error("Blocked: %s", e)
+
+    log.info("")
+    log.info(
+        "Note: the FIRE_AND_FORGET log may have appeared after the result "
+        "— that is expected behavior."
+    )

docs/examples/plugins/payload_modification.py

@@ -0,0 +1,106 @@
+# pytest: ollama, llm
+#
+# Payload modification — how to modify payloads in hooks.
+#
+# This example demonstrates:
+#   1. Using modify() to change writable payload fields
+#   2. Using model_copy(update={...}) directly for fine-grained control
+#   3. What happens when you try to modify a non-writable field (silently discarded)
+#
+# Run:
+#   uv run python docs/examples/plugins/payload_modification.py
+
+import logging
+
+from mellea import start_session
+from mellea.plugins import (
+    HookType,
+    PluginMode,
+    PluginResult,
+    hook,
+    modify,
+    plugin_scope,
+)
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    datefmt="%Y-%m-%dT%H:%M:%S",
+)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("fancy_logger").setLevel(logging.ERROR)
+log = logging.getLogger("payload_modification")
+
+
+# ---------------------------------------------------------------------------
+# Hook 1: Inject a max_tokens cap via modify() helper
+#
+# generation_pre_call writable fields include: model_options, format, tool_calls
+# ---------------------------------------------------------------------------
+
+
+@hook(HookType.GENERATION_PRE_CALL, mode=PluginMode.SEQUENTIAL, priority=10)
+async def cap_max_tokens(payload, ctx):
+    """Cap max_tokens to 256 on every generation call."""
+    opts = dict(payload.model_options or {})
+    if opts.get("max_tokens", float("inf")) > 256:
+        log.info("[cap_max_tokens] capping max_tokens to 256")
+        opts["max_tokens"] = 256
+        return modify(payload, model_options=opts)
+    log.info("[cap_max_tokens] max_tokens already within cap")
+
+
+# ---------------------------------------------------------------------------
+# Hook 2: Inject default model options via modify()
+#
+# component_pre_execute writable fields include: requirements, model_options, ...
+# This shows model_copy(update={...}) for fine-grained control.
+# ---------------------------------------------------------------------------
+
+
+@hook(HookType.COMPONENT_PRE_EXECUTE, mode=PluginMode.SEQUENTIAL, priority=10)
+async def inject_default_options(payload, ctx):
+    """Ensure a default temperature is set on every component execution."""
+    opts = dict(payload.model_options or {})
+    if "temperature" not in opts:
+        log.info("[inject_default_options] setting default temperature=0.7")
+        opts["temperature"] = 0.7
+        return modify(payload, model_options=opts)
+    log.info("[inject_default_options] temperature already set")
+
+
+# ---------------------------------------------------------------------------
+# Hook 3: Attempt to modify a non-writable field (observe it is discarded)
+#
+# generation_pre_call does NOT include 'action' or 'context' as writable.
+# This hook tries to modify 'context' — the change will be silently discarded
+# by the payload policy enforcement, and the original context will be used.
+# ---------------------------------------------------------------------------
+
+
+@hook(HookType.GENERATION_PRE_CALL, mode=PluginMode.SEQUENTIAL, priority=20)
+async def attempt_non_writable(payload, ctx):
+    """Try to modify a non-writable field — change will be silently discarded."""
+    log.info("[attempt_non_writable] attempting to modify 'hook' (non-writable)")
+    # This modification will be filtered out by the payload policy
+    modified = payload.model_copy(update={"hook": "tampered"})
+    return PluginResult(continue_processing=True, modified_payload=modified)
+
+
+if __name__ == "__main__":
+    log.info("--- Payload modification example ---")
+    log.info("")
+
+    with start_session() as m:
+        with plugin_scope(cap_max_tokens, inject_default_options, attempt_non_writable):
+            result = m.instruct(
+                "Summarize the benefits of open-source software in one sentence."
+            )
+            log.info("")
+            log.info("Result: %s", result)
+
+    log.info("")
+    log.info(
+        "Note: the 'hook' field modification in attempt_non_writable was silently "
+        "discarded by the payload policy — only writable fields are accepted."
+    )