penguiflow 2.0.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

penguiflow/remote.py

@@ -0,0 +1,486 @@
+"""Remote transport protocol and helper node for PenguiFlow."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol
+
+from pydantic import BaseModel
+
+from .core import TraceCancelled
+from .node import Node, NodePolicy
+from .state import RemoteBinding
+from .types import Message
+
+if TYPE_CHECKING:  # pragma: no cover - import for typing only
+    from .core import Context, PenguiFlow
+
+
+@dataclass(slots=True)
+class RemoteCallRequest:
+    """Input to :class:`RemoteTransport` implementations."""
+
+    message: Message
+    skill: str
+    agent_url: str
+    agent_card: Mapping[str, Any] | None = None
+    metadata: Mapping[str, Any] | None = None
+    timeout_s: float | None = None
+
+
+@dataclass(slots=True)
+class RemoteCallResult:
+    """Return value for :meth:`RemoteTransport.send`."""
+
+    result: Any
+    context_id: str | None = None
+    task_id: str | None = None
+    agent_url: str | None = None
+    meta: Mapping[str, Any] | None = None
+
+
+@dataclass(slots=True)
+class RemoteStreamEvent:
+    """Streaming event yielded by :meth:`RemoteTransport.stream`."""
+
+    text: str | None = None
+    done: bool = False
+    meta: Mapping[str, Any] | None = None
+    context_id: str | None = None
+    task_id: str | None = None
+    agent_url: str | None = None
+    result: Any | None = None
+
+
+class RemoteTransport(Protocol):
+    """Protocol describing the minimal remote invocation surface."""
+
+    async def send(self, request: RemoteCallRequest) -> RemoteCallResult:
+        """Perform a unary remote call."""
+
+    def stream(self, request: RemoteCallRequest) -> AsyncIterator[RemoteStreamEvent]:
+        """Perform a remote call that yields streaming events."""
+
+    async def cancel(self, *, agent_url: str, task_id: str) -> None:
+        """Cancel a remote task identified by ``task_id`` at ``agent_url``."""
+
+
+def _json_default(value: Any) -> Any:
+    if isinstance(value, BaseModel):
+        return value.model_dump(mode="json")
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    return repr(value)
+
+
+def _estimate_bytes(value: Any) -> int | None:
+    """Best-effort size estimation for observability metrics."""
+
+    if value is None:
+        return None
+    try:
+        if isinstance(value, BaseModel):
+            payload = value.model_dump(mode="json")
+        else:
+            payload = value
+        encoded = json.dumps(payload, default=_json_default).encode("utf-8")
+    except Exception:
+        try:
+            encoded = str(value).encode("utf-8")
+        except Exception:
+            return None
+    return len(encoded)
+
+
+def _text_bytes(text: str | None) -> int:
+    if text is None:
+        return 0
+    return len(text.encode("utf-8"))
+
+
+def _merge_remote_extra(
+    base: Mapping[str, Any],
+    *,
+    agent_url: str | None,
+    context_id: str | None,
+    task_id: str | None,
+    additional: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    extra = dict(base)
+    if agent_url is not None:
+        extra["remote_agent_url"] = agent_url
+    if context_id is not None:
+        extra["remote_context_id"] = context_id
+    if task_id is not None:
+        extra["remote_task_id"] = task_id
+    if additional:
+        for key, value in additional.items():
+            if value is not None:
+                extra[key] = value
+    return extra
+
+
+def RemoteNode(
+    *,
+    transport: RemoteTransport,
+    skill: str,
+    agent_url: str,
+    name: str,
+    agent_card: Mapping[str, Any] | None = None,
+    policy: NodePolicy | None = None,
+    streaming: bool = False,
+    record_binding: bool = True,
+) -> Node:
+    """Create a node that proxies work to a remote agent via ``transport``."""
+
+    node_policy = policy or NodePolicy()
+
+    async def _record_binding(
+        *,
+        runtime: PenguiFlow,
+        context: Context,
+        node_owner: Node,
+        trace_id: str,
+        context_id: str | None,
+        task_id: str | None,
+        agent_url_override: str | None,
+        base_extra: Mapping[str, Any],
+    ) -> tuple[asyncio.Task[None], asyncio.Event] | None:
+        if context_id is None or task_id is None:
+            return None
+
+        agent_ref = agent_url_override or agent_url
+
+        if record_binding:
+            binding = RemoteBinding(
+                trace_id=trace_id,
+                context_id=context_id,
+                task_id=task_id,
+                agent_url=agent_ref,
+            )
+            await runtime.save_remote_binding(binding)
+
+        cancel_event = runtime.ensure_trace_event(trace_id)
+
+        async def _issue_cancel(reason: str) -> None:
+            start_cancel = time.perf_counter()
+            try:
+                await transport.cancel(agent_url=agent_ref, task_id=task_id)
+            except Exception as exc:  # pragma: no cover - defensive logging
+                latency = (time.perf_counter() - start_cancel) * 1000
+                extra = _merge_remote_extra(
+                    base_extra,
+                    agent_url=agent_ref,
+                    context_id=context_id,
+                    task_id=task_id,
+                    additional={
+                        "remote_cancel_reason": reason,
+                        "remote_error": repr(exc),
+                        "remote_status": "cancel_error",
+                    },
+                )
+                await runtime.record_remote_event(
+                    event="remote_cancel_error",
+                    node=node_owner,
+                    context=context,
+                    trace_id=trace_id,
+                    latency_ms=latency,
+                    level=logging.ERROR,
+                    extra=extra,
+                )
+                return
+
+            latency = (time.perf_counter() - start_cancel) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=agent_ref,
+                context_id=context_id,
+                task_id=task_id,
+                additional={
+                    "remote_cancel_reason": reason,
+                    "remote_status": "cancelled",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_cancelled",
+                node=node_owner,
+                context=context,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.INFO,
+                extra=extra,
+            )
+
+        if cancel_event.is_set():
+            await _issue_cancel("pre_cancelled")
+            raise TraceCancelled(trace_id)
+
+        async def _mirror_cancel() -> None:
+            try:
+                await cancel_event.wait()
+            except asyncio.CancelledError:
+                return
+            await _issue_cancel("trace_cancel")
+
+        cancel_task = asyncio.create_task(_mirror_cancel())
+        runtime.register_external_task(trace_id, cancel_task)
+        return cancel_task, cancel_event
+
+    async def _remote_impl(message: Message, ctx: Context) -> Any:
+        if not isinstance(message, Message):
+            raise TypeError("Remote nodes require penguiflow.types.Message inputs")
+
+        runtime = ctx.runtime
+        if runtime is None:
+            raise RuntimeError("Context is not bound to a running PenguiFlow")
+
+        owner = ctx.owner
+        if not isinstance(owner, Node):  # pragma: no cover - defensive safety
+            raise RuntimeError("Remote context owner must be a Node")
+
+        trace_id = message.trace_id
+        cancel_task: asyncio.Task[None] | None = None
+        cancel_event: asyncio.Event | None = None
+        binding_registered = False
+
+        remote_context_id: str | None = None
+        remote_task_id: str | None = None
+        remote_agent_url_final = agent_url
+        response_bytes = 0
+        stream_events = 0
+
+        base_extra: dict[str, Any] = {
+            "remote_skill": skill,
+            "remote_transport": type(transport).__name__,
+            "remote_streaming": streaming,
+        }
+        request_bytes = _estimate_bytes(message)
+        if request_bytes is not None:
+            base_extra["remote_request_bytes"] = request_bytes
+
+        request = RemoteCallRequest(
+            message=message,
+            skill=skill,
+            agent_url=agent_url,
+            agent_card=agent_card,
+            metadata=message.meta,
+            timeout_s=node_policy.timeout_s,
+        )
+
+        async def _ensure_binding(
+            *,
+            context_id: str | None,
+            task_id: str | None,
+            agent_url_override: str | None,
+        ) -> None:
+            nonlocal cancel_task, cancel_event, binding_registered
+            nonlocal remote_context_id, remote_task_id, remote_agent_url_final
+            if context_id is not None:
+                remote_context_id = context_id
+            if task_id is not None:
+                remote_task_id = task_id
+            if agent_url_override is not None:
+                remote_agent_url_final = agent_url_override
+            if binding_registered:
+                return
+            if context_id is None or task_id is None:
+                return
+            record = await _record_binding(
+                runtime=runtime,
+                context=ctx,
+                node_owner=owner,
+                trace_id=trace_id,
+                context_id=context_id,
+                task_id=task_id,
+                agent_url_override=agent_url_override,
+                base_extra=base_extra,
+            )
+            if record is None:
+                return
+            cancel_task, cancel_event = record
+            binding_registered = True
+
+        async def _cleanup_cancel_task() -> None:
+            if cancel_task is not None:
+                try:
+                    if cancel_event is not None and cancel_event.is_set():
+                        await cancel_task
+                        return
+                    if not cancel_task.done():
+                        cancel_task.cancel()
+                    await cancel_task
+                except BaseException:  # pragma: no cover - cleanup guard
+                    pass
+
+        async def _run_stream() -> Any | None:
+            nonlocal response_bytes, stream_events, remote_agent_url_final
+            final_result: Any | None = None
+            stream_idx = 0
+            async for event in transport.stream(request):
+                stream_events = stream_idx + 1
+                await _ensure_binding(
+                    context_id=event.context_id,
+                    task_id=event.task_id,
+                    agent_url_override=event.agent_url,
+                )
+                if event.agent_url is not None:
+                    remote_agent_url_final = event.agent_url
+
+                chunk_bytes = 0
+                if event.text is not None:
+                    meta = dict(event.meta) if event.meta is not None else None
+                    chunk_bytes += _text_bytes(event.text)
+                    meta_bytes = _estimate_bytes(event.meta)
+                    if meta_bytes is not None:
+                        chunk_bytes += meta_bytes
+                    await ctx.emit_chunk(
+                        parent=message,
+                        text=event.text,
+                        done=event.done,
+                        meta=meta,
+                    )
+
+                if runtime is not None:
+                    meta_keys = None
+                    if event.meta:
+                        meta_keys = sorted(event.meta.keys())
+                    extra = _merge_remote_extra(
+                        base_extra,
+                        agent_url=remote_agent_url_final,
+                        context_id=remote_context_id,
+                        task_id=remote_task_id,
+                        additional={
+                            "remote_stream_seq": stream_idx,
+                            "remote_chunk_bytes": chunk_bytes if chunk_bytes else None,
+                            "remote_chunk_done": event.done,
+                            "remote_chunk_meta_keys": meta_keys,
+                        },
+                    )
+                    await runtime.record_remote_event(
+                        event="remote_stream_event",
+                        node=owner,
+                        context=ctx,
+                        trace_id=trace_id,
+                        latency_ms=(time.perf_counter() - call_start) * 1000,
+                        level=logging.DEBUG,
+                        extra=extra,
+                    )
+
+                if chunk_bytes:
+                    response_bytes += chunk_bytes
+
+                if event.result is not None:
+                    result_bytes = _estimate_bytes(event.result)
+                    if result_bytes is not None:
+                        response_bytes += result_bytes
+                    final_result = event.result
+
+                stream_idx += 1
+
+            return final_result
+
+        call_start = time.perf_counter()
+
+        await runtime.record_remote_event(
+            event="remote_call_start",
+            node=owner,
+            context=ctx,
+            trace_id=trace_id,
+            latency_ms=0.0,
+            level=logging.DEBUG,
+            extra=_merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=None,
+                task_id=None,
+            ),
+        )
+
+        try:
+            if streaming:
+                final_result = await _run_stream()
+                result_payload = final_result
+            else:
+                result = await transport.send(request)
+                await _ensure_binding(
+                    context_id=result.context_id,
+                    task_id=result.task_id,
+                    agent_url_override=result.agent_url,
+                )
+                if result.context_id is not None:
+                    remote_context_id = result.context_id
+                if result.task_id is not None:
+                    remote_task_id = result.task_id
+                if result.agent_url is not None:
+                    remote_agent_url_final = result.agent_url
+                result_payload = result.result
+                response_size = _estimate_bytes(result_payload)
+                if response_size is not None:
+                    response_bytes += response_size
+        except TraceCancelled:
+            raise
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            latency = (time.perf_counter() - call_start) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=remote_context_id,
+                task_id=remote_task_id,
+                additional={
+                    "remote_error": repr(exc),
+                    "remote_status": "error",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_error",
+                node=owner,
+                context=ctx,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.ERROR,
+                extra=extra,
+            )
+            raise
+        else:
+            latency = (time.perf_counter() - call_start) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=remote_context_id,
+                task_id=remote_task_id,
+                additional={
+                    "remote_response_bytes": response_bytes,
+                    "remote_stream_events": stream_events,
+                    "remote_status": "success",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_success",
+                node=owner,
+                context=ctx,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.INFO,
+                extra=extra,
+            )
+            return result_payload
+        finally:
+            await _cleanup_cancel_task()
+
+    return Node(_remote_impl, name=name, policy=node_policy)
+
+
+__all__ = [
+    "RemoteCallRequest",
+    "RemoteCallResult",
+    "RemoteStreamEvent",
+    "RemoteTransport",
+    "RemoteNode",
+]

penguiflow/state.py

@@ -0,0 +1,64 @@
+"""State store protocol and helpers for PenguiFlow."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from .metrics import FlowEvent
+
+
+@dataclass(slots=True)
+class StoredEvent:
+    """Representation of a runtime event persisted by a state store."""
+
+    trace_id: str | None
+    ts: float
+    kind: str
+    node_name: str | None
+    node_id: str | None
+    payload: Mapping[str, Any]
+
+    @classmethod
+    def from_flow_event(cls, event: FlowEvent) -> StoredEvent:
+        """Create a stored representation from a :class:`FlowEvent`."""
+
+        return cls(
+            trace_id=event.trace_id,
+            ts=event.ts,
+            kind=event.event_type,
+            node_name=event.node_name,
+            node_id=event.node_id,
+            payload=event.to_payload(),
+        )
+
+
+@dataclass(slots=True)
+class RemoteBinding:
+    """Association between a trace and a remote worker/agent."""
+
+    trace_id: str
+    context_id: str
+    task_id: str
+    agent_url: str
+
+
+class StateStore(Protocol):
+    """Protocol for durable state adapters used by PenguiFlow."""
+
+    async def save_event(self, event: StoredEvent) -> None:
+        """Persist a runtime event.
+
+        Implementations may choose any storage backend (Postgres, Redis, etc.).
+        The method must be idempotent since retries can emit duplicate events.
+        """
+
+    async def load_history(self, trace_id: str) -> Sequence[StoredEvent]:
+        """Return the ordered history for a trace id."""
+
+    async def save_remote_binding(self, binding: RemoteBinding) -> None:
+        """Persist the mapping between a trace and an external worker."""
+
+
+__all__ = ["StateStore", "StoredEvent", "RemoteBinding"]

