Coverage for agentos/orchestration/graph_executor.py: 32%
185 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 01:20 +0800
1"""
2Agent Graph — DAG-based multi-agent execution engine.
4Build complex agent pipelines as directed acyclic graphs where each node
5is an agent invocation and edges define data flow dependencies.
7v1.16.1: Integrated MetricsCollector for per-node execution tracking.
8"""
10from __future__ import annotations
12import time
13from collections import deque
14from collections.abc import Callable
15from dataclasses import dataclass, field
16from enum import Enum
17from typing import Any
19from agentos.tools.metrics import MetricsCollector
22class GraphNodeState(Enum):
23 """Execution state of a graph orchestrator node."""
25 PENDING = "pending"
26 RUNNING = "running"
27 COMPLETED = "completed"
28 FAILED = "failed"
29 SKIPPED = "skipped"
32@dataclass
33class GraphNode:
34 """A single node in the agent execution graph."""
36 name: str
37 agent_type: str
38 task_template: str
39 """Template string with {input} or {node_name.output} placeholders."""
41 depends_on: list[str] = field(default_factory=list)
42 """Node names this node depends on."""
44 timeout_seconds: float = 120.0
45 retry_count: int = 0
46 on_failure: str = "abort"
47 """Action on failure: 'abort', 'skip', 'continue'."""
49 state: GraphNodeState = GraphNodeState.PENDING
50 output: Any = None
51 error: str | None = None
52 latency_ms: float = 0.0
54 def resolve_task(self, node_outputs: dict[str, Any]) -> str:
55 """Resolve template placeholders using outputs from completed nodes."""
56 task = self.task_template
57 task = task.replace("{input}", str(node_outputs.get("__input__", "")))
58 for name, output in node_outputs.items():
59 placeholder = f"{{{name}.output}}"
60 if placeholder in task:
61 task = task.replace(placeholder, str(output))
62 return task
65@dataclass
66class GraphResult:
67 """Result of graph execution."""
69 node_results: dict[str, GraphNode] = field(default_factory=dict)
70 execution_order: list[str] = field(default_factory=list)
71 total_latency_ms: float = 0.0
72 success: bool = True
73 error: str | None = None
76class AgentGraph:
77 """
78 DAG-based multi-agent execution engine.
80 Define execution graphs declaratively, resolve dependencies automatically,
81 execute nodes in topological order with parallelism for independent nodes.
83 Example::
85 graph = AgentGraph()
86 graph.add_node(GraphNode(
87 name="research",
88 agent_type="researcher",
89 task_template="Research: {input}"
90 ))
91 graph.add_node(GraphNode(
92 name="summarize",
93 agent_type="summarizer",
94 task_template="Summarize: {research.output}",
95 depends_on=["research"]
96 ))
97 result = graph.execute("quantum computing advances")
98 """
100 def __init__(
101 self,
102 executor: Callable[[str, str], Any] | None = None,
103 metrics: MetricsCollector | None = None,
104 ):
105 """
106 Args:
107 executor: Callable(agent_type, task) -> output. If not provided,
108 subclasses must override _execute_node.
109 metrics: Optional MetricsCollector for per-node execution tracking.
110 """
111 self._nodes: dict[str, GraphNode] = {}
112 self._executor = executor
113 self._metrics = metrics
115 def add_node(self, node: GraphNode) -> None:
116 """Add a node to the graph. Raises ValueError on duplicate name."""
117 if node.name in self._nodes:
118 raise ValueError(f"Duplicate node name: {node.name}")
119 self._nodes[node.name] = node
121 def remove_node(self, name: str) -> None:
122 """Remove a node and all edges referencing it."""
123 if name not in self._nodes:
124 raise KeyError(f"Node not found: {name}")
125 del self._nodes[name]
126 for node in self._nodes.values():
127 node.depends_on = [d for d in node.depends_on if d != name]
129 def validate(self) -> list[str]:
130 """
131 Validate graph integrity.
133 Returns:
134 List of error messages (empty if valid).
135 """
136 errors: list[str] = []
138 for name, node in self._nodes.items():
139 for dep in node.depends_on:
140 if dep not in self._nodes:
141 errors.append(f"Node '{name}' depends on unknown node '{dep}'")
142 if dep == name:
143 errors.append(f"Node '{name}' cannot depend on itself")
145 # Check for cycles using topological sort
146 if not errors:
147 try:
148 self._topological_order()
149 except ValueError as e:
150 errors.append(str(e))
152 return errors
154 def _topological_order(self) -> list[str]:
155 """Return nodes in topological order. Raises ValueError on cycle."""
156 in_degree: dict[str, int] = {name: 0 for name in self._nodes}
157 adjacency: dict[str, list[str]] = {name: [] for name in self._nodes}
159 for name, node in self._nodes.items():
160 for dep in node.depends_on:
161 adjacency[dep].append(name)
162 in_degree[name] += 1
164 queue = deque([name for name, deg in in_degree.items() if deg == 0])
165 order: list[str] = []
167 while queue:
168 current = queue.popleft()
169 order.append(current)
170 for neighbor in adjacency[current]:
171 in_degree[neighbor] -= 1
172 if in_degree[neighbor] == 0:
173 queue.append(neighbor)
175 if len(order) != len(self._nodes):
176 remaining = set(self._nodes) - set(order)
177 raise ValueError(f"Cycle detected involving nodes: {remaining}")
179 return order
181 def execute(self, input_data: str) -> GraphResult:
182 """
183 Execute the graph with given input.
185 Args:
186 input_data: Initial task input, accessible as {input} in templates.
188 Returns:
189 GraphResult with per-node outputs and execution metadata.
190 """
191 errors = self.validate()
192 if errors:
193 return GraphResult(success=False, error="; ".join(errors))
195 t0 = time.perf_counter()
196 node_outputs: dict[str, Any] = {"__input__": input_data}
197 results: dict[str, GraphNode] = {}
198 order: list[str] = []
200 # Reset all nodes
201 for node in self._nodes.values():
202 node.state = GraphNodeState.PENDING
203 node.output = None
204 node.error = None
205 node.latency_ms = 0.0
207 try:
208 topo = self._topological_order()
209 except ValueError as e:
210 return GraphResult(success=False, error=str(e))
212 abort = False
213 for name in topo:
214 if abort:
215 self._nodes[name].state = GraphNodeState.SKIPPED
216 results[name] = self._nodes[name]
217 continue
219 node = self._nodes[name]
220 results[name] = node
221 order.append(name)
223 # Check dependencies
224 deps_failed = False
225 for dep in node.depends_on:
226 if results[dep].state == GraphNodeState.FAILED:
227 deps_failed = True
228 break
230 if deps_failed:
231 node.state = GraphNodeState.SKIPPED
232 continue
234 task = node.resolve_task(node_outputs)
235 node_t0 = time.perf_counter()
237 try:
238 node.state = GraphNodeState.RUNNING
239 output = self._execute_node(node.agent_type, task)
240 node.output = output
241 node.state = GraphNodeState.COMPLETED
242 node_outputs[name] = output
243 except Exception as exc:
244 node.state = GraphNodeState.FAILED
245 node.error = f"{type(exc).__name__}: {exc}"
246 node_outputs[name] = None
247 if node.on_failure == "abort":
248 abort = True
250 node.latency_ms = (time.perf_counter() - node_t0) * 1000
252 # Metrics: track per-node execution
253 if self._metrics is not None:
254 self._metrics.get_counter("graph_nodes_total").inc(node.agent_type)
255 self._metrics.get_counter(f"graph_node_{node.state.value}").inc(node.agent_type)
256 self._metrics.get_timer("graph_node_latency_ms").record(node.latency_ms)
258 success = all(
259 n.state in (GraphNodeState.COMPLETED, GraphNodeState.SKIPPED) for n in results.values()
260 )
261 total_latency = (time.perf_counter() - t0) * 1000
263 return GraphResult(
264 node_results=results,
265 execution_order=order,
266 total_latency_ms=total_latency,
267 success=success,
268 )
270 def _execute_node(self, agent_type: str, task: str) -> Any:
271 """Execute a single node. Override or provide executor callback."""
272 if self._executor:
273 return self._executor(agent_type, task)
274 raise NotImplementedError(
275 "No executor provided. Pass executor to __init__ or override _execute_node."
276 )
278 def to_mermaid(self) -> str:
279 """Export graph as Mermaid flowchart."""
280 lines = ["graph TD"]
281 for name, node in self._nodes.items():
282 safe = name.replace("-", "_").replace(" ", "_")
283 lines.append(f' {safe}["{name}\\n({node.agent_type})"]')
284 for name, node in self._nodes.items():
285 safe = name.replace("-", "_").replace(" ", "_")
286 for dep in node.depends_on:
287 safe_dep = dep.replace("-", "_").replace(" ", "_")
288 lines.append(f" {safe_dep} --> {safe}")
289 return "\n".join(lines)
291 @property
292 def node_count(self) -> int:
293 return len(self._nodes)
295 @property
296 def edge_count(self) -> int:
297 return sum(len(n.depends_on) for n in self._nodes.values())
300@dataclass
301class GraphRecipe:
302 """Declarative graph definition (YAML-friendly)."""
304 name: str
305 description: str = ""
306 nodes: list[dict[str, Any]] = field(default_factory=list)
307 """List of node dicts with keys: name, agent_type, task_template, depends_on, timeout_seconds, on_failure."""
309 @classmethod
310 def from_dict(cls, data: dict) -> GraphRecipe:
311 return cls(
312 name=data.get("name", "unnamed"),
313 description=data.get("description", ""),
314 nodes=data.get("nodes", []),
315 )
317 def build(self, executor: Callable | None = None) -> AgentGraph:
318 """Build an AgentGraph from this recipe."""
319 graph = AgentGraph(executor=executor)
320 for spec in self.nodes:
321 graph.add_node(
322 GraphNode(
323 name=spec["name"],
324 agent_type=spec.get("agent_type", "default"),
325 task_template=spec["task_template"],
326 depends_on=spec.get("depends_on", []),
327 timeout_seconds=spec.get("timeout_seconds", 120.0),
328 retry_count=spec.get("retry_count", 0),
329 on_failure=spec.get("on_failure", "abort"),
330 )
331 )
332 return graph