File size: 2,196 Bytes
cdc87cb daca5bd cdc87cb daca5bd cdc87cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Tracing helpers shared by every node."""
import time
from functools import wraps
from typing import Callable, Optional
def trace_event(node: str, *, status: str = "ok", duration_ms: float = 0.0,
summary: str = "", payload: Optional[dict] = None) -> dict:
"""Build a single TraceEvent dict matching the state schema."""
return {
"node": node,
"status": status,
"duration_ms": round(duration_ms, 2),
"summary": summary,
"payload": payload or {},
}
def traced(node_name: str) -> Callable:
"""
Decorator: time a node, swallow exceptions into a trace event, and
guarantee the node always returns a dict containing a `trace` list.
The wrapped function should return a partial state dict WITHOUT a `trace`
field — this decorator injects the timing/status event automatically. If
the wrapped function returns its own `trace` list, those events are kept
and the timing event is appended.
"""
def decorator(fn):
@wraps(fn)
def wrapper(state, *args, **kwargs):
t0 = time.perf_counter()
try:
result = fn(state, *args, **kwargs) or {}
dt_ms = (time.perf_counter() - t0) * 1000
summary = result.pop("_summary", "")
payload = result.pop("_payload", {})
event = trace_event(
node_name,
status="ok",
duration_ms=dt_ms,
summary=summary,
payload=payload,
)
existing = result.get("trace", [])
result["trace"] = existing + [event]
return result
except Exception as e:
dt_ms = (time.perf_counter() - t0) * 1000
event = trace_event(
node_name,
status="error",
duration_ms=dt_ms,
summary=f"{type(e).__name__}: {e}",
payload={"error": str(e)},
)
return {"trace": [event], "error": f"{node_name}: {e}"}
return wrapper
return decorator
|