{penguiflow-2.0.0.dist-info → penguiflow-2.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: penguiflow
-Version: 2.0.0
+Version: 2.1.0
 Summary: Async agent orchestration primitives.
 Author: PenguiFlow Team
 License: MIT License
@@ -37,6 +37,10 @@ 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: ruff>=0.2; extra == "dev"
+Requires-Dist: fastapi>=0.110; extra == "dev"
+Requires-Dist: httpx>=0.27; extra == "dev"
+Provides-Extra: a2a-server
+Requires-Dist: fastapi>=0.110; extra == "a2a-server"
 Dynamic: license-file
 
 # PenguiFlow 🐧❄️
@@ -77,6 +81,18 @@ It provides:
 * **Observability hooks** (`FlowEvent` callbacks for logging, MLflow, or custom metrics sinks)
 * **Policy-driven routing** (optional policies steer routers without breaking existing flows)
 * **Traceable exceptions** (`FlowError` captures node/trace metadata and optionally emits to Rookery)
+* **Distribution hooks (opt-in)** — plug a `StateStore` to persist trace history and a
+  `MessageBus` to publish floe traffic for remote workers without changing existing flows.
+* **Remote calls (opt-in)** — `RemoteNode` bridges the runtime to external agents through a
+  pluggable `RemoteTransport` interface (A2A-ready) while propagating streaming chunks and
+  cancellation.
+* **A2A server adapter (opt-in)** — wrap a PenguiFlow graph in a FastAPI surface using
+  `penguiflow_a2a.A2AServerAdapter` so other agents can call `message/send`,
+  `message/stream`, and `tasks/cancel` while reusing the runtime's backpressure and
+  cancellation semantics.
+* **Observability & ops polish** — remote calls emit structured metrics (latency, payload
+  sizes, cancel reasons) and the `penguiflow-admin` CLI replays trace history from any
+  configured `StateStore` for debugging.
 
 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.
@@ -168,6 +184,10 @@ print(out.payload)            # PackOut(...)
 await flow.stop()
 ```
 
+> **Opt-in distribution:** pass `state_store=` and/or `message_bus=` when calling
+> `penguiflow.core.create(...)` to persist trace history and publish floe traffic
+> without changing node logic.
+
 ---
 
 ## 🧭 Design Principles
@@ -222,6 +242,60 @@ sacrificing backpressure or ordering guarantees.  The helper wraps the payload i
 increments per-stream sequence numbers.  See `tests/test_streaming.py` and
 `examples/streaming_llm/` for an end-to-end walk-through.
 
+### Remote orchestration
+
+Phase 2 introduces `RemoteNode` and the `RemoteTransport` protocol so flows can delegate
+work to remote agents (e.g., the A2A JSON-RPC/SSE ecosystem) without changing existing
+nodes.  The helper records remote bindings via the `StateStore`, mirrors streaming
+partials back into the graph, and propagates per-trace cancellation to remote tasks via
+`RemoteTransport.cancel`.  See `tests/test_remote.py` for reference in-memory transports.
+
+### Exposing a flow over A2A
+
+Install the optional extra to expose PenguiFlow as an A2A-compatible FastAPI service:
+
+```bash
+pip install "penguiflow[a2a-server]"
+```
+
+Create the adapter and mount the routes:
+
+```python
+from penguiflow import Message, Node, create
+from penguiflow_a2a import A2AAgentCard, A2AServerAdapter, A2ASkill, create_a2a_app
+
+async def orchestrate(message: Message, ctx):
+    await ctx.emit_chunk(parent=message, text="thinking...")
+    return {"result": "done"}
+
+node = Node(orchestrate, name="main")
+flow = create(node.to())
+
+card = A2AAgentCard(
+    name="Main Agent",
+    description="Primary entrypoint for orchestration",
+    version="2.1.0",
+    skills=[A2ASkill(name="orchestrate", description="Handles orchestration")],
+)
+
+adapter = A2AServerAdapter(
+    flow,
+    agent_card=card,
+    agent_url="https://agent.example",
+)
+app = create_a2a_app(adapter)
+```
+
+The generated FastAPI app implements:
+
+* `GET /agent` for discovery (Agent Card)
+* `POST /message/send` for unary execution
+* `POST /message/stream` for SSE streaming
+* `POST /tasks/cancel` to mirror cancellation into PenguiFlow traces
+
+`A2AServerAdapter` reuses the runtime's `StateStore` hooks, so bindings between trace IDs
+and external `taskId`/`contextId` pairs are persisted automatically.
+
 ### Reliability & guardrails
 
 PenguiFlow enforces reliability boundaries out of the box:
@@ -478,9 +552,15 @@ docs or diagramming pipelines.
 * **Structured `FlowEvent`s**: every node event carries `{ts, trace_id, node_name, event,
   latency_ms, q_depth_in, q_depth_out, attempt}` plus a mutable `extra` map for custom
   annotations.
+* **Remote call telemetry**: `RemoteNode` executions emit extra metrics (latency, request
+  and response bytes, context/task identifiers, cancel reasons) so remote hops can be
+  traced end-to-end.
 * **Middleware hooks**: subscribe observers (e.g., MLflow) to the structured `FlowEvent`
   stream. See `examples/mlflow_metrics/` for an MLflow integration and
   `examples/reliability_middleware/` for a concrete timeout + retry walkthrough.
+* **`penguiflow-admin` CLI**: inspect or replay stored trace history from any configured
+  `StateStore` (`penguiflow-admin history <trace>` or `penguiflow-admin replay <trace>`)
+  when debugging distributed runs.
 
 ---
 
@@ -488,9 +568,9 @@ docs or diagramming pipelines.
 
 - **In-process runtime**: there is no built-in distribution layer yet. Long-running CPU work should be delegated to your own pools or services.
 - **Registry-driven typing**: nodes default to validation. Provide a `ModelRegistry` when calling `flow.run(...)` or set `validate="none"` explicitly for untyped hops.
-- **Observability**: structured `FlowEvent` callbacks power logs/metrics; integrations with
-  third-party stacks (OTel, Prometheus, Datadog) remain DIY. See the MLflow middleware
-  example for a lightweight pattern.
+- **Observability**: structured `FlowEvent` callbacks and the `penguiflow-admin` CLI power
+  local debugging; integrations with third-party stacks (OTel, Prometheus, Datadog) remain
+  DIY. See the MLflow middleware example for a lightweight pattern.
 - **Roadmap**: follow-up releases focus on optional distributed backends, deeper observability integrations, and additional playbook patterns. Contributions and proposals are welcome!
 
 ---