stack/apps/coordinator/src/models.py

"""Data models for mosaic-coordinator."""

from enum import StrEnum
from typing import Literal

from pydantic import BaseModel, Field, field_validator


class Capability(StrEnum):
    """Agent capability levels."""

    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class AgentName(StrEnum):
    """Available AI agents."""

    OPUS = "opus"
    SONNET = "sonnet"
    HAIKU = "haiku"
    GLM = "glm"
    MINIMAX = "minimax"


class ContextAction(StrEnum):
    """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]