Coverage for src / agent_contracts / visualizer.py: 90%

1"""ContractVisualizer - Architecture Document Generator. 

2 

3Generates comprehensive architecture documentation from NodeRegistry 

4contracts, including Mermaid diagrams for visual representation. 

5""" 

6from __future__ import annotations 

7 

8from typing import TYPE_CHECKING 

9from collections import defaultdict 

10 

11if TYPE_CHECKING: 

12 from langgraph.graph.state import CompiledStateGraph 

13 from agent_contracts.registry import NodeRegistry 

14 from agent_contracts.contracts import NodeContract, TriggerCondition 

15 

16 

17class ContractVisualizer: 

18 """Generate architecture documentation from NodeRegistry contracts. 

19  

20 Example: 

21 visualizer = ContractVisualizer(registry) 

22 doc = visualizer.generate_architecture_doc() 

23 with open("ARCHITECTURE.md", "w") as f: 

24 f.write(doc) 

25 """ 

26 

27 def __init__(self, registry: "NodeRegistry", graph: "CompiledStateGraph | None" = None): 

28 """Initialize visualizer. 

29  

30 Args: 

31 registry: NodeRegistry with registered nodes 

32 graph: Optional compiled LangGraph for flow visualization 

33 """ 

34 self.registry = registry 

35 self.graph = graph 

36 

37 def generate_architecture_doc(self) -> str: 

38 """Generate complete architecture document. 

39  

40 Returns: 

41 Markdown document with all sections 

42 """ 

43 sections = [ 

44 self._generate_header(), 

45 self.generate_langgraph_flow(), # Added LangGraph flow 

46 self.generate_state_slices_section(), 

47 self.generate_hierarchy_diagram(), 

48 self.generate_dataflow_diagram(), 

49 self.generate_trigger_hierarchy(), 

50 self.generate_nodes_reference(), 

51 self._generate_footer(), 

52 ] 

53 return "\n\n---\n\n".join(s for s in sections if s) 

54 

55 def _generate_header(self) -> str: 

56 """Generate document header.""" 

57 return """# 🏗️ Agent Architecture 

58 

59> Auto-generated from `NodeRegistry` contracts""" 

60 

61 def _generate_footer(self) -> str: 

62 """Generate document footer.""" 

63 return "_{Generated by `agent-contracts` visualizer}" 

64 

65 def generate_langgraph_flow(self) -> str: 

66 """Generate LangGraph flow visualization from compiled graph.""" 

67 if not self.graph: 

68 return "" 

69 

70 try: 

71 mermaid = self.graph.get_graph().draw_mermaid() 

72 return f"""## 🔗 LangGraph Node Flow 

73 

74> Auto-generated from compiled LangGraph 

75 

76```mermaid 

77{mermaid} 

78```""" 

79 except Exception: 

80 return "" 

81 

82 

83 # ========================================================================= 

84 # State Slices Section 

85 # ========================================================================= 

86 

87 def generate_state_slices_section(self) -> str: 

88 """Generate State Slices documentation. 

89  

90 Collects all slices from node contracts and creates documentation. 

91 """ 

92 # Collect slices from all nodes 

93 all_reads: set[str] = set() 

94 all_writes: set[str] = set() 

95 

96 for name in self.registry.get_all_nodes(): 

97 contract = self.registry.get_contract(name) 

98 if contract: 

99 all_reads.update(contract.reads) 

100 all_writes.update(contract.writes) 

101 

102 all_slices = sorted(all_reads | all_writes) 

103 

104 if not all_slices: 

105 return "" 

106 

107 # Build table 

108 lines = [ 

109 "## 📦 State Slices", 

110 "", 

111 "State is organized into isolated **slices** for separation of concerns.", 

112 "", 

113 "| Slice | Read By | Written By |", 

114 "|:------|:--------|:-----------|", 

115 ] 

116 

117 for slice_name in all_slices: 

118 readers = [n for n in self.registry.get_all_nodes() 

119 if slice_name in (self.registry.get_contract(n).reads or [])] 

120 writers = [n for n in self.registry.get_all_nodes() 

121 if slice_name in (self.registry.get_contract(n).writes or [])] 

122 

