engine.py

"""THE GRAPH — nodes, edges, waves. This file is the whole mechanism.

A graph workflow is a map of steps: each node does one job, each edge says
"when this finishes, go there". The engine runs it in waves:

    while some nodes are ready:
        run every ready node (at the same time, if there are several)
        merge what they wrote into the shared state
        fire their edges / ask their routers where to go next

That's it. Three ideas carry everything:

  state    one plain dict (the blackboard). Every node reads it, returns the
           keys it wants to merge. Parallel nodes must write DISJOINT keys —
           a collision raises instead of silently losing a write.
  routers  plain Python functions over state. Models write state; code reads
           it and picks the edge. No LLM ever decides control flow directly.
  guards   the loop's two-guardrail pattern, generalized: per-node max_visits
           (bounded cycles) + global max_steps (never spin forever). A node
           exception is recorded and surfaced, never raised out of the run —
           same surface-don't-crash rule as ToolRegistry.execute.

Waves trade a little pipelining for a lot of legibility: execution order is
deterministic, so traces read the same way twice and evals can pin the path.
"""

from __future__ import annotations

import threading
import time
from collections import defaultdict
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Any

from prometheus.loop.agent import error_text

# Same observer protocol as the loop: notify(kind, event). Graph runs emit
# graph_start / node_start / node_end / route / graph_end, and pass through
# whatever a node emits (an agent_node's llm/tool events) tagged with node=.
Observer = Callable[[str, dict], None]

START = "START"
END = "END"

NodeFn = Callable[[dict], dict]  # reads state, returns keys to merge
RouteFn = Callable[[dict], str]  # reads state, returns a target label


class GraphStateCollision(Exception):
    """Two nodes in the same wave wrote the same key — a graph bug, not a race."""


@dataclass
class Node:
    name: str
    fn: NodeFn
    kind: str = "fn"          # tool / llm / agent / router label for describe()
    max_visits: int = 1       # >1 only on nodes inside an intended cycle
    on_error: str | None = None  # node to jump to if fn raises (default: drain to END)


@dataclass
class Graph:
    name: str
    nodes: dict[str, Node] = field(default_factory=dict)
    edges: list[tuple[str, str]] = field(default_factory=list)
    routers: dict[str, tuple[RouteFn, dict[str, str]]] = field(default_factory=dict)

    def add_node(self, node: Node) -> None:
        if node.name in (START, END):
            raise ValueError(f"'{node.name}' is reserved")
        self.nodes[node.name] = node

    def add_edge(self, src: str, dst: str) -> None:
        """Unconditional: when src finishes, dst gets one step closer to ready.
        Add nodes before edges — unknown endpoints fail here, not mid-run."""
        for end in (src, dst):
            if end not in self.nodes and end not in (START, END):
                raise ValueError(f"unknown node '{end}'")
        self.edges.append((src, dst))

    def add_router(self, src: str, route: RouteFn, targets: dict[str, str]) -> None:
        """Conditional: when src finishes, `route(state)` returns a label and
        execution jumps to targets[label]. Routers are code, never models."""
        if src not in self.nodes:
            raise ValueError(f"unknown node '{src}'")
        for label, dst in targets.items():
            if dst not in self.nodes and dst != END:
                raise ValueError(f"router target '{dst}' (label '{label}') is unknown")
        self.routers[src] = (route, targets)

    def describe(self) -> dict[str, Any]:
        """The topology as data — what the dashboard draws. Rendering from this
        (never from a hand-copied picture) is what keeps the chart honest."""
        edges = [{"src": s, "dst": d, "conditional": False} for s, d in self.edges]
        for src, (_route, targets) in self.routers.items():
            edges += [{"src": src, "dst": dst, "conditional": True}
                      for dst in dict.fromkeys(targets.values())]
        return {"name": self.name,
                "nodes": [{"name": n.name, "kind": n.kind} for n in self.nodes.values()],
                "edges": edges}


