Delegation MCP

"""Task-to-agent mapping with intelligent suggestions. This module provides functionality for mapping task categories to agents with intelligent suggestions based on agent capabilities and routing strategies. """ import logging from typing import TypedDict from rich.console import Console from rich.prompt import Confirm, Prompt from rich.table import Table from .agent_profiles import ( get_agent_profile, ROUTING_PRESETS, DEFAULT_ROUTING_RULES, RoutingPreset ) logger = logging.getLogger(__name__) console = Console() class TaskCategory(TypedDict): """Definition of a task category.""" key: str name: str description: str pattern_examples: list[str] # Task categories for delegation TASK_CATEGORIES: list[TaskCategory] = [ { "key": "security_audit", "name": "Security Audit", "description": "Security audits, vulnerability scans, safety checks", "pattern_examples": [ "security", "vulnerability", "audit", "CVE", "harden", "secure", "protect", "lock down", "access control", "permissions", "rules", "firestore rules", "authentication", "authorization", "encrypt", "expose", "leak", "breach", "attack", "threat", "OWASP", "XSS", "injection", "sanitize", "exploit", ], }, { "key": "code_review", "name": "Code Review", "description": "Code quality review, best practices analysis", "pattern_examples": [ "review", "code quality", "best practices", "lint", "improve", "clean up", "tech debt", "smell", "anti-pattern", "convention", "standards", "maintainability", "readability", "code analysis", ], }, { "key": "architecture", "name": "Architecture", "description": "System design, architecture planning, complex reasoning", "pattern_examples": [ "architecture", "design", "system design", "structure", "organize", "plan", "approach", "strategy", "pattern", "blueprint", "diagram", "flow", "schema", ], }, { "key": "refactoring", "name": "Refactoring", "description": "Code refactoring, cleanup, optimization", "pattern_examples": [ "refactor", "cleanup", "optimize code", "rename", "restructure", "reorganize", "simplify", "DRY", "extract", "inline", "consolidate", "modularize", ], }, { "key": "quick_fix", "name": "Quick Fixes", "description": "Rapid bug fixes, small code changes", "pattern_examples": [ "fix", "bug", "quick change", "error", "crash", "broken", "not working", "issue", "problem", "patch", "hotfix", ], }, { "key": "documentation", "name": "Documentation", "description": "README files, API docs, code comments", "pattern_examples": [ "documentation", "docs", "README", "comments", "comment", "explain", "describe", "guide", "tutorial", "how-to", "API docs", "docstring", "examples", ], }, { "key": "testing", "name": "Testing", "description": "Unit tests, integration tests, test coverage", "pattern_examples": [ "test", "testing", "coverage", "unit test", "integration test", "e2e", "spec", "assertion", "mock", "stub", "test case", "test suite", ], }, { "key": "performance", "name": "Performance", "description": "Performance analysis and optimization", "pattern_examples": [ "performance", "optimize", "speed", "slow", "latency", "throughput", "bottleneck", "profiling", "benchmark", "memory", "CPU", "scalability", ], }, { "key": "browser_interaction", "name": "Browser Interaction", "description": "Browser automation, web scraping, UI testing", "pattern_examples": ["browser", "selenium", "playwright", "chrome"], }, { "key": "git_operations", "name": "Git Operations", "description": "Git workflows, repository management", "pattern_examples": [ "git", "commit", "merge", "branch", "push", "pull", "rebase", "cherry-pick", "stash", "tag", "history", "checkout", "reset", "revert", ], }, { "key": "shell_tasks", "name": "Shell/Terminal", "description": "Shell scripting, terminal commands", "pattern_examples": [ "shell", "terminal", "bash", "script", "command", "CLI", "automation", "cron", "env", "environment variables", "path", "execute", ], }, { "key": "exploration", "name": "Code Exploration", "description": "Code exploration, dependency tracing, implementation analysis", "pattern_examples": ["how does", "trace the flow", "what files implement", "understand the implementation", "map dependencies"], }, { "key": "debugging", "name": "Debugging", "description": "Bug investigation, error analysis, root cause identification", "pattern_examples": ["debug", "why is failing", "investigate", "find the cause", "troubleshoot"], }, { "key": "impact_analysis", "name": "Impact Analysis", "description": "Dependency analysis, usage finding, breaking change assessment", "pattern_examples": ["what would break", "find all usages", "what depends on", "impact of changing", "all references"], }, { "key": "general", "name": "General Tasks", "description": "Default for tasks that don't fit specific categories", "pattern_examples": ["general", "misc", "other"], }, ] class TaskMapper: """Manages task-to-agent mapping during installation.""" def __init__(self): """Initialize the task mapper.""" self.task_mappings: dict[str, str] = {} self.selected_strategy: str = "balanced" def select_strategy(self) -> str: """ Prompt user to select a routing strategy. Returns: Selected strategy key """ console.print("\n[bold]Delegation Strategy[/bold]") console.print("Choose how tasks should be distributed among agents:\n") table = Table(show_header=True, header_style="bold magenta") table.add_column("Option", style="cyan", justify="center") table.add_column("Strategy", style="green") table.add_column("Description", style="white") table.add_column("Priorities", style="yellow") strategies = list(ROUTING_PRESETS.items()) for i, (key, preset) in enumerate(strategies, 1): priorities = f"Cost: {preset['cost_priority']}, Quality: {preset['quality_priority']}" table.add_row(str(i), preset["name"], preset["description"], priorities) console.print(table) console.print("\n") choices = [str(i) for i in range(1, len(strategies) + 1)] default_idx = [k for k, _ in strategies].index("balanced") + 1 selection = Prompt.ask( "Select strategy", choices=choices, default=str(default_idx) ) selected_key = strategies[int(selection) - 1][0] self.selected_strategy = selected_key console.print(f"\n[green]✓[/green] Selected: {ROUTING_PRESETS[selected_key]['name']}\n") return selected_key def suggest_mappings(self, agent_names: list[str], strategy_key: str) -> dict[str, tuple[str, str]]: """ Generate intelligent mapping suggestions based on strategy and agent capabilities. Args: agent_names: List of available agent names strategy_key: Key of the selected routing strategy Returns: Dictionary of task_key -> (suggested_agent, reasoning) """ suggestions: dict[str, tuple[str, str]] = {} preset = ROUTING_PRESETS[strategy_key] # Helper to find best agent from a list of preferred ones def find_best_available(preferred: list[str], fallback_reason: str) -> tuple[str, str]: for agent in preferred: if agent in agent_names: # Find specific reason from rules if available return agent, fallback_reason # Fallback logic based on strategy if preset["cost_priority"] == "high": # Prefer free/local agents for agent in agent_names: profile = get_agent_profile(agent) if profile["cost_tier"] == "free": return agent, "Selected for cost efficiency" if preset["quality_priority"] == "high": # Prefer Claude/Gemini for agent in ["claude", "gemini"]: if agent in agent_names: return agent, "Selected for high quality" # Default to first available return agent_names[0], "Best available option" for category in TASK_CATEGORIES: task_key = category["key"] # Get default rule rule = DEFAULT_ROUTING_RULES.get(task_key) if not rule: suggestions[task_key] = (agent_names[0], "Default assignment") continue # Apply strategy overrides preferred = rule["preferred"] reason = rule["reason"] if strategy_key == "cost_optimized": # Prioritize free agents free_agents = [a for a in agent_names if get_agent_profile(a)["cost_tier"] == "free"] if free_agents: preferred = free_agents + preferred reason = "Cost optimized choice" elif strategy_key == "speed_first": # Prioritize fast agents fast_agents = [a for a in agent_names if get_agent_profile(a)["response_speed"] == "fast"] if fast_agents: preferred = fast_agents + preferred reason = "Optimized for speed" elif strategy_key == "token_saver": # Prioritize large context or concise agents # (Simplified logic: prefer Gemini for context, Aider for conciseness) if task_key in ["architecture", "exploration"]: preferred = ["gemini"] + preferred reason = "Large context window" else: preferred = ["aider"] + preferred reason = "Concise responses" # Find best agent agent, final_reason = find_best_available(preferred, reason) suggestions[task_key] = (agent, final_reason) return suggestions def display_suggestions( self, suggestions: dict[str, tuple[str, str]], agent_names: list[str] ) -> None: """ Display mapping suggestions in a formatted table. Args: suggestions: Dictionary of task_key -> (agent, reasoning) agent_names: List of available agent names for context """ table = Table( title=f"Suggested Mappings ({ROUTING_PRESETS[self.selected_strategy]['name']})", show_header=True, header_style="bold magenta" ) table.add_column("Task Category", style="cyan", no_wrap=True) table.add_column("Suggested Agent", style="green") table.add_column("Reasoning", style="yellow") # Create task key to category mapping for lookup category_map = {cat["key"]: cat for cat in TASK_CATEGORIES} for task_key, (agent, reasoning) in suggestions.items(): category = category_map.get(task_key) if category: task_name = category["name"] table.add_row(task_name, agent, reasoning) console.print("\n") console.print(table) console.print("\n") def prompt_task_assignments( self, agent_names: list[str], suggestions: dict[str, tuple[str, str]] ) -> dict[str, str]: """ Interactive prompt for task-to-agent assignment. Args: agent_names: List of available agent names suggestions: Pre-computed suggestions Returns: Dictionary of task_key -> agent_name """ self.display_suggestions(suggestions, agent_names) console.print("[bold]Task Assignment Configuration[/bold]") console.print("You can accept all suggestions or customize individual mappings.\n") # Ask if user wants to use all suggestions accept_all = Confirm.ask( "Accept all suggested mappings?", default=True ) if accept_all: self.task_mappings = { task_key: agent for task_key, (agent, _) in suggestions.items() } console.print("\n[green]✓[/green] Using all suggested mappings\n") return self.task_mappings # Custom assignment console.print("\nCustomize task assignments:\n") self.task_mappings = {} # Create task key to category mapping category_map = {cat["key"]: cat for cat in TASK_CATEGORIES} for task_key, (suggested_agent, reasoning) in suggestions.items(): category = category_map.get(task_key) if not category: continue task_name = category["name"] description = category["description"] console.print(f"\n[cyan]{task_name}[/cyan]: {description}") console.print(f" Suggested: [green]{suggested_agent}[/green] ({reasoning})") # Ask if user wants to change use_suggestion = Confirm.ask( f" Use {suggested_agent} for {task_name}?", default=True ) if use_suggestion: self.task_mappings[task_key] = suggested_agent console.print(f" [green]✓[/green] Assigned to {suggested_agent}") else: # Let user pick an agent console.print(f" Available agents: {', '.join(agent_names)}") while True: chosen_agent = Prompt.ask( f" Select agent for {task_name}", choices=agent_names, default=suggested_agent ) if chosen_agent in agent_names: self.task_mappings[task_key] = chosen_agent console.print(f" [green]✓[/green] Assigned to {chosen_agent}") break else: console.print(f" [red]✗[/red] Invalid agent. Please choose from: {', '.join(agent_names)}") console.print(f"\n[green]✓[/green] Task assignment configuration complete\n") return self.task_mappings def get_task_mappings(self) -> dict[str, str]: """ Get the task-to-agent mappings. Returns: Dictionary of task_key -> agent_name """ return self.task_mappings def map_tasks(self, agent_names: list[str]) -> dict[str, str]: """ Complete task mapping flow with suggestions and user input. Args: agent_names: List of available agent names Returns: Dictionary of task_key -> agent_name """ if len(agent_names) < 2: logger.warning("Need at least 2 agents for task mapping") return {} # Select strategy strategy_key = self.select_strategy() # Generate suggestions based on strategy suggestions = self.suggest_mappings(agent_names, strategy_key) # Get user assignments return self.prompt_task_assignments(agent_names, suggestions)