penguiflow 1.0.0__py3-none-any.whl

penguiflow/node.py

@@ -0,0 +1,117 @@
+"""Node abstractions for PenguiFlow runtime."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import uuid
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from pydantic.type_adapter import TypeAdapter
+
+    from .core import Context
+    from .registry import ModelRegistry
+
+
+@dataclass(slots=True)
+class NodePolicy:
+    """Execution policy configuration placeholder."""
+
+    validate: str = "both"
+    timeout_s: float | None = None
+    max_retries: int = 0
+    backoff_base: float = 0.5
+    backoff_mult: float = 2.0
+    max_backoff: float | None = None
+
+    def __post_init__(self) -> None:
+        if self.validate not in {"both", "in", "out", "none"}:
+            raise ValueError("validate must be one of 'both', 'in', 'out', 'none'")
+
+
+@dataclass(slots=True)
+class Node:
+    """Wraps an async callable with metadata used by the runtime."""
+
+    func: Callable[..., Awaitable[Any]]
+    name: str | None = None
+    policy: NodePolicy = field(default_factory=NodePolicy)
+    allow_cycle: bool = False
+    node_id: str = field(init=False)
+
+    def __post_init__(self) -> None:
+        if not asyncio.iscoroutinefunction(self.func):
+            raise TypeError("Node function must be declared with async def")
+
+        self.name = self.name or self.func.__name__
+        assert self.name is not None  # narrow for type-checkers
+        self.node_id = uuid.uuid4().hex
+
+        signature = inspect.signature(self.func)
+        params = list(signature.parameters.values())
+        if len(params) != 2:
+            raise ValueError(
+                f"Node '{self.name}' must accept exactly two parameters "
+                f"(message, ctx); got {len(params)}"
+            )
+
+        ctx_param = params[1]
+        if ctx_param.kind not in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        ):
+            raise ValueError("Context parameter must be positional")
+
+    def _maybe_validate(
+        self,
+        adapter: TypeAdapter[Any] | None,
+        value: Any,
+        *,
+        enforce: bool,
+    ) -> Any:
+        if not enforce or adapter is None:
+            return value
+        return adapter.validate_python(value)
+
+    async def invoke(
+        self,
+        message: Any,
+        ctx: Context,
+        *,
+        registry: ModelRegistry | None,
+    ) -> Any:
+        """Invoke the underlying coroutine, applying optional validation."""
+
+        adapter_in: TypeAdapter[Any] | None = None
+        adapter_out: TypeAdapter[Any] | None = None
+
+        if registry is not None and self.policy.validate != "none":
+            node_name = self.name
+            assert node_name is not None
+            adapter_in, adapter_out = registry.adapters(node_name)
+
+        enforce_in = self.policy.validate in {"in", "both"}
+        enforce_out = self.policy.validate in {"out", "both"}
+
+        validated_msg = self._maybe_validate(adapter_in, message, enforce=enforce_in)
+        result = await self.func(validated_msg, ctx)
+
+        if result is None:
+            return None
+
+        return self._maybe_validate(adapter_out, result, enforce=enforce_out)
+
+    def to(self, *nodes: Node) -> tuple[Node, tuple[Node, ...]]:
+        return self, nodes
+
+    def __hash__(self) -> int:
+        return hash(self.node_id)
+
+    def __repr__(self) -> str:  # pragma: no cover - debug helper
+        return f"Node(name={self.name!r}, node_id={self.node_id})"
+
+
+__all__ = ["Node", "NodePolicy"]

penguiflow/patterns.py

@@ -0,0 +1,142 @@
+"""Common orchestration patterns for PenguiFlow."""
+
+from __future__ import annotations
+
+import asyncio
+from collections import defaultdict
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from typing import Any, TypeVar, cast
+
+from pydantic import BaseModel
+from pydantic.type_adapter import TypeAdapter
+
+from .node import Node, NodePolicy
+from .types import Message
+
+PayloadT = TypeVar("PayloadT")
+ResultT = TypeVar("ResultT")
+
+__all__ = [
+    "map_concurrent",
+    "join_k",
+    "predicate_router",
+    "union_router",
+]
+
+
+async def map_concurrent(
+    items: Iterable[PayloadT],
+    worker: Callable[[PayloadT], Awaitable[ResultT]],
+    *,
+    max_concurrency: int = 8,
+) -> list[ResultT]:
+    """Run the async *worker* across *items* with bounded concurrency."""
+
+    items_list = list(items)
+    semaphore = asyncio.Semaphore(max(1, max_concurrency))
+    results: list[ResultT | None] = [None] * len(items_list)
+
+    async def run(index: int, item: PayloadT) -> None:
+        async with semaphore:
+            results[index] = await worker(item)
+
+    await asyncio.gather(*(run(idx, item) for idx, item in enumerate(items_list)))
+    return [cast(ResultT, result) for result in results]
+
+
+def predicate_router(
+    name: str,
+    predicate: Callable[[Any], Sequence[Node | str] | Node | str | None],
+) -> Node:
+    """Create a node that routes messages based on predicate outputs."""
+
+    async def router(msg: Any, ctx) -> None:
+        targets = predicate(msg)
+        if targets is None:
+            return
+
+        normalized = _normalize_targets(ctx, targets)
+        if normalized:
+            await ctx.emit(msg, to=normalized)
+
+    return Node(router, name=name, policy=NodePolicy(validate="none"))
+
+
+def union_router(
+    name: str,
+    union_model: type[BaseModel],
+) -> Node:
+    """Route based on a discriminated union Pydantic model."""
+
+    adapter = TypeAdapter(union_model)
+
+    async def router(msg: BaseModel, ctx) -> None:
+        validated = adapter.validate_python(msg)
+
+        target = getattr(validated, "kind", validated.__class__.__name__)
+        normalized = _normalize_targets(ctx, target)
+        if not normalized:
+            raise KeyError(f"No successor matches '{target}'")
+        await ctx.emit(validated, to=normalized)
+
+    return Node(router, name=name, policy=NodePolicy(validate="none"))
+
+
+def join_k(name: str, k: int) -> Node:
+    """Aggregate *k* messages per trace_id and emit the grouped payloads."""
+
+    if k <= 0:
+        raise ValueError("k must be positive")
+
+    buckets: defaultdict[str, list[Any]] = defaultdict(list)
+
+    async def aggregator(msg: Any, ctx) -> Any:
+        trace_id = getattr(msg, "trace_id", None)
+        if trace_id is None:
+            raise ValueError("join_k requires messages with trace_id")
+
+        bucket = buckets[trace_id]
+        bucket.append(msg)
+        if len(bucket) < k:
+            return None
+
+        buckets.pop(trace_id, None)
+        batch = list(bucket)
+        first = batch[0]
+        if isinstance(first, Message):
+            payloads = [item.payload for item in batch]
+            aggregated = first.model_copy(update={"payload": payloads})
+            return aggregated
+        return batch
+
+    return Node(aggregator, name=name, policy=NodePolicy(validate="none"))
+
+
+def _normalize_targets(context, targets) -> list[Node]:
+    if isinstance(targets, Node):
+        target_list: Sequence[Node | str] = [targets]
+    elif isinstance(targets, str):
+        target_list = [targets]
+    else:
+        target_list = list(targets)
+
+    normalized: list[Node] = []
+    candidates = list(getattr(context, "_outgoing", {}).keys())
+    for target in target_list:
+        if isinstance(target, Node):
+            normalized.append(target)
+            continue
+
+        if not isinstance(target, str):
+            raise TypeError("Targets must be Node or str")
+
+        matched = None
+        for node in candidates:
+            if isinstance(node, Node) and node.name == target:
+                matched = node
+                break
+        if matched is None:
+            raise KeyError(f"No successor named '{target}'")
+        normalized.append(matched)
+
+    return normalized

penguiflow/registry.py

@@ -0,0 +1,49 @@
+"""Model registry for PenguiFlow."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+from pydantic import BaseModel
+from pydantic.type_adapter import TypeAdapter
+
+ModelT = TypeVar("ModelT", bound=BaseModel)
+
+
+@dataclass(slots=True)
+class RegistryEntry:
+    in_adapter: TypeAdapter[Any]
+    out_adapter: TypeAdapter[Any]
+
+
+class ModelRegistry:
+    """Stores per-node type adapters for validation."""
+
+    def __init__(self) -> None:
+        self._entries: dict[str, RegistryEntry] = {}
+
+    def register(
+        self,
+        node_name: str,
+        in_model: type[BaseModel],
+        out_model: type[BaseModel],
+    ) -> None:
+        if node_name in self._entries:
+            raise ValueError(f"Node '{node_name}' already registered")
+        if not issubclass(in_model, BaseModel) or not issubclass(out_model, BaseModel):
+            raise TypeError("Models must inherit from pydantic.BaseModel")
+        self._entries[node_name] = RegistryEntry(
+            TypeAdapter(in_model),
+            TypeAdapter(out_model),
+        )
+
+    def adapters(self, node_name: str) -> tuple[TypeAdapter[Any], TypeAdapter[Any]]:
+        try:
+            entry = self._entries[node_name]
+        except KeyError as exc:
+            raise KeyError(f"Node '{node_name}' not registered") from exc
+        return entry.in_adapter, entry.out_adapter
+
+
+__all__ = ["ModelRegistry"]

