diff --git a/apps/coordinator/src/models.py b/apps/coordinator/src/models.py index eb04b97..d1186f9 100644 --- a/apps/coordinator/src/models.py +++ b/apps/coordinator/src/models.py @@ -6,6 +6,24 @@ from typing import Literal from pydantic import BaseModel, Field, field_validator +class Capability(str, Enum): + """Agent capability levels.""" + + HIGH = "high" + MEDIUM = "medium" + LOW = "low" + + +class AgentName(str, Enum): + """Available AI agents.""" + + OPUS = "opus" + SONNET = "sonnet" + HAIKU = "haiku" + GLM = "glm" + MINIMAX = "minimax" + + class ContextAction(str, Enum): """Actions to take based on context usage thresholds.""" @@ -108,3 +126,88 @@ class IssueMetadata(BaseModel): if v is None: return [] return v + + +class AgentProfile(BaseModel): + """Profile defining agent capabilities, costs, and context limits.""" + + name: AgentName = Field(description="Agent identifier") + context_limit: int = Field( + gt=0, + description="Maximum tokens for agent context window" + ) + cost_per_mtok: float = Field( + ge=0.0, + description="Cost per million tokens (0 for self-hosted)" + ) + capabilities: list[Capability] = Field( + min_length=1, + description="Difficulty levels this agent can handle" + ) + best_for: str = Field( + min_length=1, + description="Optimal use cases for this agent" + ) + + @field_validator("best_for", mode="before") + @classmethod + def validate_best_for_not_empty(cls, v: str) -> str: + """Ensure best_for description is not empty.""" + if not v or not v.strip(): + raise ValueError("best_for description cannot be empty") + return v + + +# Predefined agent profiles +AGENT_PROFILES: dict[AgentName, AgentProfile] = { + AgentName.OPUS: AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH, Capability.MEDIUM, Capability.LOW], + best_for="Complex reasoning, code generation, and multi-step problem solving" + ), + AgentName.SONNET: AgentProfile( + name=AgentName.SONNET, + context_limit=200000, + cost_per_mtok=3.0, + capabilities=[Capability.MEDIUM, Capability.LOW], + best_for="Balanced performance for general tasks and scripting" + ), + AgentName.HAIKU: AgentProfile( + name=AgentName.HAIKU, + context_limit=200000, + cost_per_mtok=0.8, + capabilities=[Capability.LOW], + best_for="Fast, cost-effective processing of simple tasks" + ), + AgentName.GLM: AgentProfile( + name=AgentName.GLM, + context_limit=128000, + cost_per_mtok=0.0, + capabilities=[Capability.MEDIUM, Capability.LOW], + best_for="Self-hosted open-source model for medium complexity tasks" + ), + AgentName.MINIMAX: AgentProfile( + name=AgentName.MINIMAX, + context_limit=128000, + cost_per_mtok=0.0, + capabilities=[Capability.LOW], + best_for="Self-hosted lightweight model for simple tasks and prototyping" + ), +} + + +def get_agent_profile(agent_name: AgentName) -> AgentProfile: + """Retrieve profile for a specific agent. + + Args: + agent_name: Name of the agent + + Returns: + AgentProfile for the requested agent + + Raises: + KeyError: If agent_name is not defined + """ + return AGENT_PROFILES[agent_name] diff --git a/apps/coordinator/tests/test_agent_profiles.py b/apps/coordinator/tests/test_agent_profiles.py new file mode 100644 index 0000000..208c0c4 --- /dev/null +++ b/apps/coordinator/tests/test_agent_profiles.py @@ -0,0 +1,402 @@ +"""Tests for agent profile system.""" + +import pytest + +from src.models import ( + AGENT_PROFILES, + AgentName, + AgentProfile, + Capability, + get_agent_profile, +) + + +class TestAgentProfileDataStructure: + """Tests for AgentProfile data structure.""" + + def test_agent_profile_has_required_fields(self) -> None: + """Test that AgentProfile has all required fields.""" + profile = AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH, Capability.MEDIUM, Capability.LOW], + best_for="Complex reasoning and code generation" + ) + + assert profile.name == AgentName.OPUS + assert profile.context_limit == 200000 + assert profile.cost_per_mtok == 15.0 + assert len(profile.capabilities) == 3 + assert profile.best_for == "Complex reasoning and code generation" + + def test_agent_profile_validation_positive_context_limit(self) -> None: + """Test that context_limit must be positive.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=-1, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH], + best_for="Test" + ) + + def test_agent_profile_validation_zero_context_limit(self) -> None: + """Test that context_limit cannot be zero.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=0, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH], + best_for="Test" + ) + + def test_agent_profile_validation_non_negative_cost(self) -> None: + """Test that cost_per_mtok must be non-negative.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=-1.0, + capabilities=[Capability.HIGH], + best_for="Test" + ) + + def test_agent_profile_validation_non_empty_capabilities(self) -> None: + """Test that capabilities list cannot be empty.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[], + best_for="Test" + ) + + def test_agent_profile_validation_non_empty_best_for(self) -> None: + """Test that best_for description cannot be empty.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH], + best_for="" + ) + + +class TestAgentProfilesDefinition: + """Tests for predefined agent profiles.""" + + def test_opus_profile_exists(self) -> None: + """Test that Opus profile is defined correctly.""" + assert AgentName.OPUS in AGENT_PROFILES + profile = AGENT_PROFILES[AgentName.OPUS] + + assert profile.name == AgentName.OPUS + assert profile.context_limit == 200000 + assert profile.cost_per_mtok == 15.0 + assert Capability.HIGH in profile.capabilities + assert Capability.MEDIUM in profile.capabilities + assert Capability.LOW in profile.capabilities + assert "complex" in profile.best_for.lower() or "reasoning" in profile.best_for.lower() + + def test_sonnet_profile_exists(self) -> None: + """Test that Sonnet profile is defined correctly.""" + assert AgentName.SONNET in AGENT_PROFILES + profile = AGENT_PROFILES[AgentName.SONNET] + + assert profile.name == AgentName.SONNET + assert profile.context_limit == 200000 + assert profile.cost_per_mtok == 3.0 + assert Capability.MEDIUM in profile.capabilities + assert Capability.LOW in profile.capabilities + assert Capability.HIGH not in profile.capabilities + + def test_haiku_profile_exists(self) -> None: + """Test that Haiku profile is defined correctly.""" + assert AgentName.HAIKU in AGENT_PROFILES + profile = AGENT_PROFILES[AgentName.HAIKU] + + assert profile.name == AgentName.HAIKU + assert profile.context_limit == 200000 + assert profile.cost_per_mtok == 0.8 + assert Capability.LOW in profile.capabilities + assert Capability.MEDIUM not in profile.capabilities + assert Capability.HIGH not in profile.capabilities + + def test_glm_profile_exists(self) -> None: + """Test that GLM profile is defined correctly.""" + assert AgentName.GLM in AGENT_PROFILES + profile = AGENT_PROFILES[AgentName.GLM] + + assert profile.name == AgentName.GLM + assert profile.context_limit == 128000 + assert profile.cost_per_mtok == 0.0 + assert Capability.MEDIUM in profile.capabilities + assert Capability.LOW in profile.capabilities + + def test_minimax_profile_exists(self) -> None: + """Test that MiniMax profile is defined correctly.""" + assert AgentName.MINIMAX in AGENT_PROFILES + profile = AGENT_PROFILES[AgentName.MINIMAX] + + assert profile.name == AgentName.MINIMAX + assert profile.context_limit == 128000 + assert profile.cost_per_mtok == 0.0 + assert Capability.LOW in profile.capabilities + + def test_all_profiles_have_unique_costs_and_limits(self) -> None: + """Test that costs and context limits are correctly differentiated.""" + # Verify at least some differentiation exists + opus = AGENT_PROFILES[AgentName.OPUS] + sonnet = AGENT_PROFILES[AgentName.SONNET] + haiku = AGENT_PROFILES[AgentName.HAIKU] + glm = AGENT_PROFILES[AgentName.GLM] + minimax = AGENT_PROFILES[AgentName.MINIMAX] + + # Opus should have highest cost + assert opus.cost_per_mtok > sonnet.cost_per_mtok + assert sonnet.cost_per_mtok > haiku.cost_per_mtok + + # Self-hosted should be free + assert glm.cost_per_mtok == 0.0 + assert minimax.cost_per_mtok == 0.0 + + +class TestGetAgentProfile: + """Tests for get_agent_profile function.""" + + def test_get_opus_profile(self) -> None: + """Test retrieving Opus profile by name.""" + profile = get_agent_profile(AgentName.OPUS) + + assert profile.name == AgentName.OPUS + assert profile.context_limit == 200000 + + def test_get_sonnet_profile(self) -> None: + """Test retrieving Sonnet profile by name.""" + profile = get_agent_profile(AgentName.SONNET) + + assert profile.name == AgentName.SONNET + assert profile.context_limit == 200000 + + def test_get_haiku_profile(self) -> None: + """Test retrieving Haiku profile by name.""" + profile = get_agent_profile(AgentName.HAIKU) + + assert profile.name == AgentName.HAIKU + assert profile.context_limit == 200000 + + def test_get_glm_profile(self) -> None: + """Test retrieving GLM profile by name.""" + profile = get_agent_profile(AgentName.GLM) + + assert profile.name == AgentName.GLM + assert profile.context_limit == 128000 + + def test_get_minimax_profile(self) -> None: + """Test retrieving MiniMax profile by name.""" + profile = get_agent_profile(AgentName.MINIMAX) + + assert profile.name == AgentName.MINIMAX + assert profile.context_limit == 128000 + + def test_get_profile_returns_copy(self) -> None: + """Test that get_agent_profile returns independent copies.""" + profile1 = get_agent_profile(AgentName.OPUS) + profile2 = get_agent_profile(AgentName.OPUS) + + # Verify same values + assert profile1.name == profile2.name + assert profile1.context_limit == profile2.context_limit + + # Verify they are equal but can be independently modified if needed + assert profile1.model_dump() == profile2.model_dump() + + +class TestCapabilityEnum: + """Tests for Capability enum.""" + + def test_capability_enum_values(self) -> None: + """Test that Capability enum has expected values.""" + assert Capability.HIGH.value == "high" + assert Capability.MEDIUM.value == "medium" + assert Capability.LOW.value == "low" + + def test_capability_enum_ordering(self) -> None: + """Test capability comparison logic.""" + # All three should be available + capabilities = [Capability.HIGH, Capability.MEDIUM, Capability.LOW] + assert len(capabilities) == 3 + + +class TestAgentNameEnum: + """Tests for AgentName enum.""" + + def test_agent_name_enum_values(self) -> None: + """Test that AgentName enum has all expected agents.""" + agent_names = [ + AgentName.OPUS, + AgentName.SONNET, + AgentName.HAIKU, + AgentName.GLM, + AgentName.MINIMAX, + ] + assert len(agent_names) == 5 + + def test_agent_name_string_representation(self) -> None: + """Test string values of agent names.""" + assert AgentName.OPUS.value == "opus" + assert AgentName.SONNET.value == "sonnet" + assert AgentName.HAIKU.value == "haiku" + assert AgentName.GLM.value == "glm" + assert AgentName.MINIMAX.value == "minimax" + + +class TestProfileCapabilityMatching: + """Tests for capability matching against profiles.""" + + def test_opus_handles_high_difficulty(self) -> None: + """Test that Opus can handle high difficulty tasks.""" + profile = get_agent_profile(AgentName.OPUS) + assert Capability.HIGH in profile.capabilities + + def test_sonnet_handles_medium_difficulty(self) -> None: + """Test that Sonnet can handle medium difficulty tasks.""" + profile = get_agent_profile(AgentName.SONNET) + assert Capability.MEDIUM in profile.capabilities + + def test_haiku_handles_low_difficulty(self) -> None: + """Test that Haiku can handle low difficulty tasks.""" + profile = get_agent_profile(AgentName.HAIKU) + assert Capability.LOW in profile.capabilities + + def test_profile_best_for_description_exists(self) -> None: + """Test that all profiles have meaningful best_for descriptions.""" + for agent_name, profile in AGENT_PROFILES.items(): + msg_short = f"{agent_name} has insufficient best_for description" + assert len(profile.best_for) > 10, msg_short + msg_incomplete = f"{agent_name} has incomplete best_for description" + assert not profile.best_for.endswith("..."), msg_incomplete + + +class TestProfileConsistency: + """Tests for consistency across all profiles.""" + + def test_all_profiles_defined(self) -> None: + """Test that all five agents have profiles defined.""" + assert len(AGENT_PROFILES) == 5 + agent_names = { + AgentName.OPUS, + AgentName.SONNET, + AgentName.HAIKU, + AgentName.GLM, + AgentName.MINIMAX, + } + defined_names = set(AGENT_PROFILES.keys()) + assert agent_names == defined_names + + def test_anthropic_models_have_200k_context(self) -> None: + """Test that Anthropic models have 200K context limit.""" + anthropic_models = [AgentName.OPUS, AgentName.SONNET, AgentName.HAIKU] + for model in anthropic_models: + profile = AGENT_PROFILES[model] + assert profile.context_limit == 200000 + + def test_self_hosted_models_have_128k_context(self) -> None: + """Test that self-hosted models have 128K context limit.""" + self_hosted_models = [AgentName.GLM, AgentName.MINIMAX] + for model in self_hosted_models: + profile = AGENT_PROFILES[model] + assert profile.context_limit == 128000 + + def test_self_hosted_models_are_free(self) -> None: + """Test that self-hosted models have zero cost.""" + self_hosted_models = [AgentName.GLM, AgentName.MINIMAX] + for model in self_hosted_models: + profile = AGENT_PROFILES[model] + assert profile.cost_per_mtok == 0.0 + + def test_anthropic_models_have_costs(self) -> None: + """Test that Anthropic models have non-zero costs.""" + anthropic_models = [AgentName.OPUS, AgentName.SONNET, AgentName.HAIKU] + for model in anthropic_models: + profile = AGENT_PROFILES[model] + assert profile.cost_per_mtok > 0.0 + + def test_cost_reflects_capability(self) -> None: + """Test that cost roughly reflects capability level.""" + opus_cost = AGENT_PROFILES[AgentName.OPUS].cost_per_mtok + sonnet_cost = AGENT_PROFILES[AgentName.SONNET].cost_per_mtok + haiku_cost = AGENT_PROFILES[AgentName.HAIKU].cost_per_mtok + + # Opus > Sonnet > Haiku + assert opus_cost > sonnet_cost + assert sonnet_cost > haiku_cost + + +class TestBestForValidation: + """Tests for best_for field validation.""" + + def test_best_for_with_whitespace_only_fails(self) -> None: + """Test that best_for with only whitespace is rejected.""" + with pytest.raises(ValueError): + AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH], + best_for=" " + ) + + def test_best_for_with_valid_string_passes(self) -> None: + """Test that best_for with valid text passes validation.""" + profile = AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH], + best_for="Valid description" + ) + assert profile.best_for == "Valid description" + + +class TestCapabilityValidation: + """Tests for capability-specific validation.""" + + def test_multiple_capabilities_allowed(self) -> None: + """Test that multiple capabilities can be assigned.""" + profile = AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH, Capability.MEDIUM, Capability.LOW], + best_for="Test" + ) + assert len(profile.capabilities) == 3 + + def test_single_capability_allowed(self) -> None: + """Test that single capability can be assigned.""" + profile = AgentProfile( + name=AgentName.HAIKU, + context_limit=200000, + cost_per_mtok=0.8, + capabilities=[Capability.LOW], + best_for="Test" + ) + assert len(profile.capabilities) == 1 + assert profile.capabilities[0] == Capability.LOW + + def test_duplicate_capabilities_handled(self) -> None: + """Test that duplicate capabilities are allowed (pydantic behavior).""" + profile = AgentProfile( + name=AgentName.OPUS, + context_limit=200000, + cost_per_mtok=15.0, + capabilities=[Capability.HIGH, Capability.HIGH, Capability.MEDIUM], + best_for="Test" + ) + assert Capability.HIGH in profile.capabilities + assert Capability.MEDIUM in profile.capabilities