power-loop 0.2.0__py3-none-any.whl

power_loop/core/events.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections import defaultdict
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Any
+
+from power_loop.contracts.events import AgentEvent, AgentEventType
+
+_logger = logging.getLogger(__name__)
+
+EventHandler = Callable[[AgentEvent], Any] | Callable[[AgentEvent], Awaitable[Any]]
+
+
+@dataclass(frozen=True)
+class _SubscribedHandler:
+    handler: EventHandler
+    priority: int
+    order: int
+
+
+class AgentEventBus:
+    """In-process pub/sub event bus.
+
+    - Subscribers are ordered by (priority, registration order).
+    - Supports both sync and async handlers.
+    - If `suppress_subscriber_errors=True`, subscriber exceptions are logged
+      and do not break the publisher.
+    """
+
+    def __init__(self, *, suppress_subscriber_errors: bool = False) -> None:
+        self._handlers: defaultdict[AgentEventType, list[_SubscribedHandler]] = defaultdict(list)
+        self._global_handlers: list[_SubscribedHandler] = []
+        self._counter = 0
+        self._suppress_subscriber_errors = suppress_subscriber_errors
+
+    def subscribe(
+        self,
+        event_type: AgentEventType | None,
+        handler: EventHandler,
+        *,
+        priority: int = 0,
+    ) -> None:
+        self._counter += 1
+        wrapped = _SubscribedHandler(handler=handler, priority=priority, order=self._counter)
+        if event_type is None:
+            self._global_handlers.append(wrapped)
+            self._global_handlers.sort(key=lambda h: (h.priority, h.order))
+            return
+        self._handlers[event_type].append(wrapped)
+        self._handlers[event_type].sort(key=lambda h: (h.priority, h.order))
+
+    def unsubscribe(self, handler: EventHandler) -> None:
+        self._global_handlers = [h for h in self._global_handlers if h.handler is not handler]
+        for etype, handlers in list(self._handlers.items()):
+            self._handlers[etype] = [h for h in handlers if h.handler is not handler]
+            if not self._handlers[etype]:
+                self._handlers.pop(etype, None)
+
+    def _invoke_handler(self, sub: _SubscribedHandler, event: AgentEvent) -> Any:
+        try:
+            return sub.handler(event)
+        except Exception:
+            if self._suppress_subscriber_errors:
+                _logger.exception(
+                    "AgentEventBus: subscriber raised (event_type=%s, handler=%s)",
+                    event.type,
+                    getattr(sub.handler, "__qualname__", sub.handler),
+                )
+                return None
+            raise
+
+    async def _await_handler_result(self, result: Any) -> None:
+        if not asyncio.iscoroutine(result):
+            return
+        try:
+            await result
+        except Exception:
+            if self._suppress_subscriber_errors:
+                _logger.exception("AgentEventBus: async subscriber raised")
+                return
+            raise
+
+    def publish(self, event: AgentEvent) -> None:
+        """Publish synchronously.
+
+        Async subscriber handlers are scheduled on the running loop when available.
+        """
+
+        handlers: list[_SubscribedHandler] = list(self._global_handlers)
+        handlers.extend(self._handlers.get(event.type, []))
+
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        for sub in handlers:
+            result = self._invoke_handler(sub, event)
+            if asyncio.iscoroutine(result):
+                if loop is not None:
+                    async def _run_coro(coro: Any) -> None:
+                        await self._await_handler_result(coro)
+
+                    loop.create_task(_run_coro(result))
+                else:
+                    asyncio.run(self._await_handler_result(result))
+
+    async def publish_async(self, event: AgentEvent) -> None:
+        """Publish asynchronously (awaits async subscribers)."""
+
+        handlers: list[_SubscribedHandler] = list(self._global_handlers)
+        handlers.extend(self._handlers.get(event.type, []))
+
+        for sub in handlers:
+            result = self._invoke_handler(sub, event)
+            await self._await_handler_result(result)
+
+
+# Process-wide default bus instance (used only when user doesn't pass one).
+DEFAULT_EVENT_BUS = AgentEventBus()
+

