# -*- coding: utf-8 -*- """TaskDelegator - Subagent spawning and task delegation. Provides delegate() and delegate_parallel() for spawning subagents with separate context and memory. Supports runtime dynamic subagent definition via task_data with description, prompt, and tools. """ from __future__ import annotations import asyncio import logging import uuid from typing import Any, Awaitable, Callable, Dict, List, Optional, Union from agentscope.message import Msg logger = logging.getLogger(__name__) # Type alias for subagent specification SubagentSpec = Dict[str, Any] """Subagent specification format: { "description": "Expert code reviewer...", "prompt": "Analyze code quality...", "tools": ["Read", "Glob", "Grep"], # Optional: list of tool names "model": "gpt-4o", # Optional: model name } """ class TaskDelegator: """Delegates tasks to subagents with isolated context. Supports: - delegate(): Spawn single subagent for task - delegate_parallel(): Spawn multiple subagents concurrently - delegate_task(): Delegate with dynamic subagent definition from task_data Each subagent gets its own memory/context to prevent cross-contamination. Dynamic Subagent Definition: task_data can include an "agents" dict to define subagents inline: task_data = { "task": "Review the code changes", "agents": { "code-reviewer": { "description": "Expert code reviewer for quality and security.", "prompt": "Analyze code quality and suggest improvements.", "tools": ["Read", "Glob", "Grep"], } } } """ def __init__(self, messenger: Any, registry: Any): """Initialize TaskDelegator. Args: messenger: AgentMessenger for communication registry: AgentRegistry for agent lookup """ self._messenger = messenger self._registry = registry self._subagents: Dict[str, Any] = {} self._dynamic_subagents: Dict[str, SubagentSpec] = {} self._tasks: Dict[str, asyncio.Task] = {} async def delegate( self, agent_id: str, task: Callable[..., Awaitable[Msg]], context: Optional[Dict[str, Any]] = None, ) -> asyncio.Task: """Delegate task to a single subagent. Args: agent_id: Unique identifier for this subagent instance task: Async function representing the task context: Optional context dict for the subagent Returns: asyncio.Task for the delegated task """ async def _run_with_context(): result = await task(context or {}) return result self._tasks[agent_id] = asyncio.create_task(_run_with_context()) logger.info("Delegated task to subagent: %s", agent_id) return self._tasks[agent_id] async def delegate_parallel( self, tasks: List[Dict[str, Any]], ) -> List[asyncio.Task]: """Delegate multiple tasks in parallel. Args: tasks: List of task dicts with keys: - agent_id: Unique identifier - task: Async function to execute - context: Optional context dict Returns: List of asyncio.Task for all delegated tasks """ async def _run_task(task_def: Dict[str, Any]): agent_id = task_def["agent_id"] task_func = task_def["task"] context = task_def.get("context", {}) async def _run_with_context(): return await task_func(context) self._tasks[agent_id] = asyncio.create_task(_run_with_context()) return self._tasks[agent_id] gathered_tasks = await asyncio.gather( *[_run_task(t) for t in tasks], return_exceptions=True, ) valid_tasks = [t for t in gathered_tasks if isinstance(t, asyncio.Task)] logger.info( "Delegated %d tasks in parallel (%d succeeded)", len(tasks), len(valid_tasks), ) return valid_tasks async def wait_for(self, agent_id: str, timeout: Optional[float] = None) -> Any: """Wait for subagent task to complete. Args: agent_id: Subagent identifier timeout: Optional timeout in seconds Returns: Task result Raises: asyncio.TimeoutError: If task doesn't complete in time KeyError: If agent_id not found """ if agent_id not in self._tasks: raise KeyError(f"Unknown subagent: {agent_id}") try: return await asyncio.wait_for( self._tasks[agent_id], timeout=timeout, ) except asyncio.TimeoutError: logger.warning("Task %s timed out after %s seconds", agent_id, timeout) raise async def cancel(self, agent_id: str) -> bool: """Cancel a subagent task. Args: agent_id: Subagent identifier Returns: True if task was cancelled """ if agent_id in self._tasks: self._tasks[agent_id].cancel() del self._tasks[agent_id] logger.info("Cancelled subagent task: %s", agent_id) return True return False def list_tasks(self) -> List[str]: """List active subagent task IDs. Returns: List of agent_ids with pending tasks """ return list(self._tasks.keys()) @property def tasks(self) -> Dict[str, asyncio.Task]: """Get copy of active tasks dict.""" return dict(self._tasks) def delegate_task( self, task_type: str, task_data: Dict[str, Any], target_agent: Optional[str] = None, ) -> Dict[str, Any]: """Delegate a task with optional dynamic subagent definition. Supports runtime subagent definition via task_data["agents"]: task_data = { "task": "Review code changes", "agents": { "code-reviewer": { "description": "Expert code reviewer...", "prompt": "Analyze code quality...", "tools": ["Read", "Glob", "Grep"], } } } Args: task_type: Type of task (e.g., "analysis", "review", "research") task_data: Task payload, may include "agents" for dynamic subagent def target_agent: Optional specific agent ID to delegate to Returns: Dict with "success" and result/error """ try: # Extract dynamic subagent definitions from task_data agents_def = task_data.get("agents", {}) if agents_def: # Register dynamic subagents for agent_name, agent_spec in agents_def.items(): self._dynamic_subagents[agent_name] = agent_spec logger.info( "Registered dynamic subagent: %s (description: %s)", agent_name, agent_spec.get("description", "")[:50], ) # Determine target agent effective_target = target_agent if not effective_target: # Use first available dynamic subagent or default if agents_def: effective_target = next(iter(agents_def.keys())) else: effective_target = "default" # Execute the task task_result = self._execute_task( task_type=task_type, task_data=task_data, target_agent=effective_target, ) # Clean up dynamic subagents after execution for agent_name in agents_def.keys(): self._dynamic_subagents.pop(agent_name, None) return { "success": True, "result": task_result, "subagents_used": list(agents_def.keys()) if agents_def else [], } except Exception as e: logger.error("Task delegation failed: %s", e) return { "success": False, "error": str(e), } def _execute_task( self, task_type: str, task_data: Dict[str, Any], target_agent: str, ) -> Any: """Execute the delegated task. Args: task_type: Type of task task_data: Task payload target_agent: Target agent identifier Returns: Task execution result """ task_content = task_data.get("task", task_data.get("prompt", "")) # Check if we have a dynamic subagent spec for this target agent_spec = self._dynamic_subagents.get(target_agent) if agent_spec: logger.info( "Executing task '%s' with dynamic subagent '%s' (prompt: %s)", task_type, target_agent, agent_spec.get("prompt", "")[:50], ) # In a full implementation, this would create and run an actual agent # For now, return a structured result indicating the task was received return { "task_type": task_type, "task": task_content, "subagent": { "name": target_agent, "description": agent_spec.get("description", ""), "prompt": agent_spec.get("prompt", ""), "tools": agent_spec.get("tools", []), }, "status": "completed", "message": f"Task '{task_type}' executed with dynamic subagent '{target_agent}'", } # Fallback: execute with default behavior logger.info( "Executing task '%s' with default agent '%s'", task_type, target_agent, ) return { "task_type": task_type, "task": task_content, "target_agent": target_agent, "status": "completed", "message": f"Task '{task_type}' executed with agent '{target_agent}'", } def get_dynamic_subagent(self, name: str) -> Optional[SubagentSpec]: """Get a dynamically defined subagent specification. Args: name: Subagent name Returns: Subagent spec dict or None if not found """ return self._dynamic_subagents.get(name) def list_dynamic_subagents(self) -> List[str]: """List all registered dynamic subagent names. Returns: List of subagent names """ return list(self._dynamic_subagents.keys()) __all__ = ["TaskDelegator", "SubagentSpec"]