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:
2026-02-01 18:00:19 -06:00
parent a1b911d836
commit f0fd0bed41
2 changed files with 505 additions and 0 deletions

View File

@@ -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]