diff --git a/src/dynavec/agents/__init__.py b/src/dynavec/agents/__init__.py index 262cfa9..4e00e4c 100644 --- a/src/dynavec/agents/__init__.py +++ b/src/dynavec/agents/__init__.py @@ -1,4 +1,4 @@ -"""Agent primitives and execution loops for dynaflow.""" +"""Agent primitives, execution loops, and graph orchestration for dynaflow.""" from __future__ import annotations @@ -22,11 +22,30 @@ mcp_tools_from_session, mcp_tools_from_session_async, ) +from .graph import ( + AgentGraph, + AgentState, + BaseNode, + Edge, + FlowGraph, + FlowState, + FunctionNode, + GraphCycleError, + GraphError, + GraphResult, + GraphState, + GraphValidationError, + MaxIterationsExceededError, + NodeTrace, + append_reducer, + default_reducer, +) from .planner import Planner from .react import ReActAgent from .registry import ToolRegistry, default_registry __all__ = [ + # Existing agent primitives "AgentResult", "AgentStep", "AgentTool", @@ -34,15 +53,34 @@ "PlanStep", "Planner", "ReActAgent", + "tool", + # Tool registry & connectors "ToolRegistry", "create_http_tool", "create_python_code_tool", "default_registry", "mcp_tools_from_session", "mcp_tools_from_session_async", - "tool", + # Checkpointer "Checkpoint", "BaseCheckpointer", "MemoryCheckpointer", "DynamoDBCheckpointer", + # Graph orchestration engine (#277 / #278) + "AgentGraph", + "AgentState", + "BaseNode", + "Edge", + "FlowGraph", + "FlowState", + "FunctionNode", + "GraphCycleError", + "GraphError", + "GraphResult", + "GraphState", + "GraphValidationError", + "MaxIterationsExceededError", + "NodeTrace", + "append_reducer", + "default_reducer", ] diff --git a/src/dynavec/agents/graph.py b/src/dynavec/agents/graph.py new file mode 100644 index 0000000..dfa0ce9 --- /dev/null +++ b/src/dynavec/agents/graph.py @@ -0,0 +1,750 @@ +"""Agent graph orchestration engine for dynaflow. + +This module provides the graph execution runtime: shared state with reducers, +nodes, conditional routing edges, parallel fan-out and join (barrier) execution, +and bounded loops with termination guards and cycle validation. +""" + +from __future__ import annotations + +import ast +import re +import time +from abc import ABC, abstractmethod +from collections import defaultdict, deque +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Any, Callable + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + + +class GraphError(Exception): + """Base exception for agent graph execution and validation errors.""" + + +class GraphValidationError(GraphError): + """Raised when a graph structure is invalid or malformed.""" + + +class GraphCycleError(GraphError): + """Raised when a graph contains an un-guarded or infinite cycle.""" + + +class MaxIterationsExceededError(GraphError): + """Raised when a loop execution exceeds the maximum allowed iterations.""" + + +# --------------------------------------------------------------------------- +# GraphState & Reducers +# --------------------------------------------------------------------------- + +ReducerFn = Callable[[Any, Any], Any] + + +def default_reducer(current: Any, update: Any) -> Any: + """Default state reducer: replaces current value with the update.""" + return update + + +def append_reducer(current: Any, update: Any) -> list[Any]: + """Reducer that appends or extends lists.""" + if current is None: + return list(update) if isinstance(update, (list, tuple)) else [update] + if isinstance(current, list): + if isinstance(update, list): + return current + update + return [*current, update] + return [current, update] + + +class GraphState(dict[str, Any]): + """Typed shared state that flows through the agent graph. + + Supports isolated branch copies and deterministic merging using + custom reducer functions per key. + """ + + def __init__( + self, + initial: dict[str, Any] | None = None, + *, + reducers: dict[str, ReducerFn] | None = None, + ) -> None: + super().__init__(initial or {}) + self._reducers: dict[str, ReducerFn] = reducers or {} + + @property + def reducers(self) -> dict[str, ReducerFn]: + """Mapping of field keys to their reducer functions.""" + return self._reducers + + def register_reducer(self, key: str, reducer: ReducerFn) -> GraphState: + """Register a custom reducer for a state key.""" + self._reducers[key] = reducer + return self + + def set(self, key: str, value: Any) -> GraphState: + """Set a key and return self for fluent chaining.""" + self[key] = value + return self + + def copy(self) -> GraphState: + """Create an isolated deep-copy of the state for parallel branches.""" + new_state = GraphState(dict(self), reducers=dict(self._reducers)) + return new_state + + def merge(self, updates: dict[str, Any] | GraphState) -> None: + """Merge a dictionary of updates using defined reducers.""" + for key, val in updates.items(): + if key in self and key in self._reducers: + self[key] = self._reducers[key](self[key], val) + else: + self[key] = val + + +# Backwards compatibility alias +AgentState = GraphState + + +# --------------------------------------------------------------------------- +# BaseNode & FunctionNode +# --------------------------------------------------------------------------- + + +class BaseNode(ABC): + """Abstract base class for all nodes in an agent graph.""" + + node_type: str = "base" + + def __init__(self, name: str = "") -> None: + self.name = name + + @abstractmethod + def execute(self, state: GraphState) -> GraphState | dict[str, Any] | None: + """Execute the node logic against current state. + + Returns updated state or dict of updates, or None if in-place mutation. + """ + ... + + def to_dict(self) -> dict[str, Any]: + """Serialize node metadata for canvas / inspection.""" + return {"name": self.name, "type": self.node_type} + + +class FunctionNode(BaseNode): + """Node wrapping a callable function: ``fn(state) -> dict | None``.""" + + node_type: str = "function" + + def __init__( + self, + fn: Callable[[GraphState], GraphState | dict[str, Any] | None], + name: str = "", + ) -> None: + super().__init__(name=name or getattr(fn, "__name__", "fn")) + self.fn = fn + + def execute(self, state: GraphState) -> GraphState | dict[str, Any] | None: + return self.fn(state) + + +# --------------------------------------------------------------------------- +# Safe Expression Evaluation +# --------------------------------------------------------------------------- + +_COMPARE_RE = re.compile(r"^([\w][\w.]*)\s*(==|!=|>=|<=|>|<|in|not in)\s*(.+)$") + + +def safe_eval_expr(expr: str, state: dict[str, Any]) -> bool: + """Safely evaluate comparison expressions against state without eval(). + + Supports: + field == value field != value + field > value field >= value + field < value field <= value + field in value field not in value + expr and expr expr or expr + """ + expr = expr.strip() + + # Precedence: 'or' first + if " or " in expr: + left, right = expr.split(" or ", 1) + return safe_eval_expr(left, state) or safe_eval_expr(right, state) + + # Then 'and' + if " and " in expr: + left, right = expr.split(" and ", 1) + return safe_eval_expr(left, state) and safe_eval_expr(right, state) + + m = _COMPARE_RE.match(expr) + if not m: + raise ValueError(f"Cannot parse condition expression: {expr!r}") + + field_name, op, raw_value = m.group(1), m.group(2), m.group(3).strip() + + # Resolve dotted key access (e.g. "result.score") + actual: Any = state + for part in field_name.split("."): + if isinstance(actual, dict): + actual = actual.get(part) + else: + actual = getattr(actual, part, None) + + try: + expected = ast.literal_eval(raw_value) + except (ValueError, SyntaxError) as exc: + raise ValueError(f"Cannot parse value {raw_value!r} in expression: {exc}") from exc + + ops: dict[str, Callable[[Any, Any], bool]] = { + "==": lambda a, b: a == b, + "!=": lambda a, b: a != b, + ">": lambda a, b: a is not None and a > b, + ">=": lambda a, b: a is not None and a >= b, + "<": lambda a, b: a is not None and a < b, + "<=": lambda a, b: a is not None and a <= b, + "in": lambda a, b: a in b, + "not in": lambda a, b: a not in b, + } + return ops[op](actual, expected) + + +# --------------------------------------------------------------------------- +# Edge +# --------------------------------------------------------------------------- + + +@dataclass +class Edge: + """Directed connection between two nodes with optional gating condition. + + Condition can be: + * None: Unconditional transition. + * Callable[[GraphState], bool]: Predicate evaluated against state. + * str: Safe expression string like "confidence >= 0.8". + """ + + source: str + target: str + condition: Callable[[GraphState], bool] | str | None = None + label: str = "" + + def evaluate(self, state: GraphState) -> bool: + """Evaluate if the edge condition is satisfied.""" + if self.condition is None: + return True + if callable(self.condition): + return bool(self.condition(state)) + if isinstance(self.condition, str): + return safe_eval_expr(self.condition, state) + raise TypeError(f"Unsupported edge condition type: {type(self.condition)}") + + def is_conditional(self) -> bool: + """True if edge is guarded by a condition.""" + return self.condition is not None + + def to_dict(self) -> dict[str, Any]: + d: dict[str, Any] = {"source": self.source, "target": self.target} + if self.label: + d["label"] = self.label + if isinstance(self.condition, str): + d["condition"] = self.condition + elif self.condition is not None: + d["condition"] = f"" + return d + + +# --------------------------------------------------------------------------- +# Observability Traces & Results +# --------------------------------------------------------------------------- + + +@dataclass +class NodeTrace: + """Execution telemetry for a single node.""" + + node_name: str + node_type: str + started_at: float + ended_at: float + duration_ms: float + input_keys: list[str] = field(default_factory=list) + output_keys: list[str] = field(default_factory=list) + error: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "node_name": self.node_name, + "node_type": self.node_type, + "duration_ms": round(self.duration_ms, 3), + "input_keys": self.input_keys, + "output_keys": self.output_keys, + "error": self.error, + } + + +@dataclass +class GraphResult: + """Result returned upon completion of an AgentGraph run.""" + + state: GraphState + trace: list[NodeTrace] = field(default_factory=list) + nodes_executed: int = 0 + total_ms: float = 0.0 + graph_name: str = "" + success: bool = True + error: str | None = None + + @property + def output(self) -> GraphState: + """Convenience alias for the final state.""" + return self.state + + def to_dict(self) -> dict[str, Any]: + return { + "graph_name": self.graph_name, + "success": self.success, + "total_ms": round(self.total_ms, 3), + "nodes_executed": self.nodes_executed, + "trace": [t.to_dict() for t in self.trace], + "state": dict(self.state), + "error": self.error, + } + + +# --------------------------------------------------------------------------- +# AgentGraph Orchestration Engine +# --------------------------------------------------------------------------- + + +class AgentGraph: + """Directed graph execution engine supporting conditional routing, + parallel fan-out/join barriers, and bounded cycles with iteration guards. + """ + + def __init__( + self, + name: str = "dynaflow-graph", + *, + description: str = "", + max_iterations: int = 25, + ) -> None: + self.name = name + self.description = description + self.max_iterations = max_iterations + self._nodes: dict[str, BaseNode] = {} + self._edges: list[Edge] = [] + self._entry: str | None = None + self._join_nodes: set[str] = set() + + # ---------------------------------------------------------------- Builder + + def add_node( + self, + name: str, + node: BaseNode | Callable[[GraphState], GraphState | dict[str, Any] | None], + ) -> AgentGraph: + """Register a node in the graph.""" + if name in self._nodes: + raise GraphValidationError(f"Duplicate node name: {name!r}") + if isinstance(node, BaseNode): + node.name = name + self._nodes[name] = node + elif callable(node): + self._nodes[name] = FunctionNode(node, name=name) + else: + raise TypeError(f"Expected BaseNode or callable, got {type(node).__name__}") + return self + + def add_edge( + self, + source: str, + target: str, + *, + condition: Callable[[GraphState], bool] | str | None = None, + label: str = "", + ) -> AgentGraph: + """Add a directed edge between source and target, optionally gated by condition.""" + self._edges.append( + Edge( + source=source, + target=target, + condition=condition, + label=label, + ) + ) + return self + + def add_conditional_edges( + self, + source: str, + router: Callable[[GraphState], str | list[str]] | dict[Any, str], + *, + branches: dict[str, str] | None = None, + ) -> AgentGraph: + """Add conditional routing edges from a source node using a router callable or map.""" + if callable(router): + if branches: + for branch_key, target in branches.items(): + + def _make_cond(key: str) -> Callable[[GraphState], bool]: + def _cond(s: GraphState) -> bool: + res = router(s) + if isinstance(res, list): + return key in res + return res == key + + return _cond + + self.add_edge( + source, + target, + condition=_make_cond(branch_key), + label=str(branch_key), + ) + else: + # Direct router returning node name(s) + # We register an edge to all potential nodes evaluated at runtime + raise GraphValidationError( + "add_conditional_edges with a router function requires the 'branches' mapping." + ) + elif isinstance(router, dict): + for condition_expr, target in router.items(): + if isinstance(condition_expr, str): + self.add_edge(source, target, condition=condition_expr) + else: + raise TypeError("Router mapping keys must be string condition expressions") + return self + + def add_parallel_branches( + self, + source: str, + targets: list[str], + *, + join_at: str | None = None, + ) -> AgentGraph: + """Convenience method to fan-out from source to targets and join at join_at.""" + for target in targets: + self.add_edge(source, target) + if join_at: + self._join_nodes.add(join_at) + for target in targets: + self.add_edge(target, join_at) + return self + + def register_join_barrier(self, node_name: str) -> AgentGraph: + """Designate a node as a join barrier that must wait for all inbound incoming edges.""" + self._join_nodes.add(node_name) + return self + + def set_entry(self, name: str) -> AgentGraph: + """Designate the entry point node of the graph.""" + if name not in self._nodes: + raise GraphValidationError(f"Entry node {name!r} not found in registered nodes") + self._entry = name + return self + + @property + def node_names(self) -> list[str]: + return list(self._nodes.keys()) + + @property + def edges(self) -> list[Edge]: + return list(self._edges) + + # ------------------------------------------------------------- Validation + + def validate(self) -> list[str]: + """Validate graph topology, catching disconnected nodes and infinite malformed cycles.""" + errors: list[str] = [] + + if self._entry is None: + errors.append("No entry node set (call set_entry())") + elif self._entry not in self._nodes: + errors.append(f"Entry node {self._entry!r} not found") + + for edge in self._edges: + if edge.source not in self._nodes: + errors.append(f"Edge source {edge.source!r} not found") + if edge.target not in self._nodes: + errors.append(f"Edge target {edge.target!r} not found") + + if errors: + return errors + + # Detect cycles and distinguish valid guarded loops from infinite malformed cycles + adj: dict[str, list[Edge]] = defaultdict(list) + for e in self._edges: + adj[e.source].append(e) + + # Find cycles using DFS path tracking + visited: set[str] = set() + path: list[str] = [] + path_set: set[str] = set() + + def _check_cycles(node: str) -> None: + visited.add(node) + path.append(node) + path_set.add(node) + + for edge in adj.get(node, []): + neighbor = edge.target + if neighbor in path_set: + # Cycle found! Extract cycle path + idx = path.index(neighbor) + cycle_nodes = path[idx:] + [neighbor] + + # Check if this cycle is guarded: + # A cycle is valid if ANY edge in the cycle has a condition, + # OR any node in the cycle has an outgoing edge exiting the cycle. + cycle_node_set = set(cycle_nodes) + has_exit_or_condition = False + + for cn in cycle_node_set: + for outgoing_edge in adj.get(cn, []): + if outgoing_edge.is_conditional(): + has_exit_or_condition = True + break + if outgoing_edge.target not in cycle_node_set: + has_exit_or_condition = True + break + if has_exit_or_condition: + break + + if not has_exit_or_condition: + cycle_repr = " -> ".join(cycle_nodes) + errors.append( + f"Malformed infinite cycle detected: {cycle_repr} with no exit condition or guard." + ) + elif neighbor not in visited: + _check_cycles(neighbor) + + path.pop() + path_set.remove(node) + + for n in self._nodes: + if n not in visited: + _check_cycles(n) + + return errors + + # -------------------------------------------------------------- Execution + + def run( + self, + initial_state: dict[str, Any] | GraphState | None = None, + *, + max_steps: int = 100, + max_workers: int = 4, + ) -> GraphResult: + """Execute the graph from entry to completion.""" + errors = self.validate() + if errors: + if any("Malformed infinite cycle" in e for e in errors): + raise GraphCycleError("; ".join(errors)) + raise GraphValidationError("; ".join(errors)) + + state: GraphState = ( + initial_state.copy() + if isinstance(initial_state, GraphState) + else GraphState(initial_state or {}) + ) + + traces: list[NodeTrace] = [] + node_visit_counts: dict[str, int] = defaultdict(int) + + # Adjacency and inbound edge mappings + adj: dict[str, list[Edge]] = defaultdict(list) + for e in self._edges: + adj[e.source].append(e) + + # Barrier tracking: for join nodes, track pending inputs + # Auto-detect join barriers: any node with multiple inbound edges is a barrier candidate + inbound_count: dict[str, int] = defaultdict(int) + for e in self._edges: + inbound_count[e.target] += 1 + join_barriers = set(self._join_nodes) | { + target for target, count in inbound_count.items() if count > 1 + } + + barrier_states: dict[str, list[GraphState]] = defaultdict(list) + barrier_expected: dict[str, int] = defaultdict(int) + + t0 = time.perf_counter() + steps = 0 + error: str | None = None + + # Execution queue of ready nodes + # Each item is: node_name + ready_queue: deque[str] = deque([self._entry]) # type: ignore[list-item] + + while ready_queue and steps < max_steps: + # Batch of nodes ready to execute in parallel + current_batch: list[str] = [] + while ready_queue: + current_batch.append(ready_queue.popleft()) + + # Check loop iterations on all nodes in this batch + for node_name in current_batch: + node_visit_counts[node_name] += 1 + if node_visit_counts[node_name] > self.max_iterations: + raise MaxIterationsExceededError( + f"Loop iteration limit ({self.max_iterations}) exceeded on node {node_name!r}" + ) + + steps += len(current_batch) + + # Execute batch: in parallel if len > 1, else sequentially + if len(current_batch) == 1: + node_name = current_batch[0] + node_trace, next_state = self._execute_node(self._nodes[node_name], state) + traces.append(node_trace) + if node_trace.error: + error = node_trace.error + break + state = next_state + # Resolve outgoing edges + self._dispatch_outgoing( + node_name, + state, + adj, + join_barriers, + barrier_states, + barrier_expected, + ready_queue, + ) + else: + # Parallel fan-out + with ThreadPoolExecutor(max_workers=min(len(current_batch), max_workers)) as pool: + futures = [ + pool.submit(self._execute_node, self._nodes[name], state.copy()) + for name in current_batch + ] + batch_results = [f.result() for f in futures] + + # Check errors & merge states + for name, (node_trace, branch_state) in zip(current_batch, batch_results): + traces.append(node_trace) + if node_trace.error: + error = node_trace.error + break + # Merge branch state updates back into main state via reducers + state.merge(branch_state) + + # Dispatch outgoing + self._dispatch_outgoing( + name, + branch_state, + adj, + join_barriers, + barrier_states, + barrier_expected, + ready_queue, + ) + + if error is not None: + break + + total_ms = (time.perf_counter() - t0) * 1000 + success = error is None and all(t.error is None for t in traces) + + return GraphResult( + state=state, + trace=traces, + nodes_executed=steps, + total_ms=total_ms, + graph_name=self.name, + success=success, + error=error, + ) + + def _execute_node( + self, + node: BaseNode, + state: GraphState, + ) -> tuple[NodeTrace, GraphState]: + """Execute a single node safely with telemetry timing.""" + input_keys = list(state.keys()) + t_start = time.perf_counter() + err: str | None = None + new_state = state + + try: + res = node.execute(state) + if isinstance(res, GraphState): + new_state = res + elif isinstance(res, dict): + new_state = state.copy() + new_state.merge(res) + except Exception as exc: # noqa: BLE001 + err = f"{type(exc).__name__}: {exc}" + + t_end = time.perf_counter() + output_keys = list(new_state.keys()) if err is None else input_keys + + trace = NodeTrace( + node_name=node.name, + node_type=node.node_type, + started_at=t_start, + ended_at=t_end, + duration_ms=(t_end - t_start) * 1000, + input_keys=input_keys, + output_keys=output_keys, + error=err, + ) + return trace, new_state + + def _dispatch_outgoing( + self, + source: str, + current_state: GraphState, + adj: dict[str, list[Edge]], + join_barriers: set[str], + barrier_states: dict[str, list[GraphState]], + barrier_expected: dict[str, int], + ready_queue: deque[str], + ) -> None: + """Evaluate outgoing edges from source and schedule next nodes.""" + outgoing = adj.get(source, []) + active_targets: list[str] = [] + + for edge in outgoing: + if edge.evaluate(current_state): + active_targets.append(edge.target) + + for target in active_targets: + if target in join_barriers: + # Barrier synchronization: + # Count inbound edges that lead to this barrier from current or sibling sources + barrier_states[target].append(current_state) + # Count total active predecessors leading to this barrier + total_inbound = sum(1 for e in self._edges if e.target == target) + if len(barrier_states[target]) >= total_inbound or len(active_targets) == 1: + # All incoming branches arrived at barrier + if target not in ready_queue: + ready_queue.append(target) + else: + if target not in ready_queue: + ready_queue.append(target) + + # ---------------------------------------------------------- Serialization + + def to_dict(self) -> dict[str, Any]: + """Serialize graph definition to JSON dict (ReactFlow compatible).""" + return { + "name": self.name, + "description": self.description, + "entry": self._entry, + "max_iterations": self.max_iterations, + "nodes": {name: node.to_dict() for name, node in self._nodes.items()}, + "edges": [edge.to_dict() for edge in self._edges], + } + + +# Fluent / LangGraph style aliases +FlowGraph = AgentGraph +FlowState = GraphState diff --git a/tests/test_flow.py b/tests/test_flow.py new file mode 100644 index 0000000..cb235e7 --- /dev/null +++ b/tests/test_flow.py @@ -0,0 +1,314 @@ +"""Unit tests for dynaflow agent graph engine and control flow (Issues #277 & #278).""" + +from __future__ import annotations + +import time +from typing import Any + +import pytest + +from dynavec.agents import ( + AgentGraph, + BaseNode, + GraphCycleError, + GraphResult, + GraphState, + GraphValidationError, + MaxIterationsExceededError, + append_reducer, +) + +# --------------------------------------------------------------------------- +# Test Helpers & Custom Nodes +# --------------------------------------------------------------------------- + + +class IncrementNode(BaseNode): + """Simple node that increments a counter in state.""" + + node_type = "increment" + + def __init__(self, key: str = "count", amount: int = 1, name: str = "") -> None: + super().__init__(name=name) + self.key = key + self.amount = amount + + def execute(self, state: GraphState) -> dict[str, Any]: + return {self.key: state.get(self.key, 0) + self.amount} + + +# --------------------------------------------------------------------------- +# 1. Core End-to-End Execution (Issue #277) +# --------------------------------------------------------------------------- + + +def test_linear_three_node_graph_execution() -> None: + """Acceptance criteria #277: A graph of >=3 nodes executes end-to-end with state passed.""" + + def step1(state: GraphState) -> dict[str, Any]: + return {"query": state["query"].strip().lower()} + + def step2(state: GraphState) -> dict[str, Any]: + return {"tokens": state["query"].split()} + + def step3(state: GraphState) -> dict[str, Any]: + return {"token_count": len(state["tokens"])} + + graph = ( + AgentGraph("test-linear") + .add_node("clean", step1) + .add_node("tokenize", step2) + .add_node("count", step3) + .add_edge("clean", "tokenize") + .add_edge("tokenize", "count") + .set_entry("clean") + ) + + result = graph.run({"query": " Dynavec Vector DB "}) + + assert isinstance(result, GraphResult) + assert result.success is True + assert result.nodes_executed == 3 + assert result.state["query"] == "dynavec vector db" + assert result.state["tokens"] == ["dynavec", "vector", "db"] + assert result.state["token_count"] == 3 + assert len(result.trace) == 3 + + +def test_state_reducers_and_isolation() -> None: + """Acceptance criteria #277: State updates are isolated and merged via defined reducers.""" + state = GraphState( + {"items": ["a"], "count": 10}, + reducers={ + "items": append_reducer, + "count": lambda a, b: a + b, + }, + ) + + def node_append(s: GraphState) -> dict[str, Any]: + return {"items": ["b", "c"], "count": 5} + + graph = AgentGraph("test-reducers").add_node("append", node_append).set_entry("append") + + result = graph.run(state) + assert result.success is True + # items merged via append_reducer + assert result.state["items"] == ["a", "b", "c"] + # count merged via sum reducer + assert result.state["count"] == 15 + + +# --------------------------------------------------------------------------- +# 2. Conditional Routing & Branches (Issue #278) +# --------------------------------------------------------------------------- + + +def test_conditional_edge_branch_on_predicate() -> None: + """Acceptance criteria #278: A graph can branch on a predicate.""" + graph = ( + AgentGraph("test-branching") + .add_node("start", lambda s: {"val": s.get("input", 0)}) + .add_node("high_path", lambda s: {"route": "HIGH", "score": s["val"] * 2}) + .add_node("low_path", lambda s: {"route": "LOW", "score": s["val"] // 2}) + .add_edge("start", "high_path", condition="val >= 10") + .add_edge("start", "low_path", condition="val < 10") + .set_entry("start") + ) + + # Test high path + res_high = graph.run({"input": 15}) + assert res_high.success is True + assert res_high.state["route"] == "HIGH" + assert res_high.state["score"] == 30 + assert [t.node_name for t in res_high.trace] == ["start", "high_path"] + + # Test low path + res_low = graph.run({"input": 4}) + assert res_low.success is True + assert res_low.state["route"] == "LOW" + assert res_low.state["score"] == 2 + assert [t.node_name for t in res_low.trace] == ["start", "low_path"] + + +def test_conditional_edge_callable_predicate() -> None: + """Verify conditional edges evaluate Python callable predicates.""" + graph = ( + AgentGraph("test-callable-edge") + .add_node("eval", lambda s: {"score": 0.85}) + .add_node("pass_node", lambda s: {"verdict": "APPROVED"}) + .add_node("fail_node", lambda s: {"verdict": "REJECTED"}) + .add_edge("eval", "pass_node", condition=lambda s: s["score"] >= 0.8) + .add_edge("eval", "fail_node", condition=lambda s: s["score"] < 0.8) + .set_entry("eval") + ) + + result = graph.run() + assert result.state["verdict"] == "APPROVED" + assert [t.node_name for t in result.trace] == ["eval", "pass_node"] + + +# --------------------------------------------------------------------------- +# 3. Parallel Branches (Fan-Out + Join Barrier) (Issue #278) +# --------------------------------------------------------------------------- + + +def test_parallel_branches_fan_out_and_join() -> None: + """Acceptance criteria #278: Run two branches in parallel and join.""" + + def branch_a(s: GraphState) -> dict[str, Any]: + time.sleep(0.01) + return {"results": ["from_A"]} + + def branch_b(s: GraphState) -> dict[str, Any]: + time.sleep(0.01) + return {"results": ["from_B"]} + + def join_node(s: GraphState) -> dict[str, Any]: + return {"joined": True, "total_results": len(s.get("results", []))} + + state = GraphState({"results": []}, reducers={"results": append_reducer}) + + graph = ( + AgentGraph("test-parallel-join") + .add_node("start", lambda s: {"initialized": True}) + .add_node("branch_a", branch_a) + .add_node("branch_b", branch_b) + .add_node("join", join_node) + .add_parallel_branches("start", ["branch_a", "branch_b"], join_at="join") + .set_entry("start") + ) + + result = graph.run(state) + + assert result.success is True + assert result.state["joined"] is True + assert sorted(result.state["results"]) == ["from_A", "from_B"] + assert result.state["total_results"] == 2 + executed_nodes = [t.node_name for t in result.trace] + assert "start" == executed_nodes[0] + assert "join" == executed_nodes[-1] + assert set(executed_nodes[1:3]) == {"branch_a", "branch_b"} + + +# --------------------------------------------------------------------------- +# 4. Loops & Cycles with Termination Guards (Issue #278) +# --------------------------------------------------------------------------- + + +def test_bounded_loop_terminates_on_condition() -> None: + """Acceptance criteria #278: A loop terminates on condition and respects iteration cap.""" + # Graph that increments counter in a loop until count >= 3 + graph = ( + AgentGraph("test-loop-termination", max_iterations=10) + .add_node("init", lambda s: {"count": 0}) + .add_node("step", IncrementNode("count", amount=1, name="step")) + .add_node("done", lambda s: {"status": "FINISHED"}) + .add_edge("init", "step") + # Loop back if count < 3 + .add_edge("step", "step", condition="count < 3") + # Exit to done when count >= 3 + .add_edge("step", "done", condition="count >= 3") + .set_entry("init") + ) + + result = graph.run() + + assert result.success is True + assert result.state["count"] == 3 + assert result.state["status"] == "FINISHED" + assert [t.node_name for t in result.trace] == [ + "init", + "step", + "step", + "step", + "done", + ] + + +def test_loop_exceeds_max_iterations_raises_error() -> None: + """Acceptance criteria #278: Respects iteration cap and halts runaway loops.""" + # Never terminates because target is impossible + graph = ( + AgentGraph("test-infinite-loop", max_iterations=5) + .add_node("step", IncrementNode("count", 1, name="step")) + .add_node("exit", lambda s: {"done": True}) + .add_edge("step", "step", condition="count < 100") + .add_edge("step", "exit", condition="count >= 100") + .set_entry("step") + ) + + with pytest.raises(MaxIterationsExceededError) as exc_info: + graph.run({"count": 0}) + + assert "exceeded" in str(exc_info.value).lower() + assert "step" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# 5. Cycle Detection & Validation Errors (Issue #278) +# --------------------------------------------------------------------------- + + +def test_cycle_detection_flags_unconditional_infinite_cycle() -> None: + """Acceptance criteria #278: Cycle detection and clear errors for malformed graphs.""" + # Node A -> Node B -> Node A with no condition or exit + graph = ( + AgentGraph("test-malformed-cycle") + .add_node("node_a", lambda s: s) + .add_node("node_b", lambda s: s) + .add_edge("node_a", "node_b") + .add_edge("node_b", "node_a") # Unconditional back-edge + .set_entry("node_a") + ) + + with pytest.raises(GraphCycleError) as exc_info: + graph.run() + + assert "Malformed infinite cycle detected" in str(exc_info.value) + assert "node_a" in str(exc_info.value) + assert "node_b" in str(exc_info.value) + + +def test_graph_validation_catches_missing_entry_or_targets() -> None: + """Verify validation catches missing nodes or disconnected entries.""" + graph = AgentGraph("invalid-graph") + graph.add_node("n1", lambda s: s) + + # Missing entry + with pytest.raises(GraphValidationError) as exc: + graph.run() + assert "No entry node set" in str(exc.value) + + # Non-existent target in edge + graph.set_entry("n1") + graph.add_edge("n1", "missing_node") + with pytest.raises(GraphValidationError) as exc2: + graph.run() + assert "Edge target 'missing_node' not found" in str(exc2.value) + + +# --------------------------------------------------------------------------- +# 6. Serialization (ReactFlow Compatibility) +# --------------------------------------------------------------------------- + + +def test_graph_serialization_to_dict() -> None: + """Verify graph serialization to dictionary format compatible with UI canvas.""" + graph = ( + AgentGraph("serialize-test", description="Demo workflow") + .add_node("fetch", lambda s: s) + .add_node("format", lambda s: s) + .add_edge("fetch", "format", condition="ready == True", label="on-ready") + .set_entry("fetch") + ) + + serialized = graph.to_dict() + assert serialized["name"] == "serialize-test" + assert serialized["entry"] == "fetch" + assert "fetch" in serialized["nodes"] + assert "format" in serialized["nodes"] + assert len(serialized["edges"]) == 1 + assert serialized["edges"][0]["source"] == "fetch" + assert serialized["edges"][0]["target"] == "format" + assert serialized["edges"][0]["condition"] == "ready == True" + assert serialized["edges"][0]["label"] == "on-ready"