Coverage for src / agent_contracts / supervisor.py: 86%
132 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-09 00:42 +0900
1"""GenericSupervisor - Generic Supervisor.
3Has no node-specific routing logic, determines routing via
4Registry trigger conditions and LLM.
5"""
6from __future__ import annotations
8from typing import Optional
10from pydantic import BaseModel, Field
12from langchain_core.language_models import BaseChatModel
13from langchain_core.runnables import RunnableConfig
15from agent_contracts.registry import NodeRegistry, get_node_registry
16from agent_contracts.config import get_config
17from agent_contracts.utils.logging import get_logger
18from agent_contracts.routing import MatchedRule, RoutingReason, RoutingDecision
20logger = get_logger("agent_contracts.supervisor")
23class SupervisorDecision(BaseModel):
24 """Supervisor decision result."""
25 next_node: str = Field(description="Next node name, or 'done'")
26 reasoning: str = Field(default="", description="Decision reasoning")
29class GenericSupervisor:
30 """Generic Supervisor.
32 Has no node-specific rule-based logic,
33 evaluates conditions from NodeRegistry.
35 Example:
36 supervisor = GenericSupervisor("shopping", llm=llm)
37 decision = await supervisor.decide(state)
38 """
40 def __init__(
41 self,
42 supervisor_name: str,
43 llm: BaseChatModel | None = None,
44 registry: NodeRegistry | None = None,
45 max_iterations: int | None = None,
46 terminal_response_types: set[str] | None = None,
47 ):
48 """Initialize.
50 Args:
51 supervisor_name: Supervisor type
52 llm: LangChain LLM
53 registry: Node registry (uses global if omitted)
54 max_iterations: Max iterations (uses config if omitted)
55 terminal_response_types: Terminal response types (uses config if omitted)
56 """
57 self.name = supervisor_name
58 self.llm = llm
59 self.registry = registry or get_node_registry()
60 self.logger = logger
62 # Load from config or use defaults
63 config = get_config()
64 self.max_iterations = max_iterations or config.supervisor.max_iterations
65 self.terminal_response_types = terminal_response_types or set(
66 config.supervisor.terminal_response_types
67 )
69 async def run(
70 self,
71 state: dict,
72 config: Optional[RunnableConfig] = None,
73 ) -> dict:
74 """Execute Supervisor node.
76 Called as LangGraph node.
77 """
78 # Iteration management
79 internal = state.get("_internal", {})
80 iteration_key = f"{self.name}_iteration"
81 current_iteration = internal.get(iteration_key, 0)
83 # Infinite loop prevention
84 if current_iteration >= self.max_iterations:
85 self.logger.warning(f"Max iterations ({self.max_iterations}) reached for {self.name}")
86 return {
87 "_internal": {
88 **internal,
89 "decision": "done",
90 iteration_key: current_iteration,
91 }
92 }
94 # Decide
95 decision = await self.decide(state, config)
97 self.logger.info(
98 f"{self.name} supervisor decision: {decision.next_node} ({decision.reasoning})"
99 )
101 return {
102 "_internal": {
103 **internal,
104 "decision": decision.next_node,
105 iteration_key: current_iteration + 1,
106 }
107 }
109 async def decide(
110 self,
111 state: dict,
112 config: Optional[RunnableConfig] = None,
113 ) -> SupervisorDecision:
114 """Determine next node.
116 This is a convenience wrapper around decide_with_trace() that returns
117 a simplified SupervisorDecision instead of the full RoutingDecision.
119 For debugging and detailed routing information, use decide_with_trace().
120 """
121 routing_decision = await self.decide_with_trace(state, config)
122 return routing_decision.to_supervisor_decision()
124 def _check_immediate_rules(self, state: dict) -> str | None:
125 """Check if should exit immediately.
127 Returns 'done' for user input waiting or final states.
128 """
129 response = state.get("response", {})
130 response_type = response.get("response_type")
132 if response_type in self.terminal_response_types:
133 return "done"
135 return None
137 def _select_top_matches(self, matches: list[tuple[int, str]]) -> list[str]:
138 """Select top candidates handling ties (Top 3 + Ties)."""
139 if not matches:
140 return []
142 selected = []
143 limit = 3
144 last_prio = -1
146 for i, (prio, name) in enumerate(matches):
147 if i < limit:
148 selected.append(name)
149 last_prio = prio
150 elif prio == last_prio:
151 selected.append(name)
152 else:
153 break
154 return selected
156 async def _decide_with_llm(
157 self,
158 state: dict,
159 rule_candidates: list[str],
160 child_decision: str | None,
161 config: Optional[RunnableConfig] = None,
162 ) -> SupervisorDecision | None:
163 """Decide using LLM."""
164 try:
165 prompt = self.registry.build_llm_prompt(self.name, state)
167 # Add current state info
168 request = state.get("request", {})
170 context = f"""
171Current action: {request.get('action', 'unknown')}
172User message: {request.get('message', 'None')}
174High priority system rules suggest: {rule_candidates}
175Last active node suggested: {child_decision or 'None'}
176"""
177 full_prompt = f"{prompt}\n\nContext:\n{context}"
179 # Use LangChain structured output
180 structured_llm = self.llm.with_structured_output(SupervisorDecision)
181 result = await structured_llm.ainvoke(
182 f"System: You are a decision-making supervisor for a {self.name} flow. "
183 f"If 'High priority system rules' are provided, you MUST select one of them. "
184 f"Otherwise, prioritize user intent.\n\n{full_prompt}",
185 config=config,
186 )
188 # Validate LLM decision against valid nodes
189 valid_nodes = set(self.registry.get_supervisor_nodes(self.name))
190 valid_nodes.add("done")
192 if result.next_node not in valid_nodes:
193 self.logger.warning(
194 f"LLM returned invalid node: {result.next_node}, "
195 f"valid nodes: {valid_nodes}"
196 )
197 # If rule candidates exist, use the top one
198 if rule_candidates:
199 return SupervisorDecision(
200 next_node=rule_candidates[0],
201 reasoning=f"LLM returned invalid '{result.next_node}', using rule candidate"
202 )
203 # Otherwise return None to trigger fallback
204 return None
206 return result
208 except Exception as e:
209 self.logger.error(f"LLM decision failed: {e}")
210 return None
212 def _build_matched_rules(
213 self,
214 matches: list[tuple[int, str]],
215 ) -> list[MatchedRule]:
216 """Build MatchedRule list from trigger matches."""
217 matched_rules = []
219 for priority, node_name in matches:
220 contract = self.registry.get_contract(node_name)
221 if not contract:
222 continue
224 # Find the matching condition description
225 condition_str = ""
226 for condition in contract.trigger_conditions:
227 if condition.priority == priority:
228 if condition.when:
229 parts = [f"{k}={v}" for k, v in condition.when.items()]
230 condition_str = " AND ".join(parts)
231 elif condition.when_not:
232 parts = [f"NOT {k}={v}" for k, v in condition.when_not.items()]
233 condition_str = " AND ".join(parts)
234 else:
235 condition_str = "(always)"
236 break
238 matched_rules.append(MatchedRule(
239 node=node_name,
240 condition=condition_str or "(unknown)",
241 priority=priority,
242 ))
244 return matched_rules
246 async def decide_with_trace(
247 self,
248 state: dict,
249 config: Optional[RunnableConfig] = None,
250 ) -> RoutingDecision:
251 """Determine next node with full traceability.
253 Returns RoutingDecision with detailed reasoning.
254 Use this for debugging and explainability.
256 Example:
257 decision = await supervisor.decide_with_trace(state)
258 print(f"Selected: {decision.selected_node}")
259 print(f"Type: {decision.reason.decision_type}")
260 for rule in decision.reason.matched_rules:
261 print(f" - {rule.node} (P{rule.priority}): {rule.condition}")
262 """
263 # Enhance trace config (create new config to avoid mutation)
264 base_config = config or {}
265 existing_metadata = base_config.get("metadata", {})
266 existing_tags = base_config.get("tags", [])
267 config = {
268 **base_config,
269 "metadata": {
270 **existing_metadata,
271 "supervisor_name": self.name,
272 "supervisor_iteration": state.get("_internal", {}).get(f"{self.name}_iteration", 0),
273 },
274 "tags": [*existing_tags, "supervisor_decision"],
275 }
277 # Phase 0: Immediate exit check (terminal state)
278 immediate = self._check_immediate_rules(state)
279 if immediate:
280 return RoutingDecision(
281 selected_node=immediate,
282 reason=RoutingReason(decision_type="terminal_state")
283 )
285 # Phase 0.5: Explicit Routing (Return to Sender)
286 req = state.get("request", {})
287 action = req.get("action") if isinstance(req, dict) else None
289 if action == "answer":
290 interview = state.get("interview", {})
291 lq = interview.get("last_question") if isinstance(interview, dict) else None
292 node_id = None
293 if lq:
294 if isinstance(lq, dict):
295 node_id = lq.get("node_id")
296 else:
297 node_id = getattr(lq, "node_id", None)
299 if node_id:
300 return RoutingDecision(
301 selected_node=node_id,
302 reason=RoutingReason(decision_type="explicit_routing")
303 )
305 # Phase 1: Rule-based evaluation
306 matches = self.registry.evaluate_triggers(self.name, state)
307 matched_rules = self._build_matched_rules(matches)
309 # Smart selection for LLM context (Top 3 + Ties)
310 rule_candidates = self._select_top_matches(matches)
312 # Child node suggestion
313 internal = state.get("_internal", {})
314 previous_decision = internal.get("decision")
316 child_decision = None
317 if previous_decision and previous_decision != "done":
318 child_decision = previous_decision
320 # Phase 2: LLM decision
321 if self.llm:
322 llm_result = await self._decide_with_llm(
323 state,
324 rule_candidates,
325 child_decision,
326 config=config,
327 )
328 if llm_result:
329 return RoutingDecision(
330 selected_node=llm_result.next_node,
331 reason=RoutingReason(
332 decision_type="llm_decision",
333 matched_rules=matched_rules,
334 llm_used=True,
335 llm_reasoning=llm_result.reasoning,
336 )
337 )
339 # Phase 3: Fallback
340 if matches:
341 return RoutingDecision(
342 selected_node=matches[0][1],
343 reason=RoutingReason(
344 decision_type="rule_match",
345 matched_rules=matched_rules,
346 )
347 )
349 if child_decision:
350 return RoutingDecision(
351 selected_node=child_decision,
352 reason=RoutingReason(decision_type="fallback")
353 )
355 return RoutingDecision(
356 selected_node="done",
357 reason=RoutingReason(decision_type="fallback")
358 )
360 async def __call__(
361 self,
362 state: dict,
363 config: Optional[RunnableConfig] = None,
364 ) -> dict:
365 """LangGraph-compatible Callable."""
366 return await self.run(state, config)