Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-agents/src/lexigram/ai/agents/strategies/supervisor.py: 18%

153 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Supervisor strategy for multi-agent orchestration. 

2 

3Implements a delegation-based reasoning loop where a supervisor LLM 

4decides which specialized sub-agent should handle each part of a task: 

5 

6 1. **ASSESS** — The supervisor LLM receives the objective and a 

7 list of available sub-agents (presented as tools). 

8 2. **DELEGATE** — The supervisor selects a sub-agent and provides 

9 a sub-task message. 

10 3. **REVIEW** — The sub-agent's response is fed back to the 

11 supervisor as an observation. 

12 4. **DECIDE** — The supervisor either delegates again, or produces 

13 a final synthesized answer. 

14 

15This strategy replaces CrewAI's hierarchical process and AutoGen's 

16conversation patterns with a clean, contract-based implementation. 

17 

18Example:: 

19 

20 from lexigram.ai.agents.strategies import SupervisorStrategy 

21 

22 supervisor_strategy = SupervisorStrategy( 

23 sub_agents={"billing": billing_agent, "technical": tech_agent}, 

24 executor=executor, 

25 max_delegations=5, 

26 ) 

27""" 

28 

29from __future__ import annotations 

30 

31import asyncio 

32import time 

33from typing import TYPE_CHECKING, Any, cast 

34 

35from lexigram.ai.agents.delegation.agent_tool import AgentAsToolAdapter 

36from lexigram.ai.agents.strategies.base import AbstractStrategy 

37from lexigram.ai.agents.strategies.token_utils import TokenAccumulator 

38from lexigram.ai.agents.types import ReasoningStep, ToolExecutionRecord 

39from lexigram.contracts.ai.agents import AgentError, AgentResponse 

40from lexigram.contracts.ai.llm import ChatMessage, Role 

41from lexigram.logging import ( 

42 get_logger, 

43) 

44from lexigram.result import Err, Ok, Result 

45 

46if TYPE_CHECKING: 

47 from lexigram.contracts.ai.agents import ( 

48 AgentExecutorProtocol, 

49 AgentProtocol, 

50 ToolProtocol, 

51 ) 

52 from lexigram.contracts.ai.llm import CompletionProtocol, LLMClientProtocol 

53 

54logger = get_logger(__name__) 

55 

56 

57# --------------------------------------------------------------------------- 

58# Prompt templates 

59# --------------------------------------------------------------------------- 

60 

61_SUPERVISOR_SYSTEM = """ 

62You are a supervisor agent that delegates tasks to specialized sub-agents. 

63 

64## Your Role 

65- Analyze the user's request and decide which sub-agent is best suited. 

66- Delegate tasks by calling the appropriate agent tool. 

67- Review sub-agent responses and decide: delegate again or provide a final answer. 

68- You may delegate to multiple agents sequentially if the task requires it. 

69 

70## Output Format 

71 

72### Option A — Delegate to a sub-agent 

73``` 

74THOUGHT: <your reasoning about which agent to use and why> 

75ACTION: delegate_to_<agent_name> 

76ACTION_INPUT: {{"message": "<the sub-task to delegate>"}} 

77``` 

78 

79### Option B — Final Answer (after reviewing sub-agent responses) 

80``` 

81THOUGHT: <your reasoning about why you have enough information> 

82FINAL_ANSWER: <your complete, synthesized answer> 

83``` 

84 

85## Rules 

86- Always start with THOUGHT. 

87- Delegate ONE task at a time. 

88- After receiving an observation, decide whether to delegate again or answer. 

89- Synthesize information from multiple agents if needed. 

90- When you have a complete answer, use FINAL_ANSWER. 

91 

92## Available Agents 

93{agent_descriptions} 

94""" 

95 

96_OBSERVATION_TEMPLATE = "OBSERVATION from {agent_name}: {observation}" 

97 

98 

99class SupervisorStrategy(AbstractStrategy): 

100 """Supervisor strategy that orchestrates multiple sub-agents. 

101 

102 The supervisor uses an LLM to decide which sub-agent should handle 

103 each part of a task. Sub-agents are exposed as tools using 

104 ``AgentAsToolAdapter``, allowing the supervisor to delegate via 

105 the standard tool-calling mechanism. 

106 

107 This enables hierarchical multi-agent patterns: 

108 - Customer support routing (classify → specialist agent) 

109 - Research tasks (delegate to searcher, analyzer, summarizer) 

110 - Quality assurance (generate → review → revise) 