penguiflow/types.py

@@ -0,0 +1,58 @@
+"""Typed message and controller models for PenguiFlow."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class Headers(BaseModel):
+    tenant: str
+    topic: str | None = None
+    priority: int = 0
+
+
+class Message(BaseModel):
+    payload: Any
+    headers: Headers
+    trace_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    ts: float = Field(default_factory=time.time)
+    deadline_s: float | None = None
+
+
+class PlanStep(BaseModel):
+    kind: Literal["retrieve", "web", "sql", "summarize", "route", "stop"]
+    args: dict[str, Any] = Field(default_factory=dict)
+    max_concurrency: int = 1
+
+
+class Thought(BaseModel):
+    steps: list[PlanStep]
+    rationale: str
+    done: bool = False
+
+
+class WM(BaseModel):
+    query: str
+    facts: list[Any] = Field(default_factory=list)
+    hops: int = 0
+    budget_hops: int = 8
+    confidence: float = 0.0
+
+
+class FinalAnswer(BaseModel):
+    text: str
+    citations: list[str] = Field(default_factory=list)
+
+
+__all__ = [
+    "Headers",
+    "Message",
+    "PlanStep",
+    "Thought",
+    "WM",
+    "FinalAnswer",
+]

penguiflow/viz.py

@@ -0,0 +1,5 @@
+"""Visualization helpers for PenguiFlow graphs."""
+
+from __future__ import annotations
+
+__all__: list[str] = []