123 readers_str = ", ".join(f"`{r}`" for r in readers) 

124 writers_str = ", ".join(f"`{w}`" for w in writers) 

125 

126 lines.append(f"| `{slice_name}` | {readers_str or '-'} | {writers_str or '-'} |") 

127 

128 return "\n".join(lines) 

129 

130 def _build_slice_relationships(self) -> list[tuple[str, str, str]]: 

131 """Build slice relationships from node I/O patterns.""" 

132 # Count how many nodes transform from one slice to another 

133 transform_counts: dict[tuple[str, str], int] = defaultdict(int) 

134 

135 for name in self.registry.get_all_nodes(): 

136 contract = self.registry.get_contract(name) 

137 if not contract: 

138 continue 

139 

140 for read_slice in contract.reads: 

141 for write_slice in contract.writes: 

142 if read_slice != write_slice: 

143 transform_counts[(read_slice, write_slice)] += 1 

144 

145 # Sort by count and take top relationships 

146 sorted_transforms = sorted( 

147 transform_counts.items(), 

148 key=lambda x: x[1], 

149 reverse=True 

150 )[:10] # Limit to 10 most common to avoid diagram clutter 

151 

152 return [ 

153 (src, dst, f"{count} nodes") 

154 for (src, dst), count in sorted_transforms 

155 ] 

156 

157 # ========================================================================= 

158 # Hierarchy Diagram 

159 # ========================================================================= 

160 

161 def generate_hierarchy_diagram(self) -> str: 

162 """Generate Supervisor-Node hierarchy diagram.""" 

163 # Group nodes by supervisor 

164 supervisors: dict[str, list[str]] = defaultdict(list) 

165 

166 for name in self.registry.get_all_nodes(): 

167 contract = self.registry.get_contract(name) 

168 if contract: 

169 supervisors[contract.supervisor].append(name) 

170 

171 if not supervisors: 

172 return "" 

173 

174 lines = [ 

175 "## 🎯 System Hierarchy", 

176 "", 

177 "```mermaid", 

178 "flowchart TB", 

179 ] 

180 

181 # Generate subgraphs for each supervisor 

182 for sup_name, nodes in sorted(supervisors.items()): 

183 safe_sup = self._safe_id(sup_name) 

184 lines.append(f' subgraph {safe_sup}["🎯 {sup_name.replace("_", " ").title()}"]') 

185 lines.append(" direction LR") 

186 

187 for node_name in nodes: 

188 contract = self.registry.get_contract(node_name) 

189 icon = self._get_node_icon(contract) 

190 safe_node = self._safe_id(node_name) 

191 lines.append(f' {safe_node}["{icon} {node_name}"]') 

192 

193 lines.append(" end") 

194 

195 # Add terminal node styling 

196 terminal_nodes = [ 

197 self._safe_id(n) for n in self.registry.get_all_nodes() 

198 if self.registry.get_contract(n) and self.registry.get_contract(n).is_terminal 

199 ] 

200 if terminal_nodes: 

201 lines.append("") 

202 lines.append(" classDef terminal fill:#e94560,stroke:#16213e,color:#fff") 

203 lines.append(f" class {','.join(terminal_nodes)} terminal") 

204 

205 lines.append("```") 

206 

207 return "\n".join(lines) 

208 

209 def _get_node_icon(self, contract: "NodeContract | None") -> str: 

210 """Get emoji icon for node based on its properties. 

211  

212 Combines icons: [LLM?] + [Terminal/Domain Type] 

213 """ 

214 if not contract: 

215 return "📦" 

216 

217 # 1. Determine base icon (Type/Domain) 

218 if contract.is_terminal: 

219 base_icon = "🔚" 

220 elif "search" in contract.name.lower(): 

221 base_icon = "🔍" 

222 elif "interview" in contract.name.lower(): 

223 base_icon = "💬" 

224 elif "like" in contract.name.lower() or "heart" in contract.name.lower(): 

225 base_icon = "❤️" 

226 elif "card" in contract.name.lower(): 

227 base_icon = "🃏" 

228 else: 

229 base_icon = "📦" 

230 

231 # 2. Add LLM indicator if applicable 

232 if contract.requires_llm: 

233 return f"🤖{base_icon}" 

234 

