diff --git a/apps/coordinator/src/agent_assignment.py b/apps/coordinator/src/agent_assignment.py new file mode 100644 index 0000000..1ac72d5 --- /dev/null +++ b/apps/coordinator/src/agent_assignment.py @@ -0,0 +1,177 @@ +"""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 diff --git a/apps/coordinator/tests/test_agent_assignment.py b/apps/coordinator/tests/test_agent_assignment.py new file mode 100644 index 0000000..2114ba5 --- /dev/null +++ b/apps/coordinator/tests/test_agent_assignment.py @@ -0,0 +1,261 @@ +"""Tests for agent assignment algorithm. + +Test scenarios: +1. Assignment for low/medium/high difficulty issues +2. Context capacity filtering (50% rule enforcement) +3. Cost optimization logic +4. Error handling for impossible assignments +""" + +import pytest + +from src.agent_assignment import NoCapableAgentError, assign_agent +from src.models import AgentName, AGENT_PROFILES + + +class TestAgentAssignment: + """Test the intelligent agent assignment algorithm.""" + + def test_assign_low_difficulty_prefers_cheapest(self) -> None: + """Test that low difficulty issues get assigned to cheapest capable agent.""" + # For low difficulty with small context (25K tokens), expect cheapest self-hosted + # Both GLM and minimax are cost=0, GLM comes first alphabetically + assigned = assign_agent( + estimated_context=25000, + difficulty="easy" + ) + assert assigned == AgentName.GLM + + def test_assign_low_difficulty_large_context_uses_haiku(self) -> None: + """Test that low difficulty with larger context uses Haiku.""" + # minimax and GLM have 128K limit (can handle up to 64K) + # 100K * 2 (50% rule) = needs 200K capacity + # Should use Haiku (200K context, cheapest commercial for low) + assigned = assign_agent( + estimated_context=100000, + difficulty="easy" + ) + assert assigned == AgentName.HAIKU + + def test_assign_low_difficulty_within_self_hosted_uses_glm(self) -> None: + """Test that low difficulty within self-hosted capacity uses GLM.""" + # 60K tokens needs 120K capacity (50% rule) + # GLM has 128K limit (can handle up to 64K) + # Should use GLM (self-hosted, cost=0) + assigned = assign_agent( + estimated_context=60000, + difficulty="easy" + ) + assert assigned == AgentName.GLM + + def test_assign_medium_difficulty_prefers_glm(self) -> None: + """Test that medium difficulty prefers self-hosted GLM when possible.""" + # GLM is self-hosted (cost=0) and can handle medium difficulty + assigned = assign_agent( + estimated_context=30000, + difficulty="medium" + ) + assert assigned == AgentName.GLM + + def test_assign_medium_difficulty_large_context_uses_sonnet(self) -> None: + """Test that medium difficulty with large context uses Sonnet.""" + # 80K tokens needs 160K capacity (50% rule) + # GLM has 128K limit (can handle up to 64K) + # Should use Sonnet (200K context, cheapest commercial for medium) + assigned = assign_agent( + estimated_context=80000, + difficulty="medium" + ) + assert assigned == AgentName.SONNET + + def test_assign_high_difficulty_uses_opus(self) -> None: + """Test that high difficulty always uses Opus.""" + # Only Opus can handle high difficulty + assigned = assign_agent( + estimated_context=50000, + difficulty="hard" + ) + assert assigned == AgentName.OPUS + + def test_assign_high_difficulty_large_context_uses_opus(self) -> None: + """Test that high difficulty with large context still uses Opus.""" + # Even with large context, Opus is the only option for high difficulty + assigned = assign_agent( + estimated_context=90000, + difficulty="hard" + ) + assert assigned == AgentName.OPUS + + def test_fifty_percent_rule_enforced(self) -> None: + """Test that 50% context capacity rule is strictly enforced.""" + # 65K tokens needs 130K capacity (50% rule) + # GLM has 128K limit, so can't handle this + # Should use Sonnet (200K limit, can handle up to 100K) + assigned = assign_agent( + estimated_context=65000, + difficulty="medium" + ) + assert assigned == AgentName.SONNET + + def test_self_hosted_preferred_when_capable(self) -> None: + """Test that self-hosted agents are preferred over commercial when capable.""" + # For medium difficulty with 30K context: + # GLM (self-hosted, cost=0) can handle it + # Sonnet (commercial, cost=3.0) can also handle it + # Should prefer GLM + assigned = assign_agent( + estimated_context=30000, + difficulty="medium" + ) + assert assigned == AgentName.GLM + + def test_impossible_assignment_raises_error(self) -> None: + """Test that impossible assignments raise NoCapableAgentError.""" + # No agent can handle 150K tokens (needs 300K capacity with 50% rule) + # Max capacity is 200K (Opus, Sonnet, Haiku) + with pytest.raises(NoCapableAgentError) as exc_info: + assign_agent( + estimated_context=150000, + difficulty="medium" + ) + assert "No capable agent found" in str(exc_info.value) + assert "150000" in str(exc_info.value) + + def test_impossible_assignment_high_difficulty_massive_context(self) -> None: + """Test error when even Opus cannot handle the context.""" + # Opus has 200K limit, so can handle up to 100K with 50% rule + # This should fail + with pytest.raises(NoCapableAgentError) as exc_info: + assign_agent( + estimated_context=120000, + difficulty="hard" + ) + assert "No capable agent found" in str(exc_info.value) + + def test_edge_case_exact_fifty_percent(self) -> None: + """Test edge case where context exactly meets 50% threshold.""" + # 100K tokens needs exactly 200K capacity + # Haiku, Sonnet, Opus all have 200K + # For low difficulty, should use Haiku (cheapest) + assigned = assign_agent( + estimated_context=100000, + difficulty="easy" + ) + # GLM can only handle 64K (128K / 2), so needs commercial + assert assigned == AgentName.HAIKU + + def test_agent_selection_by_cost_ordering(self) -> None: + """Test that agents are selected by cost when multiple are capable.""" + # For low difficulty with 20K context, multiple agents qualify: + # - GLM (cost=0, 128K limit) - comes first alphabetically + # - minimax (cost=0, 128K limit) + # - Haiku (cost=0.8, 200K limit) + # - Sonnet (cost=3.0, 200K limit) + # Should pick cheapest: GLM (cost=0, alphabetically first) + assigned = assign_agent( + estimated_context=20000, + difficulty="easy" + ) + # GLM selected due to alphabetical ordering when costs are equal + assert assigned == AgentName.GLM + + def test_capability_filtering_excludes_incapable_agents(self) -> None: + """Test that agents without required capability are excluded.""" + # For medium difficulty: + # - minimax cannot handle medium (only LOW) + # - Haiku cannot handle medium (only LOW) + # Valid options: GLM, Sonnet, Opus + # Should prefer GLM (self-hosted, cost=0) + assigned = assign_agent( + estimated_context=30000, + difficulty="medium" + ) + assert assigned == AgentName.GLM + assert assigned not in [AgentName.MINIMAX, AgentName.HAIKU] + + def test_zero_context_estimate(self) -> None: + """Test assignment with zero context estimate.""" + # Zero context should work with any agent + # For low difficulty, should get cheapest (GLM comes first alphabetically) + assigned = assign_agent( + estimated_context=0, + difficulty="easy" + ) + assert assigned == AgentName.GLM + + def test_small_context_estimate(self) -> None: + """Test assignment with very small context estimate.""" + # 1K tokens should work with any agent (GLM comes first alphabetically) + assigned = assign_agent( + estimated_context=1000, + difficulty="easy" + ) + assert assigned == AgentName.GLM + + +class TestAgentAssignmentEdgeCases: + """Test edge cases and boundary conditions.""" + + def test_difficulty_case_insensitive(self) -> None: + """Test that difficulty matching is case-insensitive.""" + # Should handle different casings of difficulty + assigned_lower = assign_agent(estimated_context=30000, difficulty="easy") + assigned_title = assign_agent(estimated_context=30000, difficulty="easy") + assert assigned_lower == assigned_title + + def test_max_capacity_for_each_agent(self) -> None: + """Test maximum handleable context for each agent type.""" + # minimax: 128K / 2 = 64K max + assigned = assign_agent(estimated_context=64000, difficulty="easy") + assert assigned in [AgentName.MINIMAX, AgentName.GLM] + + # GLM: 128K / 2 = 64K max + assigned = assign_agent(estimated_context=64000, difficulty="medium") + assert assigned == AgentName.GLM + + # Opus: 200K / 2 = 100K max + assigned = assign_agent(estimated_context=100000, difficulty="hard") + assert assigned == AgentName.OPUS + + def test_negative_context_raises_error(self) -> None: + """Test that negative context raises appropriate error.""" + with pytest.raises(ValueError) as exc_info: + assign_agent(estimated_context=-1000, difficulty="easy") + assert "negative" in str(exc_info.value).lower() + + def test_invalid_difficulty_raises_error(self) -> None: + """Test that invalid difficulty raises appropriate error.""" + with pytest.raises(ValueError) as exc_info: + assign_agent(estimated_context=30000, difficulty="invalid") # type: ignore + assert "difficulty" in str(exc_info.value).lower() + + +class TestAgentAssignmentIntegration: + """Integration tests with actual agent profiles.""" + + def test_uses_actual_agent_profiles(self) -> None: + """Test that assignment uses actual AGENT_PROFILES data.""" + assigned = assign_agent(estimated_context=30000, difficulty="medium") + assert assigned in AGENT_PROFILES + profile = AGENT_PROFILES[assigned] + assert profile.context_limit >= 60000 # 30K * 2 for 50% rule + + def test_all_difficulty_levels_have_assignments(self) -> None: + """Test that all difficulty levels can be assigned for reasonable contexts.""" + # Test each difficulty level + easy_agent = assign_agent(estimated_context=30000, difficulty="easy") + assert easy_agent in AGENT_PROFILES + + medium_agent = assign_agent(estimated_context=30000, difficulty="medium") + assert medium_agent in AGENT_PROFILES + + hard_agent = assign_agent(estimated_context=30000, difficulty="hard") + assert hard_agent in AGENT_PROFILES + + def test_cost_optimization_verified_with_profiles(self) -> None: + """Test that cost optimization actually selects cheaper agents.""" + # For medium difficulty with 30K context: + # GLM (cost=0) should be selected over Sonnet (cost=3.0) + assigned = assign_agent(estimated_context=30000, difficulty="medium") + assigned_cost = AGENT_PROFILES[assigned].cost_per_mtok + assert assigned_cost == 0.0 # Self-hosted