power_loop/core/hooks.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from power_loop.contracts.hook_contexts import BaseHookCtx
+from power_loop.contracts.hooks import HookContext, HookDirective, HookPoint, HookResult
+
+HookHandlerFn = Callable[..., Any]
+"""A hook handler callable.
+
+For **typed** hooks (the recommended style) the signature is::
+
+    def handler(ctx: SomeHookCtx) -> None | HookDirective
+
+Handlers mutate *ctx* in-place and optionally return a ``HookDirective``.
+
+Legacy handlers that receive ``HookContext`` and return
+``HookContext | HookResult | dict | None`` are still supported via
+:meth:`AgentHooks.run` / :meth:`AgentHooks.run_async`.
+"""
+
+
+@dataclass
+class _HookEntry:
+    handler: HookHandlerFn
+    order: int
+
+
+class AgentHooks:
+    """Hook manager with ordered sync/async handlers.
+
+    **Typed API** (recommended — all pipeline hook points use this):
+
+    Handlers receive a strongly-typed ``*Ctx`` dataclass, mutate it in place,
+    and optionally return ``HookDirective`` or set ``ctx.directive``.
+
+    **Legacy API** (still supported for unit-test ergonomics):
+
+    Handlers receive ``HookContext`` and may return ``HookContext``,
+    ``HookResult``, ``dict``, or ``None``.
+    """
+
+    def __init__(self) -> None:
+        self._handlers: dict[str, list[_HookEntry]] = {}
+
+    def register(self, hook_point: HookPoint | str, handler: HookHandlerFn, *, order: int = 0) -> None:
+        key = str(hook_point)
+        self._handlers.setdefault(key, []).append(_HookEntry(handler=handler, order=order))
+        self._handlers[key].sort(key=lambda e: e.order)
+
+    def clear(self, hook_point: HookPoint | str | None = None) -> None:
+        if hook_point is None:
+            self._handlers.clear()
+            return
+        self._handlers.pop(str(hook_point), None)
+
+    # ── Legacy dict-based API ──
+
+    @staticmethod
+    def _apply(context: HookContext, directive: HookDirective, result: Any) -> tuple[HookContext, HookDirective]:
+        if result is None:
+            return context, directive
+        if isinstance(result, HookResult):
+            return result.context, result.directive
+        if isinstance(result, HookContext):
+            return result, directive
+        if isinstance(result, dict):
+            context.values = result
+            return context, directive
+        return context, directive
+
+    def run(self, hook_point: HookPoint | str, context: HookContext) -> HookResult:
+        key = str(hook_point)
+        directive = HookDirective.CONTINUE
+        for entry in self._handlers.get(key, []):
+            result = entry.handler(context)
+            context, directive = self._apply(context, directive, result)
+        return HookResult(context=context, directive=directive)
+
+    async def run_async(self, hook_point: HookPoint | str, context: HookContext) -> HookResult:
+        key = str(hook_point)
+        directive = HookDirective.CONTINUE
+        for entry in self._handlers.get(key, []):
+            result = entry.handler(context)
+            if hasattr(result, "__await__"):
+                result = await result
+            context, directive = self._apply(context, directive, result)
+        return HookResult(context=context, directive=directive)
+
+    # ── Typed context API (all pipeline hooks use this path) ──
+
+    @staticmethod
+    def _apply_typed(ctx: BaseHookCtx, result: Any) -> None:
+        """Apply handler return value to the typed context."""
+        if result is None:
+            return
+        if isinstance(result, HookDirective):
+            ctx.directive = result
+            return
+        if isinstance(result, HookResult):
+            ctx.directive = result.directive
+            return
+
+    def run_typed(self, hook_point: HookPoint | str, ctx: BaseHookCtx) -> None:
+        """Run handlers with a typed context.  Handlers mutate *ctx* in place."""
+        for entry in self._handlers.get(str(hook_point), []):
+            result = entry.handler(ctx)
+            self._apply_typed(ctx, result)
+
+    async def run_typed_async(self, hook_point: HookPoint | str, ctx: BaseHookCtx) -> None:
+        """Async version of :meth:`run_typed`."""
+        for entry in self._handlers.get(str(hook_point), []):
+            result = entry.handler(ctx)
+            if hasattr(result, "__await__"):
+                result = await result
+            self._apply_typed(ctx, result)
+
+
+DEFAULT_HOOKS = AgentHooks()
+

