stack/apps/coordinator/src/agent_assignment.py

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