235 return base_icon 

236 

237 def _safe_id(self, name: str) -> str: 

238 """Convert name to safe Mermaid ID.""" 

239 return name.replace("-", "_").replace(" ", "_") 

240 

241 # ========================================================================= 

242 # Data Flow Diagram 

243 # ========================================================================= 

244 

245 def generate_dataflow_diagram(self) -> str: 

246 """Generate data flow diagram showing key data paths. 

247  

248 Shows: 

249 - Slices as central data stores 

250 - Entry points (nodes reading request) 

251 - Terminal nodes (marked with is_terminal) 

252 - Cross-supervisor data connections 

253 """ 

254 # Group nodes by supervisor 

255 supervisors: dict[str, list[str]] = defaultdict(list) 

256 for name in self.registry.get_all_nodes(): 

257 contract = self.registry.get_contract(name) 

258 if contract: 

259 supervisors[contract.supervisor].append(name) 

260 

261 if not supervisors: 

262 return "" 

263 

264 lines = [ 

265 "## 🔀 Data Flow", 

266 "", 

267 "> Key data paths through the system", 

268 "", 

269 "```mermaid", 

270 "flowchart TB", 

271 ] 

272 

273 # Find key nodes: entry points (read request) and terminals 

274 entry_nodes: list[str] = [] 

275 terminal_nodes: list[str] = [] 

276 for name in self.registry.get_all_nodes(): 

277 contract = self.registry.get_contract(name) 

278 if contract: 

279 if "request" in contract.reads: 

280 entry_nodes.append(name) 

281 if contract.is_terminal: 

282 terminal_nodes.append(name) 

283 

284 # Add slices (excluding _internal) 

285 key_slices = ["request", "response"] 

286 for name in self.registry.get_all_nodes(): 

287 contract = self.registry.get_contract(name) 

288 if contract: 

289 for s in contract.reads + contract.writes: 

290 if not s.startswith('_') and s not in key_slices: 

291 key_slices.append(s) 

292 

293 # Slices subgraph 

294 lines.append(' subgraph slices["📦 State"]') 

295 for slice_name in sorted(set(key_slices)): 

296 icon = "📥" if slice_name == "request" else "📤" if slice_name == "response" else "📁" 

297 slice_id = f"slice_{self._safe_id(slice_name)}" 

298 lines.append(f' {slice_id}[("{icon} {slice_name}")]') 

299 lines.append(" end") 

300 lines.append("") 

301 

302 # Add supervisor subgraphs with their nodes 

303 for sup_name, nodes in sorted(supervisors.items()): 

304 safe_sup = f"sup_{self._safe_id(sup_name)}" 

305 lines.append(f' subgraph {safe_sup}["🎯 {sup_name}"]') 

306 lines.append(" direction LR") 

307 for node_name in nodes: 

308 contract = self.registry.get_contract(node_name) 

309 icon = self._get_node_icon(contract) 

310 safe_node = self._safe_id(node_name) 

311 lines.append(f' {safe_node}["{icon} {node_name}"]') 

312 lines.append(" end") 

313 

314 lines.append("") 

315 

316 # Add edges: request -> entry nodes 

317 lines.append(" %% Entry points") 

318 for node_name in entry_nodes: # Show all entry points 

319 lines.append(f" slice_request --> {self._safe_id(node_name)}") 

320 

321 # Add edges: nodes writing to response -> response 

322 lines.append(" %% Response outputs") 

323 response_writers: list[str] = [] 

324 for name in self.registry.get_all_nodes(): 

325 contract = self.registry.get_contract(name) 

326 if contract and "response" in contract.writes: 

327 response_writers.append(name) 

328 for node_name in response_writers: # Show all response writers 

329 lines.append(f" {self._safe_id(node_name)} --> slice_response") 

330 

331 # Add edges for other slices (excluding request/response/_internal) 

332 lines.append(" %% Slice data flows") 

333 other_slices = [s for s in key_slices if s not in ["request", "response"]] 

334 for slice_name in other_slices: 

335 slice_id = f"slice_{self._safe_id(slice_name)}" 

336 # Find readers and writers for this slice 

337 readers: list[str] = [] 

338 writers: list[str] = [] 

