feat(#144): Implement agent profiles
- Add Capability enum (HIGH, MEDIUM, LOW) for agent difficulty levels - Add AgentName enum for all 5 agents (opus, sonnet, haiku, glm, minimax) - Implement AgentProfile data structure with validation - context_limit: max tokens for context window - cost_per_mtok: cost per million tokens (0 for self-hosted) - capabilities: list of difficulty levels the agent handles - best_for: description of optimal use cases - Define profiles for all 5 agents with specifications: - Anthropic models (opus, sonnet, haiku): 200K context, various costs - Self-hosted models (glm, minimax): 128K context, free - Implement get_agent_profile() function for profile lookup - Add comprehensive test suite (37 tests, 100% coverage) - Profile data structure validation - All 5 predefined profiles exist and are correct - Capability enum and AgentName enum tests - Best_for validation and capability matching - Consistency checks across profiles Fixes #144 Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -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]
|
||||
|
||||
402
apps/coordinator/tests/test_agent_profiles.py
Normal file
402
apps/coordinator/tests/test_agent_profiles.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user