111 

112 Example:: 

113 

114 strategy = SupervisorStrategy( 

115 sub_agents={"research": research_agent, "writing": writing_agent}, 

116 executor=executor, 

117 ) 

118 result = await strategy.execute( 

119 message="Write a report about quantum computing", 

120 tools=[], # supervisor uses sub-agents, not tools 

121 history=[], 

122 llm=llm_client, 

123 ) 

124 """ 

125 

126 def __init__( 

127 self, 

128 sub_agents: dict[str, AgentProtocol], 

129 executor: AgentExecutorProtocol, 

130 *, 

131 max_delegations: int = 5, 

132 llm_timeout: float = 120.0, 

133 ) -> None: 

134 """Initialize the supervisor strategy. 

135 

136 Args: 

137 sub_agents: Named sub-agents available for delegation. 

138 executor: Agent executor used to run sub-agents. 

139 max_delegations: Maximum number of delegation rounds. 

140 llm_timeout: Timeout per LLM call in seconds. 

141 """ 

142 self._sub_agents = sub_agents 

143 self._executor = executor 

144 self.max_delegations = max_delegations 

145 self.llm_timeout = llm_timeout 

146 

147 # Build agent tools 

148 self._agent_tools: dict[str, AgentAsToolAdapter] = { 

149 name: AgentAsToolAdapter(agent=agent, executor=executor) 

150 for name, agent in sub_agents.items() 

151 } 

152 

153 async def execute( 

154 self, 

155 message: str, 

156 tools: list[ToolProtocol], 

157 history: list[dict[str, Any]], 

158 llm: LLMClientProtocol, 

159 **kwargs: Any, 

160 ) -> Result[AgentResponse, AgentError]: 

161 """Execute the supervisor delegation loop. 

162 

163 Args: 

164 message: The user's input message. 

165 tools: Additional tools (merged with agent-as-tool adapters). 

166 history: Conversation history as list of message dicts. 

167 llm: LLM client implementing ``LLMClientProtocol``. 

168 **kwargs: Additional parameters (system_prompt, etc.). 

169 

170 Returns: 

171 ``Ok(AgentResponse)`` with the supervisor's final answer 

172 and full delegation trace. ``Err(AgentError)`` on failure. 

