penguiflow 2.1.0__tar.gz → 2.2.0__tar.gz

{penguiflow-2.1.0 → penguiflow-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: penguiflow
-Version: 2.1.0
+Version: 2.2.0
 Summary: Async agent orchestration primitives.
 Author: PenguiFlow Team
 License: MIT License
@@ -36,11 +36,14 @@ Requires-Dist: pytest>=7.4; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
 Requires-Dist: coverage[toml]>=7.0; extra == "dev"
+Requires-Dist: hypothesis>=6.103; extra == "dev"
 Requires-Dist: ruff>=0.2; extra == "dev"
-Requires-Dist: fastapi>=0.110; extra == "dev"
+Requires-Dist: fastapi>=0.118; extra == "dev"
 Requires-Dist: httpx>=0.27; extra == "dev"
 Provides-Extra: a2a-server
-Requires-Dist: fastapi>=0.110; extra == "a2a-server"
+Requires-Dist: fastapi>=0.118; extra == "a2a-server"
+Provides-Extra: planner
+Requires-Dist: litellm>=1.77.3; extra == "planner"
 Dynamic: license-file
 
 # PenguiFlow 🐧❄️
@@ -56,6 +59,9 @@ Dynamic: license-file
   <a href="https://github.com/penguiflow/penguiflow">
     <img src="https://img.shields.io/badge/coverage-85%25-brightgreen" alt="Coverage">
   </a>
+  <a href="https://nightly.link/penguiflow/penguiflow/workflows/benchmarks/main/benchmarks.json.zip">
+    <img src="https://img.shields.io/badge/benchmarks-latest-orange" alt="Benchmarks">
+  </a>
   <a href="https://pypi.org/project/penguiflow/">
     <img src="https://img.shields.io/pypi/v/penguiflow.svg" alt="PyPI version">
   </a>
@@ -97,6 +103,23 @@ It provides:
 Built on pure `asyncio` (no threads), PenguiFlow is small, predictable, and repo-agnostic.
 Product repos only define **their models + node functions** — the core stays dependency-light.
 
+## Gold Standard Scorecard
+
+| Area | Metric | Target | Current |
+| --- | --- | --- | --- |
+| Hop overhead | µs per hop | ≤ 500 | 398 |
+| Streaming order | gaps/dupes | 0 | 0 |
+| Cancel leakage | orphan tasks | 0 | 0 |
+| Coverage | lines | ≥85% | 87% |
+| Deps | count | ≤2 | 2 |
+| Import time | ms | ≤220 | 203 |
+
+## 📑 Core Behavior Spec
+
+* [Core Behavior Spec](docs/core_behavior_spec.md) — single-page rundown of ordering,
+  streaming, cancellation, deadline, and fan-in invariants with pointers to regression
+  tests.
+
 ---
 
 ## ✨ Why PenguiFlow?
@@ -346,6 +369,70 @@ The new `penguiflow.testkit` module keeps unit tests tiny:
 The harness is covered by `tests/test_testkit.py` and demonstrated in
 `examples/testkit_demo/`.
 
+### JSON-only ReAct planner (Phase A)
+
+Phase A introduces a lightweight planner loop that keeps PenguiFlow typed and
+deterministic:
+
+* `penguiflow.catalog.NodeSpec` + `build_catalog` turn registered nodes into
+  tool descriptors with JSON Schemas derived from your Pydantic models.
+* `penguiflow.planner.ReactPlanner` drives a JSON-only ReAct loop over those
+  descriptors, validating every LLM action with Pydantic and replaying invalid
+  steps to request corrections.
+* LiteLLM stays optional—install `penguiflow[planner]` or inject a custom
+  `llm_client` for deterministic/offline runs.
+
+See `examples/react_minimal/` for a stubbed end-to-end run.
+
+### Trajectory summarisation & pause/resume (Phase B)
+
+Phase B adds the tools you need for longer-running, approval-driven flows:
+
+* **Token-aware summaries** — `Trajectory.compress()` keeps a compact state and
+  the planner can route summaries through a cheaper `summarizer_llm` before
+  asking for the next action.
+* **`PlannerPause` contract** — nodes can call `await ctx.pause(...)` to return a
+  typed pause payload. Resume the run later with `ReactPlanner.resume(token, user_input=...)`.
+* **Developer hints** — pass `planning_hints={...}` to enforce disallowed tools,
+  preferred ordering, or parallelism ceilings.
+
+All three features are exercised in `examples/react_pause_resume/`, which runs
+entirely offline with stubbed LLM responses.
+
+### Adaptive re-planning & budgets (Phase C)
+
+Phase C closes the loop when things go sideways:
+
+* **Structured failure feedback** — if a tool raises after exhausting its retries,
+  the planner records `{failure: {node, args, error_code, suggestion}}` and feeds
+  it back to the LLM, prompting a constrained re-plan instead of aborting.
+* **Hard guardrails** — configure wall-clock deadlines and hop budgets directly
+  on `ReactPlanner`; attempts beyond the allotted hops surface deterministic
+  violations and ultimately finish with `reason="budget_exhausted"` alongside a
+  constraint snapshot.
+* **Typed exit reasons** — runs now finish with one of
+  `answer_complete`, `no_path`, or `budget_exhausted`, keeping downstream code
+  simple and machine-checkable.
+
+The new `examples/react_replan/` sample shows a retrieval timeout automatically
+recover via a cached index without leaving the JSON-only contract.
+
+### Parallel fan-out & joins (Phase D)
+
+Phase D lets the planner propose sets of independent tool calls and join them
+without leaving the typed surface area:
+
+* **Parallel `plan` blocks** — the LLM can return `{"plan": [...]}` actions
+  where each branch is validated against the catalog and executed concurrently.
+* **Typed joins** — provide a `{"join": {"node": ...}}` descriptor and the
+  planner will aggregate results, auto-populate fields like `expect`, `results`,
+  or `failures`, and feed branch metadata through `ctx.meta` for the join node.
+* **Deterministic telemetry** — branch errors, pauses, and joins are recorded as
+  structured observations so follow-up actions can re-plan or finish cleanly.
+
+See `examples/react_parallel/` for a shard fan-out that merges responses in one
+round-trip.
+
 
 ## 🧭 Repo Structure
 
@@ -626,6 +713,8 @@ pytest -q
 * `examples/streaming_llm/`: mock LLM emitting streaming chunks to an SSE sink.
 * `examples/metadata_propagation/`: attaching and consuming `Message.meta` context.
 * `examples/visualizer/`: exports Mermaid + DOT diagrams with loop/subflow annotations.
+* `examples/react_minimal/`: JSON-only ReactPlanner loop with a stubbed LLM.
+* `examples/react_pause_resume/`: Phase B planner features with pause/resume and developer hints.
 
 ---
 

{penguiflow-2.1.0 → penguiflow-2.2.0}/README.md

@@ -11,6 +11,9 @@
   <a href="https://github.com/penguiflow/penguiflow">
     <img src="https://img.shields.io/badge/coverage-85%25-brightgreen" alt="Coverage">
   </a>
+  <a href="https://nightly.link/penguiflow/penguiflow/workflows/benchmarks/main/benchmarks.json.zip">
+    <img src="https://img.shields.io/badge/benchmarks-latest-orange" alt="Benchmarks">
+  </a>
   <a href="https://pypi.org/project/penguiflow/">
     <img src="https://img.shields.io/pypi/v/penguiflow.svg" alt="PyPI version">
   </a>
@@ -52,6 +55,23 @@ It provides:
 Built on pure `asyncio` (no threads), PenguiFlow is small, predictable, and repo-agnostic.
 Product repos only define **their models + node functions** — the core stays dependency-light.
 
+## Gold Standard Scorecard
+
+| Area | Metric | Target | Current |
+| --- | --- | --- | --- |
+| Hop overhead | µs per hop | ≤ 500 | 398 |
+| Streaming order | gaps/dupes | 0 | 0 |
+| Cancel leakage | orphan tasks | 0 | 0 |
+| Coverage | lines | ≥85% | 87% |
+| Deps | count | ≤2 | 2 |
+| Import time | ms | ≤220 | 203 |
+
+## 📑 Core Behavior Spec
+
+* [Core Behavior Spec](docs/core_behavior_spec.md) — single-page rundown of ordering,
+  streaming, cancellation, deadline, and fan-in invariants with pointers to regression
+  tests.
+
 ---
 
 ## ✨ Why PenguiFlow?
@@ -301,6 +321,70 @@ The new `penguiflow.testkit` module keeps unit tests tiny:
 The harness is covered by `tests/test_testkit.py` and demonstrated in
 `examples/testkit_demo/`.
 
+### JSON-only ReAct planner (Phase A)
+
+Phase A introduces a lightweight planner loop that keeps PenguiFlow typed and
+deterministic:
+
+* `penguiflow.catalog.NodeSpec` + `build_catalog` turn registered nodes into
+  tool descriptors with JSON Schemas derived from your Pydantic models.
+* `penguiflow.planner.ReactPlanner` drives a JSON-only ReAct loop over those
+  descriptors, validating every LLM action with Pydantic and replaying invalid
+  steps to request corrections.
+* LiteLLM stays optional—install `penguiflow[planner]` or inject a custom
+  `llm_client` for deterministic/offline runs.
+
+See `examples/react_minimal/` for a stubbed end-to-end run.
+
+### Trajectory summarisation & pause/resume (Phase B)
+
+Phase B adds the tools you need for longer-running, approval-driven flows:
+
+* **Token-aware summaries** — `Trajectory.compress()` keeps a compact state and
+  the planner can route summaries through a cheaper `summarizer_llm` before
+  asking for the next action.
+* **`PlannerPause` contract** — nodes can call `await ctx.pause(...)` to return a
+  typed pause payload. Resume the run later with `ReactPlanner.resume(token, user_input=...)`.
+* **Developer hints** — pass `planning_hints={...}` to enforce disallowed tools,
+  preferred ordering, or parallelism ceilings.
+
+All three features are exercised in `examples/react_pause_resume/`, which runs
+entirely offline with stubbed LLM responses.
+
+### Adaptive re-planning & budgets (Phase C)
+
+Phase C closes the loop when things go sideways:
+
+* **Structured failure feedback** — if a tool raises after exhausting its retries,
+  the planner records `{failure: {node, args, error_code, suggestion}}` and feeds
+  it back to the LLM, prompting a constrained re-plan instead of aborting.
+* **Hard guardrails** — configure wall-clock deadlines and hop budgets directly
+  on `ReactPlanner`; attempts beyond the allotted hops surface deterministic
+  violations and ultimately finish with `reason="budget_exhausted"` alongside a
+  constraint snapshot.
+* **Typed exit reasons** — runs now finish with one of
+  `answer_complete`, `no_path`, or `budget_exhausted`, keeping downstream code
+  simple and machine-checkable.
+
+The new `examples/react_replan/` sample shows a retrieval timeout automatically
+recover via a cached index without leaving the JSON-only contract.
+
+### Parallel fan-out & joins (Phase D)
+
+Phase D lets the planner propose sets of independent tool calls and join them
+without leaving the typed surface area:
+
+* **Parallel `plan` blocks** — the LLM can return `{"plan": [...]}` actions
+  where each branch is validated against the catalog and executed concurrently.
+* **Typed joins** — provide a `{"join": {"node": ...}}` descriptor and the
+  planner will aggregate results, auto-populate fields like `expect`, `results`,
+  or `failures`, and feed branch metadata through `ctx.meta` for the join node.
+* **Deterministic telemetry** — branch errors, pauses, and joins are recorded as
+  structured observations so follow-up actions can re-plan or finish cleanly.
+
+See `examples/react_parallel/` for a shard fan-out that merges responses in one
+round-trip.
+
 
 ## 🧭 Repo Structure
 
@@ -581,6 +665,8 @@ pytest -q
 * `examples/streaming_llm/`: mock LLM emitting streaming chunks to an SSE sink.
 * `examples/metadata_propagation/`: attaching and consuming `Message.meta` context.
 * `examples/visualizer/`: exports Mermaid + DOT diagrams with loop/subflow annotations.
+* `examples/react_minimal/`: JSON-only ReactPlanner loop with a stubbed LLM.
+* `examples/react_pause_resume/`: Phase B planner features with pause/resume and developer hints.
 
 ---
 

{penguiflow-2.1.0 → penguiflow-2.2.0}/penguiflow/__init__.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 from . import testkit
 from .bus import BusEnvelope, MessageBus
+from .catalog import NodeSpec, SideEffect, build_catalog, tool
 from .core import (
     DEFAULT_QUEUE_MAXSIZE,
     Context,
@@ -12,11 +13,19 @@ from .core import (
     call_playbook,
     create,
 )
+from .debug import format_flow_event
 from .errors import FlowError, FlowErrorCode
 from .metrics import FlowEvent
-from .middlewares import Middleware
+from .middlewares import LatencyCallback, Middleware, log_flow_events
 from .node import Node, NodePolicy
 from .patterns import join_k, map_concurrent, predicate_router, union_router
+from .planner import (
+    PlannerAction,
+    PlannerFinish,
+    ReactPlanner,
+    Trajectory,
+    TrajectoryStep,
+)
 from .policies import DictRoutingPolicy, RoutingPolicy, RoutingRequest
 from .registry import ModelRegistry
 from .remote import (
@@ -45,8 +54,15 @@ __all__ = [
     "Node",
     "NodePolicy",
     "ModelRegistry",
+    "NodeSpec",
+    "SideEffect",
+    "build_catalog",
+    "tool",
     "Middleware",
+    "log_flow_events",
+    "LatencyCallback",
     "FlowEvent",
+    "format_flow_event",
     "FlowError",
     "FlowErrorCode",
     "MessageBus",
@@ -82,6 +98,11 @@ __all__ = [
     "RemoteCallResult",
     "RemoteStreamEvent",
     "RemoteNode",
+    "ReactPlanner",
+    "PlannerAction",
+    "PlannerFinish",
+    "Trajectory",
+    "TrajectoryStep",
 ]
 
-__version__ = "2.1.0"
+__version__ = "2.2.0"

penguiflow-2.2.0/penguiflow/catalog.py

@@ -0,0 +1,146 @@
+"""Tool catalog helpers for the planner."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, cast
+
+from pydantic import BaseModel
+
+from .node import Node
+from .registry import ModelRegistry
+
+SideEffect = Literal["pure", "read", "write", "external", "stateful"]
+
+
+@dataclass(frozen=True, slots=True)
+class NodeSpec:
+    """Structured metadata describing a planner-discoverable node."""
+
+    node: Node
+    name: str
+    desc: str
+    args_model: type[BaseModel]
+    out_model: type[BaseModel]
+    side_effects: SideEffect = "pure"
+    tags: Sequence[str] = field(default_factory=tuple)
+    auth_scopes: Sequence[str] = field(default_factory=tuple)
+    cost_hint: str | None = None
+    latency_hint_ms: int | None = None
+    safety_notes: str | None = None
+    extra: Mapping[str, Any] = field(default_factory=dict)
+
+    def to_tool_record(self) -> dict[str, Any]:
+        """Convert the spec to a serialisable record for prompting."""
+
+        return {
+            "name": self.name,
+            "desc": self.desc,
+            "side_effects": self.side_effects,
+            "tags": list(self.tags),
+            "auth_scopes": list(self.auth_scopes),
+            "cost_hint": self.cost_hint,
+            "latency_hint_ms": self.latency_hint_ms,
+            "safety_notes": self.safety_notes,
+            "args_schema": self.args_model.model_json_schema(),
+            "out_schema": self.out_model.model_json_schema(),
+            "extra": dict(self.extra),
+        }
+
+
+def _normalise_sequence(value: Sequence[str] | None) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    return tuple(dict.fromkeys(value))
+
+
+def tool(
+    *,
+    desc: str | None = None,
+    side_effects: SideEffect = "pure",
+    tags: Sequence[str] | None = None,
+    auth_scopes: Sequence[str] | None = None,
+    cost_hint: str | None = None,
+    latency_hint_ms: int | None = None,
+    safety_notes: str | None = None,
+    extra: Mapping[str, Any] | None = None,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Annotate a node function with catalog metadata."""
+
+    payload: dict[str, Any] = {
+        "desc": desc,
+        "side_effects": side_effects,
+        "tags": _normalise_sequence(tags),
+        "auth_scopes": _normalise_sequence(auth_scopes),
+        "cost_hint": cost_hint,
+        "latency_hint_ms": latency_hint_ms,
+        "safety_notes": safety_notes,
+        "extra": dict(extra) if extra else {},
+    }
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        func_ref = cast(Any, func)
+        func_ref.__penguiflow_tool__ = payload
+        return func
+
+    return decorator
+
+
+def _load_metadata(func: Callable[..., Any]) -> dict[str, Any]:
+    raw = getattr(func, "__penguiflow_tool__", None)
+    if not raw:
+        return {
+            "desc": inspect.getdoc(func) or func.__name__,
+            "side_effects": "pure",
+            "tags": (),
+            "auth_scopes": (),
+            "cost_hint": None,
+            "latency_hint_ms": None,
+            "safety_notes": None,
+            "extra": {},
+        }
+    return {
+        "desc": raw.get("desc") or inspect.getdoc(func) or func.__name__,
+        "side_effects": raw.get("side_effects", "pure"),
+        "tags": tuple(raw.get("tags", ())),
+        "auth_scopes": tuple(raw.get("auth_scopes", ())),
+        "cost_hint": raw.get("cost_hint"),
+        "latency_hint_ms": raw.get("latency_hint_ms"),
+        "safety_notes": raw.get("safety_notes"),
+        "extra": dict(raw.get("extra", {})),
+    }
+
+
+def build_catalog(
+    nodes: Sequence[Node],
+    registry: ModelRegistry,
+) -> list[NodeSpec]:
+    """Derive :class:`NodeSpec` objects from runtime nodes."""
+
+    specs: list[NodeSpec] = []
+    for node in nodes:
+        node_name = node.name or node.func.__name__
+        in_model, out_model = registry.models(node_name)
+        metadata = _load_metadata(node.func)
+        specs.append(
+            NodeSpec(
+                node=node,
+                name=node_name,
+                desc=metadata["desc"],
+                args_model=in_model,
+                out_model=out_model,
+                side_effects=metadata["side_effects"],
+                tags=metadata["tags"],
+                auth_scopes=metadata["auth_scopes"],
+                cost_hint=metadata["cost_hint"],
+                latency_hint_ms=metadata["latency_hint_ms"],
+                safety_notes=metadata["safety_notes"],
+                extra=metadata["extra"],
+            )
+        )
+    return specs
+
+
+__all__ = ["NodeSpec", "SideEffect", "build_catalog", "tool"]

{penguiflow-2.1.0 → penguiflow-2.2.0}/penguiflow/core.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 import asyncio
 import logging
 import time
+import warnings
 from collections import deque
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from contextlib import suppress
@@ -686,6 +687,21 @@ class PenguiFlow:
                     trace_id,
                 )
 
+                if (
+                    result is not None
+                    and self._expects_message_output(node)
+                    and not isinstance(result, Message)
+                ):
+                    node_name = node.name or node.node_id
+                    warning_msg = (
+                        "Node "
+                        f"'{node_name}' is registered for Message -> Message outputs "
+                        f"but returned {type(result).__name__}. "
+                        "Return a penguiflow.types.Message to preserve headers, "
+                        "trace_id, and meta."
+                    )
+                    warnings.warn(warning_msg, RuntimeWarning, stacklevel=2)
+
                 if result is not None:
                     (
                         destination,
@@ -1210,6 +1226,29 @@ class PenguiFlow:
         for waiter in waiters:
             waiter.set()
 
+    def _expects_message_output(self, node: Node) -> bool:
+        registry = self._registry
+        if registry is None:
+            return False
+
+        models = getattr(registry, "models", None)
+        if models is None:
+            return False
+
+        node_name = node.name
+        if not node_name:
+            return False
+
+        try:
+            _in_model, out_model = models(node_name)
+        except Exception:  # pragma: no cover - registry without entry
+            return False
+
+        try:
+            return issubclass(out_model, Message)
+        except TypeError:
+            return False
+
     def _controller_postprocess(
         self,
         node: Node,

penguiflow-2.2.0/penguiflow/debug.py

@@ -0,0 +1,30 @@
+"""Developer-facing debugging helpers for PenguiFlow."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from .metrics import FlowEvent
+
+
+def format_flow_event(event: FlowEvent) -> dict[str, Any]:
+    """Return a structured payload ready for logging.
+
+    The returned dictionary mirrors :meth:`FlowEvent.to_payload` and flattens any
+    embedded ``FlowError`` payload so that log aggregators can index the error
+    metadata (``flow_error_code``, ``flow_error_message``, ...).
+    """
+
+    payload = dict(event.to_payload())
+    error_payload: Mapping[str, Any] | None = event.error_payload
+    if error_payload is not None:
+        # Preserve the original payload for downstream consumers.
+        payload["flow_error"] = dict(error_payload)
+        for key, value in error_payload.items():
+            payload[f"flow_error_{key}"] = value
+    return payload
+
+
+__all__ = ["format_flow_event"]
+

{penguiflow-2.1.0 → penguiflow-2.2.0}/penguiflow/metrics.py

@@ -31,6 +31,15 @@ class FlowEvent:
     def __post_init__(self) -> None:
         object.__setattr__(self, "extra", MappingProxyType(dict(self.extra)))
 
+    @property
+    def error_payload(self) -> Mapping[str, Any] | None:
+        """Return the structured ``FlowError`` payload if present."""
+
+        raw_payload = self.extra.get("flow_error")
+        if isinstance(raw_payload, Mapping):
+            return MappingProxyType(dict(raw_payload))
+        return None
+
     @property
     def queue_depth(self) -> int:
         """Return the combined depth of incoming and outgoing queues."""

penguiflow-2.2.0/penguiflow/middlewares.py

@@ -0,0 +1,87 @@
+"""Middleware hooks for PenguiFlow."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Protocol
+
+from .metrics import FlowEvent
+
+LatencyCallback = Callable[[str, float, FlowEvent], None]
+
+
+class Middleware(Protocol):
+    """Base middleware signature receiving :class:`FlowEvent` objects."""
+
+    async def __call__(self, event: FlowEvent) -> None: ...
+
+
+def log_flow_events(
+    logger: logging.Logger | None = None,
+    *,
+    start_level: int = logging.INFO,
+    success_level: int = logging.INFO,
+    error_level: int = logging.ERROR,
+    latency_callback: LatencyCallback | None = None,
+) -> Middleware:
+    """Return middleware that emits structured node lifecycle logs.
+
+    Parameters
+    ----------
+    logger:
+        Optional :class:`logging.Logger` instance. When omitted a logger named
+        ``"penguiflow.flow"`` is used.
+    start_level, success_level, error_level:
+        Logging levels for ``node_start``, ``node_success``, and
+        ``node_error`` events respectively.
+    latency_callback:
+        Optional callable invoked with ``(event_type, latency_ms, event)`` for
+        ``node_success`` and ``node_error`` events. Use this hook to connect the
+        middleware to histogram-based metrics backends without
+        re-implementing timing logic.
+    """
+
+    log = logger or logging.getLogger("penguiflow.flow")
+
+    async def _middleware(event: FlowEvent) -> None:
+        if event.event_type not in {"node_start", "node_success", "node_error"}:
+            return
+
+        payload = event.to_payload()
+        log_level = start_level
+
+        if event.event_type == "node_start":
+            log_level = start_level
+        elif event.event_type == "node_success":
+            log_level = success_level
+        else:
+            log_level = error_level
+            if event.error_payload is not None:
+                payload = dict(payload)
+                payload["error_payload"] = dict(event.error_payload)
+
+        log.log(log_level, event.event_type, extra=payload)
+
+        if (
+            latency_callback is not None
+            and event.event_type in {"node_success", "node_error"}
+            and event.latency_ms is not None
+        ):
+            try:
+                latency_callback(event.event_type, float(event.latency_ms), event)
+            except Exception:
+                log.exception(
+                    "log_flow_events_latency_callback_error",
+                    extra={
+                        "event": "log_flow_events_latency_callback_error",
+                        "node_name": event.node_name,
+                        "node_id": event.node_id,
+                        "trace_id": event.trace_id,
+                    },
+                )
+
+    return _middleware
+
+
+__all__ = ["Middleware", "FlowEvent", "log_flow_events", "LatencyCallback"]

{penguiflow-2.1.0 → penguiflow-2.2.0}/penguiflow/registry.py

@@ -15,6 +15,8 @@ ModelT = TypeVar("ModelT", bound=BaseModel)
 class RegistryEntry:
     in_adapter: TypeAdapter[Any]
     out_adapter: TypeAdapter[Any]
+    in_model: type[BaseModel]
+    out_model: type[BaseModel]
 
 
 class ModelRegistry:
@@ -36,6 +38,8 @@ class ModelRegistry:
         self._entries[node_name] = RegistryEntry(
             TypeAdapter(in_model),
             TypeAdapter(out_model),
+            in_model,
+            out_model,
         )
 
     def adapters(self, node_name: str) -> tuple[TypeAdapter[Any], TypeAdapter[Any]]:
@@ -45,5 +49,22 @@ class ModelRegistry:
             raise KeyError(f"Node '{node_name}' not registered") from exc
         return entry.in_adapter, entry.out_adapter
 
+    def models(
+        self, node_name: str
+    ) -> tuple[type[BaseModel], type[BaseModel]]:
+        """Return the registered models for ``node_name``.
+
+        Raises
+        ------
+        KeyError
+            If the node has not been registered.
+        """
+
+        try:
+            entry = self._entries[node_name]
+        except KeyError as exc:
+            raise KeyError(f"Node '{node_name}' not registered") from exc
+        return entry.in_model, entry.out_model
+
 
 __all__ = ["ModelRegistry"]