339 for name in self.registry.get_all_nodes(): 

340 contract = self.registry.get_contract(name) 

341 if contract: 

342 if slice_name in contract.reads: 

343 readers.append(name) 

344 if slice_name in contract.writes: 

345 writers.append(name) 

346 # Add writer -> slice edges 

347 for node_name in writers: 

348 lines.append(f" {self._safe_id(node_name)} -.-> {slice_id}") 

349 # Add slice -> reader edges 

350 for node_name in readers: 

351 lines.append(f" {slice_id} -.-> {self._safe_id(node_name)}") 

352 

353 # Styling 

354 lines.extend([ 

355 "", 

356 " classDef slice fill:#f5f5f5,stroke:#999", 

357 " classDef terminal fill:#e94560,stroke:#16213e,color:#fff", 

358 ]) 

359 terminal_ids = [self._safe_id(n) for n in terminal_nodes] 

360 if terminal_ids: 

361 lines.append(f" class {','.join(terminal_ids)} terminal") 

362 

363 lines.append("```") 

364 

365 # Add detailed node dependencies in collapsible section 

366 lines.extend([ 

367 "", 

368 "<details>", 

369 "<summary>📊 Detailed Node Dependencies</summary>", 

370 "", 

371 ]) 

372 

373 # Build dependency table grouped by supervisor 

374 dependencies = self.registry.analyze_data_flow() 

375 

376 for sup_name, nodes in sorted(supervisors.items()): 

377 sup_deps = [(n, dependencies.get(n, [])) for n in nodes if dependencies.get(n)] 

378 if not sup_deps: 

379 continue 

380 

381 lines.append(f"**{sup_name}**") 

382 lines.append("") 

383 lines.append("| Node | Depends On (via shared slices) |") 

384 lines.append("|:-----|:-------------------------------|") 

385 

386 for node_name, deps in sup_deps: 

387 if deps: 

388 # Find shared slices for each dependency 

389 dep_info = [] 

390 node_contract = self.registry.get_contract(node_name) 

391 for dep in deps: # Show all dependencies 

392 dep_contract = self.registry.get_contract(dep) 

393 if node_contract and dep_contract: 

394 shared = set(node_contract.reads) & set(dep_contract.writes) 

395 shared_str = ", ".join(sorted(s for s in shared if not s.startswith('_'))) 

396 if shared_str: 

397 dep_info.append(f"`{dep}` ({shared_str})") 

398 else: 

399 dep_info.append(f"`{dep}`") 

400 if dep_info: 

401 lines.append(f"| `{node_name}` | {', '.join(dep_info)} |") 

402 

403 lines.append("") 

404 

405 lines.extend([ 

406 "</details>", 

407 ]) 

408 

409 return "\n".join(lines) 

410 

411 # ========================================================================= 

412 # Trigger Hierarchy 

413 # ========================================================================= 

414 

415 def generate_trigger_hierarchy(self) -> str: 

416 """Generate trigger hierarchy documentation.""" 

417 # Group by supervisor 

418 supervisors: dict[str, list[tuple[str, "NodeContract"]]] = defaultdict(list) 

419 

420 for name in self.registry.get_all_nodes(): 

421 contract = self.registry.get_contract(name) 

422 if contract: 

423 supervisors[contract.supervisor].append((name, contract)) 

424 

425 if not supervisors: 

426 return "" 

427 

428 lines = [ 

429 "## ⚡ Trigger Hierarchy", 

430 "", 

431 "> Nodes are evaluated by **priority** (highest first)", 

432 ] 

433 

434 for sup_name, nodes in sorted(supervisors.items()): 

435 lines.extend([ 

436 "", 

437 f"### 🎯 {sup_name.replace('_', ' ').title()}", 

438 "", 

439 "| Priority | Node | Condition | Hint |", 

440 "|:--------:|:-----|:----------|:-----|", 

441 ]) 

442 

443 # Sort nodes by highest trigger priority 

444 def get_max_priority(item: tuple[str, "NodeContract"]) -> int: 

445 _, contract = item 

446 if not contract.trigger_conditions: 

447 return -1 

448 return max(c.priority for c in contract.trigger_conditions) 

449 

450 sorted_nodes = sorted(nodes, key=get_max_priority, reverse=True) 