173 """ 

174 system_prompt: str = kwargs.get("system_prompt", "") 

175 guard_pipeline = kwargs.get("guard_pipeline") 

176 steps: list[ReasoningStep] = [] 

177 tool_calls: list[ToolExecutionRecord] = [] 

178 usage = TokenAccumulator() 

179 start_time = time.monotonic() 

180 

181 # Build combined tool map: agent tools + any extra tools 

182 all_tools: dict[str, ToolProtocol] = {} 

183 for adapter in self._agent_tools.values(): 

184 all_tools[adapter.name] = adapter 

185 for t in tools: 

186 all_tools[t.name] = t 

187 

188 # Build agent descriptions for system prompt 

189 agent_descriptions = "\n".join( 

190 f"- **{adapter.name}**: {adapter.description}" 

191 for adapter in self._agent_tools.values() 

192 ) 

193 

194 full_system = system_prompt + _SUPERVISOR_SYSTEM.format( 

195 agent_descriptions=agent_descriptions 

196 ) 

197 

198 messages = self._build_messages(message, history, full_system) 

199 

200 for delegation in range(1, self.max_delegations + 1): 

201 # Ask supervisor LLM what to do 

202 completion = await self._call_llm(llm, messages) 

203 if completion is None: 

204 return Err( 

205 AgentError( 

206 f"Supervisor LLM returned empty response at delegation {delegation}" 

207 ) 

208 ) 

209 usage.add(completion) 

210 llm_text = ( 

211 completion.content 

212 if hasattr(completion, "content") 

213 else str(completion) 

214 ) 

215 

216 thought = self._extract_thought(llm_text) 

217 final_answer = self._extract_final_answer(llm_text) 

218 

219 # Check for final answer 

220 if final_answer is not None: 

221 steps.append( 

222 ReasoningStep( 

223 step_number=delegation, 

224 thought=thought, 

225 action="final_answer", 

226 observation=final_answer, 

227 ) 

228 ) 

229 

230 elapsed = (time.monotonic() - start_time) * 1000 

231 logger.info( 

232 "supervisor_final_answer", 

233 delegations=delegation, 

234 steps=len(steps), 

235 ) 

236 return Ok( 

237 AgentResponse( 

238 message=final_answer, 

239 steps=steps, 

240 tool_calls=tool_calls, 

241 total_tokens=usage.total_tokens, 

242 prompt_tokens=usage.prompt_tokens, 

243 completion_tokens=usage.completion_tokens, 

244 duration_ms=elapsed, 

245 metadata={ 

246 "strategy": "supervisor", 

247 "delegations": delegation, 

248 "agents_used": list( 

249 { 

250 s.action 

251 for s in steps 

252 if s.action and s.action.startswith("delegate_to_") 

253 } 

254 ), 

255 }, 

256 ) 

257 ) 

258 

259 # Parse delegation (tool call) 

260 tool_name, tool_args = self._extract_tool_call(llm_text) 

261 

262 if tool_name is None or tool_name not in all_tools: 

263 # No valid delegation — nudge the LLM 

264 steps.append( 

265 ReasoningStep( 

266 step_number=delegation, 

267 thought=thought, 

268 action=None, 

269 observation="[No valid delegation or final answer detected]", 

270 ) 

271 ) 

272 messages.append(ChatMessage(role=Role.ASSISTANT, content=llm_text)) 

273 messages.append( 

274 ChatMessage( 

275 role=Role.USER, 

276 content=( 

277 "Your response did not contain a valid delegation or " 

278 "FINAL_ANSWER. Please:\n" 

279 "- Use ACTION: delegate_to_<agent_name> to delegate, or\n" 

280 "- Use FINAL_ANSWER: to provide your answer." 

281 ), 

282 ) 

283 ) 

284 continue 

285 

286 # Execute delegation 

287 logger.info( 

288 "supervisor_delegation", 

289 delegation=delegation, 

290 target=tool_name, 

291 ) 

292 

293 tool = all_tools[tool_name] 

294 delegation_start = time.monotonic() 

295 

296 try: 

297 result = await tool.execute(**tool_args) 

298 duration = (time.monotonic() - delegation_start) * 1000 

299 observation = str(result) 

300 

301 record = ToolExecutionRecord( 

302 tool_name=tool_name, 

303 arguments=tool_args, 

304 result=result, 

305 duration_ms=duration, 

306 ) 

307 except (RuntimeError, TypeError, ValueError, OSError, LookupError) as exc: 

308 duration = (time.monotonic() - delegation_start) * 1000 

309 observation = f"Delegation failed: {exc}" 

310 record = ToolExecutionRecord( 

311 tool_name=tool_name, 

312 arguments=tool_args, 

313 error=str(exc), 

314 duration_ms=duration, 

315 ) 

316 

317 tool_calls.append(record) 

318 

319 # Guard before feeding back so the raw observation never hits context 

320 from lexigram.ai.agents.strategies.guard_hook import guard_observation 

321 

322 observation = await guard_observation( 

323 guard_pipeline, observation, tool_name=tool_name 

324 ) 

325 

326 # Derive agent name from tool name for the observation 

327 agent_display = ( 

328 tool_name.replace("delegate_to_", "") 

329 if tool_name.startswith("delegate_to_") 

330 else tool_name 

331 ) 

332 

333 steps.append( 

334 ReasoningStep( 

335 step_number=delegation, 

336 thought=thought, 

337 action=tool_name, 

338 tool_call=record, 

339 observation=observation, 

340 ) 

341 ) 

342 

343 # Feed observation back 

344 messages.append(ChatMessage(role=Role.ASSISTANT, content=llm_text)) 

345 messages.append( 

346 ChatMessage( 

347 role=Role.USER, 

348 content=_OBSERVATION_TEMPLATE.format( 

349 agent_name=agent_display, 

350 observation=observation, 

351 ), 

352 ) 

353 ) 

354 

355 # Max delegations reached 

356 elapsed = (time.monotonic() - start_time) * 1000 

357 logger.warning( 

358 "supervisor_max_delegations", 

359 max_delegations=self.max_delegations, 

360 steps=len(steps), 

361 ) 

362 last_obs = steps[-1].observation if steps else "No response generated" 

363 return Ok( 

364 AgentResponse( 

365 message=f"[Max delegations reached] {last_obs}", 

366 steps=steps, 

367 tool_calls=tool_calls, 

368 total_tokens=usage.total_tokens, 

369 prompt_tokens=usage.prompt_tokens, 

370 completion_tokens=usage.completion_tokens, 

371 duration_ms=elapsed, 

372 metadata={ 

373 "strategy": "supervisor", 

374 "delegations": self.max_delegations, 

375 "max_delegations_reached": True, 

376 }, 

377 ) 

378 ) 

379 

380 # ------------------------------------------------------------------ 

381 # LLM Interaction 

382 # ------------------------------------------------------------------ 

383 

384 async def _call_llm( 

385 self, 

386 llm: LLMClientProtocol, 

387 messages: list[ChatMessage], 

388 ) -> CompletionProtocol | None: 

389 """Call the LLM and return the completion, or ``None`` on failure.""" 

390 try: 

391 result = await asyncio.wait_for( 

392 llm.complete(cast("list[Any]", messages)), 

393 timeout=self.llm_timeout, 

394 ) 

395 except TimeoutError: 

396 logger.warning("supervisor_llm_timeout", timeout=self.llm_timeout) 

397 return None 

398 except (OSError, ConnectionError, RuntimeError, ValueError) as exc: 

399 logger.warning("supervisor_llm_error", error=str(exc)) 

400 return None 

401 

402 if not result.is_ok(): 

403 logger.warning("supervisor_llm_err", error=str(result.unwrap_err())) 

404 return None 

405 

406 return result.unwrap() 

407 

408 # ------------------------------------------------------------------ 

409 # Parsing Helpers 

410 # ------------------------------------------------------------------ 

411 

412 @staticmethod 

413 def _extract_thought(text: str) -> str: 

414 """Extract THOUGHT section from LLM response.""" 

415 for line in text.split("\n"): 

416 stripped = line.strip() 

417 if stripped.upper().startswith("THOUGHT:"): 

418 return stripped[len("THOUGHT:") :].strip() 

419 return text.split("\n", maxsplit=1)[0][:200] 

420 

421 @staticmethod 

422 def _extract_final_answer(text: str) -> str | None: 

423 """Extract FINAL_ANSWER if present.""" 

424 marker = "FINAL_ANSWER:" 

425 upper = text.upper() 

426 idx = upper.find(marker) 

427 if idx == -1: 

428 return None 

429 return text[idx + len(marker) :].strip() 

430 

431 @staticmethod 

432 def _extract_tool_call(text: str) -> tuple[str | None, dict[str, Any]]: 

433 """Extract ACTION and ACTION_INPUT from LLM response.""" 

434 from lexigram.serialization.backends.json import JSONDecodeError, loads 

435 

436 action_name: str | None = None 

437 action_input: dict[str, Any] = {} 

438 

439 for line in text.split("\n"): 

440 stripped = line.strip() 

441 upper = stripped.upper() 

442 if upper.startswith("ACTION:") and not upper.startswith("ACTION_INPUT:"): 

443 action_name = stripped[len("ACTION:") :].strip() 

444 elif upper.startswith("ACTION_INPUT:"): 

445 raw = stripped[len("ACTION_INPUT:") :].strip() 

446 try: 

447 action_input = loads(raw) 

448 except (ValueError, JSONDecodeError): 

449 brace_start = raw.find("{") 

450 if brace_start != -1: 

451 depth = 0 

452 for i, ch in enumerate(raw[brace_start:]): 

453 if ch == "{": 

454 depth += 1 

455 elif ch == "}": 

456 depth -= 1 

457 if depth == 0: 

458 try: 

459 action_input = loads( 

460 raw[brace_start : brace_start + i + 1] 

461 ) 

462 except (ValueError, JSONDecodeError): 

463 pass 

464 break 

465 

466 return action_name, action_input 

467 

468 # ------------------------------------------------------------------ 

469 # Message Building 

470 # ------------------------------------------------------------------ 

471 

472 @staticmethod 

473 def _build_messages( 

474 message: str, 

475 history: list[dict[str, Any]], 

476 system_prompt: str, 

477 ) -> list[ChatMessage]: 

478 """Convert history + new message into ChatMessage objects.""" 

479 messages: list[ChatMessage] = [] 

480 if system_prompt: 

481 messages.append(ChatMessage(role=Role.SYSTEM, content=system_prompt)) 

482 for entry in history: 

483 role_str = entry.get("role", "user") 

484 content = entry.get("content", "") 

485 try: 

486 role = Role(role_str) 

487 except ValueError: 

488 role = Role.USER 

489 messages.append(ChatMessage(role=role, content=content)) 

490 messages.append(ChatMessage(role=Role.USER, content=message)) 

491 return messages 

492 

493 

494__all__ = ["SupervisorStrategy"]