nodes.py had grown to 2k lines holding every integration behind a single blanket mypy exemption. It is now a package split by the outside world each node talks to, so the exemption shrinks to the four integration modules; base and mlp are type-checked, which turned up a dozen missing annotations. The senders opened a fresh connection — and, in the MQTT case, a fresh thread pool and event loop — for every single message. HTTP senders now share one pooled client, and a publisher holds one broker connection for its lifetime, fed from a bounded queue that drops the oldest value when the broker cannot keep up. An HTTP sender also no longer trips over a JSON reply that is not an object: outputs are keyed by port, so a bare scalar is a valid reply with nothing to publish. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011LF61rxW1FG5YCD2J9YqjY
306 lines
12 KiB
Python
306 lines
12 KiB
Python
"""The node base class: identity, ports, lifecycle and execution.
|
|
|
|
Every node type in this package derives from :class:`Node`; the type-specific
|
|
modules beside this one add what talking to a particular outside world means.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import Callable, Coroutine, Iterable
|
|
from typing import TYPE_CHECKING, Any, TypeAlias
|
|
|
|
from app.flow import logs
|
|
from app.flow.messages import MessageSpec, qualify
|
|
from app.flow.supervision import Supervisor
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
|
|
from app.flow.pipeline import Pipeline
|
|
from app.flow.state import StateBackend
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# What a node hands back: the pipeline's state once it is bound, since a
|
|
# trigger runs the graph, and its own outputs when it is not.
|
|
NodeResult: TypeAlias = "StateBackend | dict[str, Any] | None"
|
|
|
|
|
|
class Node:
|
|
"""
|
|
A pipeline node that wraps a function with typed inputs/outputs.
|
|
|
|
Nodes are the fundamental building blocks of a pipeline. Each node encapsulates
|
|
a function that processes data, with explicit message-based inputs (requires)
|
|
and outputs (provides). Nodes can be connected into a directed acyclic graph (DAG)
|
|
where data flows from upstream to downstream nodes.
|
|
|
|
:param f: The function to execute when the node runs.
|
|
:type f: Callable
|
|
:param requires: Input messages this node consumes. Can be a single Message
|
|
or list of Messages. Empty for source nodes.
|
|
:type requires: MessageSpec | list[MessageSpec]
|
|
:param provides: Output messages this node produces. Can be a single Message
|
|
or list of Messages.
|
|
:type provides: MessageSpec | list[MessageSpec]
|
|
:param params: Additional parameters passed to the function during execution.
|
|
:type params: dict
|
|
:param name: Optional name for the node. Defaults to function name.
|
|
:type name: str | None
|
|
|
|
:ivar synchronous: If True, node only executes when all required inputs have
|
|
new versions since last execution. Useful for synchronizing multiple streams.
|
|
:vartype synchronous: bool
|
|
|
|
:example:
|
|
>>> def process_temp(temperature, params):
|
|
... return {"celsius": temperature * 0.5 + 32}
|
|
>>>
|
|
>>> temp_node = Node(
|
|
... f=process_temp,
|
|
... requires=MessageSpec(name="temperature", dtype=DType.FLOAT),
|
|
... provides=MessageSpec(name="celsius", dtype=DType.FLOAT),
|
|
... params={},
|
|
... )
|
|
"""
|
|
|
|
__slots__ = (
|
|
"f",
|
|
"id",
|
|
"flow",
|
|
"name",
|
|
"input_ports",
|
|
"output_ports",
|
|
"requires",
|
|
"provides",
|
|
"params",
|
|
"_pipeline",
|
|
"synchronous",
|
|
"_on_health",
|
|
"supervisor",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
f: Callable[..., Any],
|
|
requires: MessageSpec | Iterable[MessageSpec] = (),
|
|
provides: MessageSpec | Iterable[MessageSpec] = (),
|
|
params: dict[str, Any] | None = None,
|
|
name: str | None = None,
|
|
):
|
|
self.f = f
|
|
self._pipeline: Pipeline | None = None
|
|
self._on_health: Callable[[Node, str, str | None], None] | None = None
|
|
# Set by the controller before start(); a node built on its own (tests,
|
|
# previews) runs its loops unsupervised.
|
|
self.supervisor: Supervisor | None = None
|
|
self.params = dict(params) if params else {}
|
|
self.synchronous = bool(self.params.get("synchronous", False))
|
|
|
|
self.input_ports = self._normalize_ports(requires)
|
|
self.output_ports = self._normalize_ports(provides)
|
|
|
|
self.flow = ""
|
|
self.name: str = name or str(getattr(f, "__name__", "node"))
|
|
self.id: str = self.name
|
|
self._index_ports()
|
|
|
|
@staticmethod
|
|
def _normalize_ports(
|
|
msgs: MessageSpec | Iterable[MessageSpec],
|
|
) -> list[MessageSpec]:
|
|
"""Accept a single spec or any iterable of them."""
|
|
if isinstance(msgs, MessageSpec):
|
|
return [msgs]
|
|
return list(msgs or ())
|
|
|
|
def _index_ports(self) -> None:
|
|
"""Index bound ports by message name; unbound ports have no wiring."""
|
|
self.requires: dict[str, MessageSpec] = {
|
|
s.name: s for s in self.input_ports if s.name
|
|
}
|
|
self.provides: dict[str, MessageSpec] = {
|
|
s.name: s for s in self.output_ports if s.name
|
|
}
|
|
|
|
def assign_flow(self, flow: str, node_id: str) -> None:
|
|
"""Place this node in a flow, qualifying its identity and messages.
|
|
|
|
Called once by the loader, before the pipeline is built.
|
|
"""
|
|
self.flow = flow
|
|
self.name = node_id
|
|
self.id = f"{flow}.{node_id}"
|
|
self.input_ports = [
|
|
s.model_copy(update={"name": qualify(flow, s.name)})
|
|
for s in self.input_ports
|
|
]
|
|
self.output_ports = [
|
|
s.model_copy(update={"name": qualify(flow, s.name)})
|
|
for s in self.output_ports
|
|
]
|
|
self._index_ports()
|
|
|
|
@property
|
|
def local_id(self) -> str:
|
|
"""The node's id within its flow."""
|
|
return self.name
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Lifecycle
|
|
#
|
|
# A plain logic node has nothing to start or stop. The ones that talk to the
|
|
# outside — subscriptions, schedules, webhooks — override these, which is
|
|
# how the controller can bring a flow up or down without knowing what any
|
|
# particular node type is.
|
|
# -------------------------------------------------------------------------
|
|
|
|
async def start(self, app: FastAPI | None = None) -> None:
|
|
"""Begin whatever this node listens to. Called when its flow starts."""
|
|
|
|
async def stop(self, app: FastAPI | None = None) -> None:
|
|
"""Undo :meth:`start`. Called before a rebuild, and must be idempotent."""
|
|
|
|
def _run_supervised(
|
|
self, name: str, factory: Callable[[], Coroutine[Any, Any, None]]
|
|
) -> asyncio.Task[None] | None:
|
|
"""Run a background loop under supervision where there is one.
|
|
|
|
Returns the task only when unsupervised, since that is the one case
|
|
where the caller has to cancel it itself.
|
|
"""
|
|
if self.supervisor is not None:
|
|
self.supervisor.spawn(f"{self.id}:{name}", self.flow, factory)
|
|
return None
|
|
return asyncio.create_task(factory())
|
|
|
|
def report_health(self, status: str, detail: str | None = None) -> None:
|
|
"""Say how this node's connection is doing: ok, degraded or down."""
|
|
if self._on_health is not None:
|
|
self._on_health(self, status, detail)
|
|
|
|
def bind(self, pipeline: Pipeline) -> None:
|
|
"""
|
|
Bind this node to a pipeline for external triggering.
|
|
|
|
Once bound, the node can trigger downstream execution when called.
|
|
This is typically done automatically during pipeline construction.
|
|
|
|
:param pipeline: The pipeline to bind this node to.
|
|
:type pipeline: Pipeline
|
|
"""
|
|
self._pipeline = pipeline
|
|
|
|
def __repr__(self) -> str:
|
|
return self.id
|
|
|
|
def __hash__(self) -> int:
|
|
return hash(self.id)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Port translation
|
|
#
|
|
# The graph speaks qualified message names; node functions speak ports.
|
|
# -------------------------------------------------------------------------
|
|
|
|
def _to_kwargs(self, inputs: dict[str, Any]) -> dict[str, Any]:
|
|
kwargs = {}
|
|
for msg_name, value in inputs.items():
|
|
spec = self.requires.get(msg_name)
|
|
if spec is None:
|
|
continue
|
|
spec.check(value)
|
|
kwargs[spec.port] = value
|
|
return kwargs
|
|
|
|
def _to_messages(self, retval: dict[str, Any] | None) -> dict[str, Any] | None:
|
|
"""Map a function's port-keyed return value onto message names."""
|
|
if not retval:
|
|
return None
|
|
by_port = {s.port: s for s in self.output_ports if s.name}
|
|
outputs = {}
|
|
for key, value in retval.items():
|
|
spec = by_port.get(key) or self.provides.get(key)
|
|
if spec is None:
|
|
continue
|
|
spec.check(value)
|
|
outputs[spec.name] = value
|
|
return outputs or None
|
|
|
|
def execute(self, inputs: dict[str, Any] | None = None) -> dict[str, Any] | None:
|
|
"""Run the node function and return its outputs by message name.
|
|
|
|
Downstream nodes are not triggered — the pipeline schedules those.
|
|
"""
|
|
kwargs = self._to_kwargs(inputs or {})
|
|
return self._to_messages(self.f(**kwargs, params=self.params))
|
|
|
|
def trigger(self, inputs: dict[str, Any] | None = None) -> NodeResult:
|
|
"""
|
|
Trigger this node externally, executing downstream nodes if dependencies are met.
|
|
|
|
This method is for nodes that receive data from upstream dependencies.
|
|
For trigger/subscriber nodes that inject data into the pipeline, use :meth:`inject`.
|
|
|
|
:param inputs: Input values matching this node's ``requires``.
|
|
:type inputs: dict | None
|
|
:returns: Result of the node execution and downstream propagation.
|
|
:rtype: dict | None
|
|
:raises RuntimeError: If node is not bound to a pipeline.
|
|
"""
|
|
if self._pipeline is None:
|
|
raise RuntimeError("Node must be bound to a pipeline to trigger")
|
|
return self(inputs)
|
|
|
|
def inject(self, outputs: dict[str, Any] | None = None) -> NodeResult:
|
|
"""
|
|
Inject data into the pipeline as if this node produced it.
|
|
|
|
This method is for trigger/subscriber nodes that receive external data
|
|
(e.g., HTTP requests, MQTT messages) and need to inject it into the pipeline.
|
|
The data is validated against this node's ``provides`` specification.
|
|
|
|
For source nodes (nodes with no ``requires``), if no outputs are provided,
|
|
the node's function will be executed to generate outputs.
|
|
|
|
:param outputs: Output values matching this node's ``provides``.
|
|
:type outputs: dict | None
|
|
:returns: Result of downstream propagation.
|
|
:rtype: dict | None
|
|
:raises RuntimeError: If node is not bound to a pipeline.
|
|
:raises TypeError: If output values don't match ``provides`` types.
|
|
:raises KeyError: If required output keys are missing.
|
|
"""
|
|
if self._pipeline is None:
|
|
raise RuntimeError("Node must be bound to a pipeline to inject")
|
|
|
|
outputs = outputs or {}
|
|
|
|
# A source node asked to inject nothing produces its own data.
|
|
if not outputs and not self.requires:
|
|
collected = logs.Collector()
|
|
with logs.capture(collected):
|
|
outputs = self.f(params=self.params) or {}
|
|
self._pipeline.publish_log(self, collected, "")
|
|
|
|
return self._pipeline.trigger(self, self._to_messages(outputs))
|
|
|
|
def __call__(self, inputs: dict[str, Any] | None = None) -> NodeResult:
|
|
"""
|
|
Execute the node and trigger downstream nodes if bound to a pipeline.
|
|
|
|
Validates inputs against the node's ``requires`` specification, executes
|
|
the wrapped function, validates outputs, and triggers downstream execution
|
|
if the node is bound to a pipeline.
|
|
|
|
:param inputs: Input values keyed by message name. Must match the node's
|
|
``requires`` specification.
|
|
:type inputs: dict | None
|
|
:returns: Node outputs if successful, or pipeline execution results if bound.
|
|
:rtype: dict | None
|
|
"""
|
|
outputs = self.execute(inputs)
|
|
return self._pipeline.trigger(self, outputs) if self._pipeline else outputs
|