power_loop/core/phase.py

@@ -0,0 +1,217 @@
+"""Declarative ``@phase`` decorator for agent pipeline methods.
+
+The decorator automatically wraps a pipeline method with:
+- **before** hook (``hooks.run_async``) — may modify context values or return a directive
+- **after** hook (``hooks.run_async``) — may modify output or return a directive
+- **error** hook (optional) — fires on exception, may suppress/retry
+- **event publishing** — emits start/complete events on the bus
+
+This keeps each pipeline method focused on pure business logic while
+cross-cutting concerns (hooks, events, directives) are handled declaratively.
+
+Usage::
+
+    class MyPipeline(AgentPipeline):
+        @phase(
+            before=HookPoint.LLM_BEFORE,
+            after=HookPoint.LLM_AFTER,
+            start_event=AgentEventType.STREAM_STARTED,
+            end_event=AgentEventType.STREAM_COMPLETED,
+        )
+        async def call_llm(self, ctx: PhaseContext) -> Any:
+            return await self.llm.complete(...)
+"""
+from __future__ import annotations
+
+import functools
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from power_loop.contracts.events import AgentEvent, AgentEventType
+from power_loop.contracts.hooks import HookContext, HookDirective, HookPoint
+
+
+@dataclass
+class PhaseContext:
+    """Mutable bag of values flowing through a phase.
+
+    Pipeline methods receive this; the ``@phase`` decorator populates it
+    from the before-hook result and passes it to the method.
+    """
+    values: dict[str, Any] = field(default_factory=dict)
+    round_index: int = 0
+    session_id: str | None = None
+
+
+@dataclass
+class PhaseResult:
+    """Return value of a ``@phase``-decorated method.
+
+    The decorator merges hook directives and the method's raw output into this.
+    The caller only needs to check ``directive`` and read ``output`` / ``values``.
+    """
+    output: Any = None
+    directive: HookDirective = HookDirective.CONTINUE
+    values: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def should_break(self) -> bool:
+        return self.directive == HookDirective.BREAK
+
+    @property
+    def should_skip(self) -> bool:
+        return self.directive == HookDirective.SKIP
+
+    @property
+    def is_short_circuit(self) -> bool:
+        return self.directive == HookDirective.SHORT_CIRCUIT
+
+
+def phase(
+    *,
+    before: HookPoint | None = None,
+    after: HookPoint | None = None,
+    error: HookPoint | None = None,
+    start_event: AgentEventType | None = None,
+    end_event: AgentEventType | None = None,
+) -> Callable:
+    """Decorator that wraps a pipeline method with hooks, events, and directive handling.
+
+    The decorated method signature must be::
+
+        async def method(self, ctx: PhaseContext) -> Any
+
+    The ``self`` must be an ``AgentPipeline`` instance (has ``.hooks``, ``.bus``,
+    ``.session_id`` attributes).
+
+    Execution flow:
+
+    1. Run **before** hook → get ``HookResult``
+       - If directive is ``SKIP``: return ``PhaseResult(directive=SKIP, values=...)`` immediately
+       - If directive is ``SHORT_CIRCUIT``: return ``PhaseResult(output=values["output"], ...)``
+       - Otherwise: merge modified values back into ``ctx``
+    2. Publish **start_event** (if set)
+    3. Call the actual method
+    4. Publish **end_event** (if set)
+    5. Run **after** hook → get ``HookResult``
+       - If directive is ``BREAK``: set ``PhaseResult.directive = BREAK``
+       - May replace output via ``values["output"]``
+    6. Return ``PhaseResult``
+
+    On exception:
+    - If **error** hook is set, run it.
+      - ``SKIP`` → swallow error, use ``values["output"]`` as fallback
+      - ``SHORT_CIRCUIT`` → retry the method once
+      - Otherwise → re-package as ``PhaseResult`` with error info
+    """
+
+    def decorator(method: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[PhaseResult]]:
+
+        @functools.wraps(method)
+        async def wrapper(self: Any, ctx: PhaseContext) -> PhaseResult:
+            hooks = self.hooks
+            bus = self.bus
+
+            event_meta: dict[str, Any] = {
+                "session_id": ctx.session_id,
+                "round_index": ctx.round_index,
+            }
+
+            # ── before hook ──
+            if before is not None:
+                hr = await hooks.run_async(before, context=HookContext(values=dict(ctx.values)))
+                # Merge hook-modified values back
+                ctx.values.update(hr.context.values)
+
+                if hr.directive == HookDirective.SKIP:
+                    return PhaseResult(
+                        output=hr.context.values.get("output"),
+                        directive=HookDirective.SKIP,
+                        values=hr.context.values,
+                    )
+                if hr.directive == HookDirective.SHORT_CIRCUIT:
+                    return PhaseResult(
+                        output=hr.context.values.get("output"),
+                        directive=HookDirective.SHORT_CIRCUIT,
+                        values=hr.context.values,
+                    )
+                if hr.directive == HookDirective.BREAK:
+                    return PhaseResult(
+                        output=hr.context.values.get("output"),
+                        directive=HookDirective.BREAK,
+                        values=hr.context.values,
+                    )
+
+            # ── start event ──
+            if start_event is not None:
+                bus.publish(AgentEvent(
+                    type=start_event,
+                    payload=ctx.values.get("event_payload", {}),
+                    **event_meta,
+                ))
+
+            # ── execute method ──
+            raw_output: Any = None
+            try:
+                raw_output = await method(self, ctx)
+            except Exception as exc:
+                if error is not None:
+                    err_hr = await hooks.run_async(
+                        error,
+                        context=HookContext(values={
+                            **ctx.values,
+                            "error": exc,
+                            "error_message": str(exc),
+                        }),
+                    )
+                    if err_hr.directive == HookDirective.SKIP:
+                        raw_output = err_hr.context.values.get("output", f"Error: {exc}")
+                    elif err_hr.directive == HookDirective.SHORT_CIRCUIT:
+                        # Retry once
+                        try:
+                            raw_output = await method(self, ctx)
+                        except Exception as retry_exc:
+                            raw_output = f"Error (retry failed): {retry_exc}"
+                            return PhaseResult(output=raw_output, values={**ctx.values, "failed": True})
+                    else:
+                        return PhaseResult(
+                            output=f"Error: {exc}",
+                            values={**ctx.values, "error": exc, "failed": True},
+                        )
+                else:
+                    raise
+
+            # ── end event ──
+            if end_event is not None:
+                bus.publish(AgentEvent(
+                    type=end_event,
+                    payload=ctx.values.get("event_payload", {}),
+                    **event_meta,
+                ))
+
+            # ── after hook ──
+            result_directive = HookDirective.CONTINUE
+            result_values = dict(ctx.values)
+            if after is not None:
+                after_ctx_vals = {**ctx.values, "output": raw_output}
+                ahr = await hooks.run_async(after, context=HookContext(values=after_ctx_vals))
+                result_values = ahr.context.values
+                result_directive = ahr.directive
+                # after hook may replace output
+                if "output" in ahr.context.values:
+                    raw_output = ahr.context.values["output"]
+
+            return PhaseResult(
+                output=raw_output,
+                directive=result_directive,
+                values=result_values,
+            )
+
+        # Preserve the original hook metadata for introspection
+        wrapper._phase_before = before  # type: ignore[attr-defined]
+        wrapper._phase_after = after  # type: ignore[attr-defined]
+        wrapper._phase_error = error  # type: ignore[attr-defined]
+        return wrapper
+
+    return decorator