def run_graph(graph: Graph, state: dict, observer: Observer | None = None,
              max_steps: int = 25) -> dict:
    """Run one graph workflow to completion. Returns the final state; errors
    land in state["errors"] and the graph_end event, they are never raised."""
    lock = threading.Lock()
    raw = observer or (lambda kind, ev: None)

    def notify(kind: str, ev: dict) -> None:
        with lock:  # pool threads share the tracer's file append — keep lines whole
            raw(kind, ev)

    deps: dict[str, set] = defaultdict(set)      # static in-edges per node
    for src, dst in graph.edges:
        if dst != END:
            deps[dst].add(src)
    fired: dict[str, set] = defaultdict(set)     # which in-edges have fired
    runs: dict[str, int] = defaultdict(int)
    path: list[str] = []
    errors: dict[str, str] = state.setdefault("errors", {})
    t0 = time.perf_counter()
    notify("graph_start", {"workflow": graph.name, "nodes": list(graph.nodes)})

    def next_wave(jumps: list[str]) -> list[str]:
        """Forcible router/error jumps + nodes whose static deps have all fired."""
        wave: list[str] = []
        for name in jumps + [n for n in graph.nodes
                             if deps[n] and deps[n] <= fired[n] and runs[n] == 0]:
            if name == END or name in wave:
                continue
            if runs[name] >= graph.nodes[name].max_visits:
                errors.setdefault(name, f"max_visits={graph.nodes[name].max_visits} reached")
                continue
            wave.append(name)
        return wave

    def run_one(name: str) -> tuple[str, dict | None, str | None, int]:
        node = graph.nodes[name]
        snapshot = dict(state)
        snapshot["_notify"] = lambda kind, ev, _n=name: notify(kind, {**ev, "node": _n})
        t = time.perf_counter()
        try:
            out = node.fn(snapshot)
            return name, out or {}, None, int((time.perf_counter() - t) * 1000)
        except Exception as exc:  # surface, don't crash — the run drains cleanly
            # error_text, not repr(exc): this string is shown to whoever ran
            # the workflow. A provider that refused the call already wrote a
            # sentence for them ("Free tier used up. Add your own key in
            # Models."); repr() would hand them
            # `Refused('Free tier used up...')` instead, and with a real SDK
            # error it drags the raw JSON error body onto the screen too.
            return name, None, error_text(exc), int((time.perf_counter() - t) * 1000)

    for src, dst in graph.edges:  # START fires its edges before the first wave
        if src == START:
            fired[dst].add(START)
    wave = next_wave([])

    while wave:
        if len(path) + len(wave) > max_steps:
            errors.setdefault("engine", f"max_steps={max_steps} reached")
            break
        for name in wave:
            runs[name] += 1
            notify("node_start", {"workflow": graph.name, "node": name, "visit": runs[name]})
        if len(wave) == 1:  # solo node runs on this thread: trace order stays exact
            results = [run_one(wave[0])]
        else:
            with ThreadPoolExecutor(max_workers=len(wave)) as pool:
                results = [f.result() for f in [pool.submit(run_one, n) for n in wave]]

        jumps: list[str] = []
        wave_writes: dict[str, str] = {}
        for name, out, error, ms in results:  # merge in wave order — deterministic
            path.append(name)
            keys = [k for k in (out or {}) if not k.startswith("_")]
            for key in keys:
                if wave_writes.get(key, name) != name:
                    raise GraphStateCollision(
                        f"'{name}' and '{wave_writes[key]}' both wrote '{key}' — "
                        f"parallel branches must write disjoint keys")
                wave_writes[key] = name
                state[key] = out[key]
            notify("node_end", {"workflow": graph.name, "node": name,
                                "ms": ms, "keys": keys, "error": error})
            if error:
                errors[name] = error
                if graph.nodes[name].on_error:
                    jumps.append(graph.nodes[name].on_error)
                continue  # no on_error → fire nothing; the run drains to END
            if name in graph.routers:
                route, targets = graph.routers[name]
                label = route(state)
                target = targets.get(label)
                notify("route", {"workflow": graph.name, "router": name,
                                 "target": target or END, "reason": label})
                if target is None:
                    errors[name] = f"router returned unknown label '{label}'"
                elif target != END:
                    jumps.append(target)  # a route is a jump, not a dependency
            else:
                for src, dst in graph.edges:
                    if src == name and dst != END:
                        fired[dst].add(name)
        wave = next_wave(jumps)

    first_error = next(iter(errors.values()), None)
    notify("graph_end", {"workflow": graph.name, "ms": int((time.perf_counter() - t0) * 1000),
                         "steps": len(path), "path": path, "error": first_error})
    return state