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"]