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]
|
||||
|
||||
Reference in New Issue
Block a user