"""Data models for mosaic-coordinator.""" from enum import Enum 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.""" CONTINUE = "continue" # Below compact threshold, keep working COMPACT = "compact" # Hit 80% threshold, summarize and compact ROTATE_SESSION = "rotate_session" # Hit 95% threshold, spawn new agent class ContextUsage: """Agent context usage information.""" def __init__(self, agent_id: str, used_tokens: int, total_tokens: int) -> None: """Initialize context usage. Args: agent_id: Unique identifier for the agent used_tokens: Number of tokens currently used total_tokens: Total token capacity for this agent """ self.agent_id = agent_id self.used_tokens = used_tokens self.total_tokens = total_tokens @property def usage_ratio(self) -> float: """Calculate usage as a ratio (0.0-1.0). Returns: Ratio of used tokens to total capacity """ if self.total_tokens == 0: return 0.0 return self.used_tokens / self.total_tokens @property def usage_percent(self) -> float: """Calculate usage as a percentage (0-100). Returns: Percentage of context used """ return self.usage_ratio * 100 def __repr__(self) -> str: """String representation.""" return ( f"ContextUsage(agent_id={self.agent_id!r}, " f"used={self.used_tokens}, total={self.total_tokens}, " f"usage={self.usage_percent:.1f}%)" ) class IssueMetadata(BaseModel): """Parsed metadata from issue body.""" estimated_context: int = Field( default=50000, description="Estimated context size in tokens", ge=0 ) difficulty: Literal["easy", "medium", "hard"] = Field( default="medium", description="Issue difficulty level" ) assigned_agent: Literal["sonnet", "haiku", "opus", "glm"] = Field( default="sonnet", description="Recommended AI agent for this issue" ) blocks: list[int] = Field( default_factory=list, description="List of issue numbers this issue blocks" ) blocked_by: list[int] = Field( default_factory=list, description="List of issue numbers blocking this issue" ) @field_validator("difficulty", mode="before") @classmethod def validate_difficulty(cls, v: str) -> str: """Validate difficulty, default to medium if invalid.""" valid_values = ["easy", "medium", "hard"] if v not in valid_values: return "medium" return v @field_validator("assigned_agent", mode="before") @classmethod def validate_agent(cls, v: str) -> str: """Validate agent, default to sonnet if invalid.""" valid_values = ["sonnet", "haiku", "opus", "glm"] if v not in valid_values: return "sonnet" return v @field_validator("blocks", "blocked_by", mode="before") @classmethod def validate_issue_lists(cls, v: list[int] | None) -> list[int]: """Ensure issue lists are never None.""" 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]