Coverage for src / agent_contracts / registry.py: 73%
117 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-09 00:42 +0900
1"""NodeRegistry - Node registration and management.
3Registers all ModularNodes and provides routing map generation,
4data flow analysis, and graph construction support.
5"""
6from __future__ import annotations
8from typing import Any
10from agent_contracts.contracts import NodeContract, TriggerCondition
11from agent_contracts.utils.logging import get_logger
13logger = get_logger("agent_contracts.registry")
16class NodeRegistry:
17 """Registry for node registration and management.
19 Example:
20 registry = NodeRegistry()
21 registry.register(LikeHandlerNode)
22 registry.register(UnlikeHandlerNode)
24 # Supervisor routing
25 candidates = registry.evaluate_triggers("shopping", state)
26 """
28 def __init__(self, valid_slices: set[str] | None = None):
29 """Initialize registry.
31 Args:
32 valid_slices: Valid slice names for validation.
33 Defaults to basic set if not provided.
34 """
35 self._nodes: dict[str, type] = {} # name -> node class
36 self._contracts: dict[str, NodeContract] = {} # name -> contract
37 self._valid_slices = valid_slices or {"request", "response", "_internal"}
39 def register(self, node_class: type) -> None:
40 """Register a node class.
42 Args:
43 node_class: ModularNode subclass with CONTRACT
44 """
45 if not hasattr(node_class, "CONTRACT"):
46 raise ValueError(f"Node class {node_class.__name__} must have CONTRACT")
48 contract = node_class.CONTRACT
49 self._validate_contract(contract)
51 if contract.name in self._nodes:
52 raise ValueError(f"Node {contract.name} is already registered")
54 self._nodes[contract.name] = node_class
55 self._contracts[contract.name] = contract
57 logger.info(f"Registered node: {contract.name} (supervisor={contract.supervisor})")
59 def _validate_contract(self, contract: NodeContract) -> None:
60 """Validate contract consistency."""
61 for slice_name in contract.reads:
62 if slice_name not in self._valid_slices:
63 logger.warning(f"Unknown slice in reads: {slice_name}")
65 for slice_name in contract.writes:
66 if slice_name not in self._valid_slices:
67 logger.warning(f"Unknown slice in writes: {slice_name}")
68 if slice_name == "request":
69 logger.warning("Writing to 'request' slice is discouraged")
71 def add_valid_slice(self, slice_name: str) -> None:
72 """Add a valid slice name."""
73 self._valid_slices.add(slice_name)
75 def get_node_class(self, name: str) -> type | None:
76 """Get node class by name."""
77 return self._nodes.get(name)
79 def get_contract(self, name: str) -> NodeContract | None:
80 """Get contract by name."""
81 return self._contracts.get(name)
83 def get_all_nodes(self) -> list[str]:
84 """Get all node names."""
85 return list(self._nodes.keys())
87 def get_supervisor_nodes(self, supervisor: str) -> list[str]:
88 """Get node names belonging to a supervisor."""
89 return [
90 name for name, contract in self._contracts.items()
91 if contract.supervisor == supervisor
92 ]
94 # =========================================================================
95 # Routing Evaluation
96 # =========================================================================
98 def evaluate_triggers(
99 self,
100 supervisor: str,
101 state: dict,
102 ) -> list[tuple[int, str]]:
103 """Evaluate all node trigger conditions and return matches.
105 Args:
106 supervisor: Supervisor name to evaluate
107 state: Current State
109 Returns:
110 Matched node names (ordered by priority) as (priority, name) tuples
111 """
112 candidates: list[tuple[int, str]] = []
114 for name in self.get_supervisor_nodes(supervisor):
115 contract = self._contracts[name]
117 # Find highest priority matching condition for this node
118 highest_priority: int | None = None
119 for condition in contract.trigger_conditions:
120 if self._evaluate_condition(condition, state):
121 if highest_priority is None or condition.priority > highest_priority:
122 highest_priority = condition.priority
124 if highest_priority is not None:
125 candidates.append((highest_priority, name))
127 # Sort by priority (descending)
128 candidates.sort(key=lambda x: x[0], reverse=True)
129 return candidates
131 def _evaluate_condition(
132 self,
133 condition: TriggerCondition,
134 state: dict,
135 ) -> bool:
136 """Evaluate a single trigger condition."""
137 def matches_expected(actual: Any, expected: Any) -> bool:
138 if expected is True:
139 return bool(actual) is True
140 if expected is False:
141 return bool(actual) is False
142 return actual == expected
144 # when conditions
145 if condition.when:
146 for key, expected in condition.when.items():
147 actual = self._get_state_value(state, key)
148 if not matches_expected(actual, expected):
149 return False
151 # when_not conditions
152 if condition.when_not:
153 for key, unexpected in condition.when_not.items():
154 actual = self._get_state_value(state, key)
155 if matches_expected(actual, unexpected):
156 return False
158 return True
160 def _get_state_value(self, state: dict, key: str) -> Any:
161 """Get value from State.
163 Key format: "slice.field" / "slice.nested.field" or "field"
164 """
165 if "." in key:
166 parts = key.split(".")
167 slice_name = parts[0]
168 value: Any = state.get(slice_name, {})
169 for part in parts[1:]:
170 if not isinstance(value, dict):
171 return None
172 value = value.get(part)
173 return value
174 else:
175 # Flat key: search all slices
176 preferred_order = ["request", "response", "_internal"]
177 ordered_slices = [
178 s for s in preferred_order if s in self._valid_slices
179 ] + sorted(self._valid_slices - set(preferred_order))
181 for slice_name in ordered_slices:
182 slice_data = state.get(slice_name, {})
183 if isinstance(slice_data, dict) and key in slice_data:
184 return slice_data[key]
185 return None
187 # =========================================================================
188 # LLM Prompt Generation
189 # =========================================================================
191 def build_llm_prompt(self, supervisor: str, state: dict) -> str:
192 """Generate LLM prompt for Supervisor.
194 Aggregates LLM hints from each node to build prompt.
195 """
196 lines = ["Choose the next action based on the current state:\n"]
198 for name in self.get_supervisor_nodes(supervisor):
199 contract = self._contracts[name]
200 hints = contract.get_llm_hints()
202 if hints:
203 hint_text = "; ".join(hints)
204 lines.append(f"- **{name}**: {contract.description} ({hint_text})")
205 else:
206 lines.append(f"- **{name}**: {contract.description}")
208 lines.append("\n- **done**: Complete the current flow\n")
209 lines.append("Return only the action name.")
211 return "\n".join(lines)
213 # =========================================================================
214 # Data Flow Analysis
215 # =========================================================================
217 def analyze_data_flow(self) -> dict[str, list[str]]:
218 """Analyze data flow dependencies between nodes.
220 Returns:
221 {node_name: [dependent_nodes], ...}
222 """
223 dependencies: dict[str, list[str]] = {}
225 for name, contract in self._contracts.items():
226 deps = []
227 for other_name, other_contract in self._contracts.items():
228 if other_name == name:
229 continue
230 # If another node writes to slices I read, there's a dependency
231 if set(contract.reads) & set(other_contract.writes):
232 deps.append(other_name)
233 dependencies[name] = deps
235 return dependencies
238# =============================================================================
239# Singleton
240# =============================================================================
242_registry: NodeRegistry | None = None
245def get_node_registry() -> NodeRegistry:
246 """Get global registry."""
247 global _registry
248 if _registry is None:
249 _registry = NodeRegistry()
250 return _registry
253def reset_registry() -> None:
254 """Reset registry (for testing)."""
255 global _registry
256 _registry = None