451 

452 for name, contract in sorted_nodes: 

453 priority = get_max_priority((name, contract)) 

454 priority_icon = self._get_priority_icon(priority) 

455 

456 # Get condition summary 

457 condition = self._summarize_conditions(contract.trigger_conditions) 

458 hint = contract.get_llm_hints()[0] if contract.get_llm_hints() else "-" 

459 

460 priority_str = f"{priority_icon} **{priority}**" if priority >= 0 else "🟢 _(default)_" 

461 lines.append(f"| {priority_str} | `{name}` | {condition} | {hint} |") 

462 

463 # Add Mermaid diagram for this supervisor 

464 lines.extend([ 

465 "", 

466 "<details>", 

467 f"<summary>📊 {sup_name} Priority Chain</summary>", 

468 "", 

469 "```mermaid", 

470 "flowchart TD", 

471 f' subgraph {self._safe_id(sup_name)}["{sup_name}"]', 

472 " direction TB", 

473 ]) 

474 

475 prev_node = None 

476 for name, contract in sorted_nodes: 

477 priority = get_max_priority((name, contract)) 

478 priority_icon = self._get_priority_icon(priority) 

479 safe_name = self._safe_id(name) 

480 lines.append(f' {safe_name}["{priority_icon} P{priority}: {name}"]') 

481 

482 if prev_node: 

483 lines.append(f' {prev_node} -->|"not matched"| {safe_name}') 

484 prev_node = safe_name 

485 

486 lines.extend([ 

487 " end", 

488 "```", 

489 "", 

490 "</details>", 

491 ]) 

492 

493 return "\n".join(lines) 

494 

495 def _get_priority_icon(self, priority: int) -> str: 

496 """Get priority indicator icon.""" 

497 if priority >= 100: 

498 return "🔴" 

499 elif priority >= 50: 

500 return "🟡" 

501 elif priority >= 10: 

502 return "🟢" 

503 else: 

504 return "⚪" 

505 

506 def _summarize_conditions(self, conditions: list["TriggerCondition"]) -> str: 

507 """Summarize trigger conditions for display.""" 

508 if not conditions: 

509 return "_(default)_" 

510 

511 # Get highest priority condition 

512 highest = max(conditions, key=lambda c: c.priority) 

513 

514 parts = [] 

515 if highest.when: 

516 for key, value in list(highest.when.items())[:2]: 

517 parts.append(f"`{key}={value}`") 

518 if highest.when_not: 

519 for key, value in list(highest.when_not.items())[:1]: 

520 parts.append(f"`{key}≠{value}`") 

521 

522 return ", ".join(parts) if parts else "_(always)_" 

523 

524 # ========================================================================= 

525 # Nodes Reference 

526 # ========================================================================= 

527 

528 def generate_nodes_reference(self) -> str: 

529 """Generate nodes reference table.""" 

530 nodes = self.registry.get_all_nodes() 

531 

532 if not nodes: 

533 return "" 

534 

535 lines = [ 

536 "## 📚 Nodes Reference", 

537 "", 

538 "| Node | Supervisor | Reads | Writes | LLM | Terminal |", 

539 "|:-----|:-----------|:------|:-------|:---:|:--------:|", 

540 ] 

541 

542 for name in sorted(nodes): 

543 contract = self.registry.get_contract(name) 

544 if not contract: 

545 continue 

546 

547 reads = ", ".join(f"`{r}`" for r in contract.reads) or "-" 

548 writes = ", ".join(f"`{w}`" for w in contract.writes) or "-" 

549 llm = "✅" if contract.requires_llm else "" 

550 terminal = "🔚" if contract.is_terminal else "" 

551 

552 lines.append(f"| `{name}` | {contract.supervisor} | {reads} | {writes} | {llm} | {terminal} |") 

553 

554 lines.extend([ 

555 "", 

556 "<details>", 

557 "<summary>🔍 Legend</summary>", 

558 "", 

559 "- ✅ = Requires LLM", 

560 "- 🔚 = Terminal node (exits to END)", 

561 "- Reads/Writes = State slices accessed", 

562 "", 

563 "</details>", 

564 ]) 

565 

566 return "\n".join(lines)