"""Intelligent agent assignment algorithm. Selects the optimal agent for an issue based on: 1. Context capacity (50% rule: agent must have 2x estimated context) 2. Difficulty capability (agent must be able to handle issue difficulty) 3. Cost optimization (prefer cheapest qualifying agent) 4. Self-hosted preference (prefer cost=0 agents when capable) """ from typing import Literal from src.models import AGENT_PROFILES, AgentName, AgentProfile, Capability class NoCapableAgentError(Exception): """Raised when no agent can handle the given requirements.""" def __init__(self, estimated_context: int, difficulty: str) -> None: """Initialize error with context details. Args: estimated_context: Required context size in tokens difficulty: Issue difficulty level """ super().__init__( f"No capable agent found for difficulty={difficulty!r} " f"with estimated_context={estimated_context} tokens. " f"Consider breaking down the issue into smaller parts." ) self.estimated_context = estimated_context self.difficulty = difficulty def _map_difficulty_to_capability(difficulty: str) -> Capability: """Map difficulty string to Capability enum. Args: difficulty: Issue difficulty level Returns: Corresponding Capability level Raises: ValueError: If difficulty is not valid """ difficulty_lower = difficulty.lower() mapping = { "easy": Capability.LOW, "low": Capability.LOW, "medium": Capability.MEDIUM, "hard": Capability.HIGH, "high": Capability.HIGH, } if difficulty_lower not in mapping: raise ValueError( f"Invalid difficulty: {difficulty!r}. " f"Must be one of: {list(mapping.keys())}" ) return mapping[difficulty_lower] def _can_handle_context(profile: AgentProfile, estimated_context: int) -> bool: """Check if agent can handle context using 50% rule. Agent must have at least 2x the estimated context to ensure adequate working room and prevent context exhaustion. Args: profile: Agent profile to check estimated_context: Estimated context requirement in tokens Returns: True if agent can handle the context, False otherwise """ required_capacity = estimated_context * 2 return profile.context_limit >= required_capacity def _can_handle_difficulty(profile: AgentProfile, capability: Capability) -> bool: """Check if agent can handle the required difficulty level. Args: profile: Agent profile to check capability: Required capability level Returns: True if agent has the required capability, False otherwise """ return capability in profile.capabilities def _filter_qualified_agents( estimated_context: int, capability: Capability ) -> list[AgentProfile]: """Filter agents that meet context and capability requirements. Args: estimated_context: Required context size in tokens capability: Required capability level Returns: List of qualified agent profiles """ qualified: list[AgentProfile] = [] for profile in AGENT_PROFILES.values(): # Check both context capacity and difficulty capability if (_can_handle_context(profile, estimated_context) and _can_handle_difficulty(profile, capability)): qualified.append(profile) return qualified def _sort_by_cost(profiles: list[AgentProfile]) -> list[AgentProfile]: """Sort agents by cost, preferring self-hosted (cost=0). Agents are sorted by: 1. Cost (ascending) - cheapest first 2. Name (for stable ordering when costs are equal) Args: profiles: List of agent profiles to sort Returns: Sorted list of profiles """ return sorted(profiles, key=lambda p: (p.cost_per_mtok, p.name.value)) def assign_agent( estimated_context: int, difficulty: Literal["easy", "medium", "hard", "low", "high"] ) -> AgentName: """Assign the optimal agent for an issue. Selection algorithm: 1. Filter agents that meet context capacity (50% rule) 2. Filter agents that can handle difficulty level 3. Sort by cost (prefer self-hosted when capable) 4. Return cheapest qualifying agent Args: estimated_context: Estimated context requirement in tokens difficulty: Issue difficulty level Returns: Name of the assigned agent Raises: ValueError: If estimated_context is negative or difficulty is invalid NoCapableAgentError: If no agent can handle the requirements """ # Validate inputs if estimated_context < 0: raise ValueError( f"estimated_context must be non-negative, got {estimated_context}" ) # Map difficulty to capability capability = _map_difficulty_to_capability(difficulty) # Filter agents that meet requirements qualified_agents = _filter_qualified_agents(estimated_context, capability) # If no agents qualify, raise error if not qualified_agents: raise NoCapableAgentError(estimated_context, difficulty) # Sort by cost and select cheapest sorted_agents = _sort_by_cost(qualified_agents) selected_agent = sorted_agents[0] return selected_agent.name