cowork-local/core/routing/models.py

"""Pydantic v2 data models for Auto Model Assessment & Routing.

These are the provider-agnostic shapes shared by every routing module — the
enricher, prober, scorer, selector and switch-controller all speak in terms of
these. They serialize cleanly to/from JSON so the assessment store and the
app config (``~/.cowork_local/…``) can round-trip them.

Terminology
-----------
* A **candidate** is a ``(provider, model_id)`` pair the app can call.
* An **assessment** is what we learned about one candidate: its static
  metadata, the dynamic probe results per task type, and the derived
  ``fit_scores`` per task type.
* A **task type** is the kind of work a message represents (qa / coding / …).
* A **policy** is how we weigh quality vs cost vs latency when scoring.
"""
from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Optional, Set

from pydantic import BaseModel, Field


# --------------------------------------------------------------------------- #
# Enums
# --------------------------------------------------------------------------- #
class TaskType(str, Enum):
    """The kinds of work a chat/agent turn can represent.

    A message is classified into exactly one of these before routing (see
    ``classifier.py``). ``BENCHMARK_TASKS`` in the prober has one fixed prompt
    per value so every candidate model is compared on the same yardstick.
    """

    QA = "qa"
    CODING = "coding"
    REASONING = "reasoning"
    SUMMARIZATION = "summarization"
    CREATIVE = "creative"


class Policy(str, Enum):
    """How to trade off quality, cost and latency when scoring a model."""

    QUALITY = "quality"
    COST = "cost"
    LATENCY = "latency"
    BALANCED = "balanced"


class SwitchMode(str, Enum):
    """Per-surface routing behaviour, driven by the Off/Auto/Manual toggle.

    * ``OFF``    — routing disabled; always use the manually-selected model.
    * ``AUTO``   — silently switch to the best model when it clears the gain
                   threshold, then run the task.
    * ``MANUAL`` — propose the switch and wait for the user to confirm before
                   running with the new model.
    """

    OFF = "off"
    AUTO = "auto"
    MANUAL = "manual"


class SwitchStatus(str, Enum):
    """Lifecycle of a :class:`PendingSwitch` awaiting user confirmation."""

    PENDING = "pending"
    CONFIRMED = "confirmed"
    REJECTED = "rejected"
    EXPIRED = "expired"


# --------------------------------------------------------------------------- #
# Static metadata + dynamic probe
# --------------------------------------------------------------------------- #
class ModelMetadata(BaseModel):
    """Static, mostly-price/capability facts about one candidate model.

    ``cost_per_1k_*`` are USD per 1,000 tokens. They are ``None`` — not a
    guess — when the price is genuinely unknown; ``metadata_incomplete`` is
    then set True so the scorer/UI can flag it rather than silently trusting a
    fabricated number (see ``metadata.py``).
    """

    provider: str
    model_id: str
    tier: Optional[str] = None  # e.g. "fast" | "powerful" — free-form, user-supplied
    cost_per_1k_input: Optional[float] = None
    cost_per_1k_output: Optional[float] = None
    max_context: Optional[int] = None
    capabilities: Set[str] = Field(default_factory=set)  # e.g. {"vision", "tools"}
    available: bool = True
    metadata_incomplete: bool = False

    @property
    def key(self) -> str:
        """Stable ``provider/model_id`` identity used as a dict key everywhere."""
        return candidate_key(self.provider, self.model_id)

    @property
    def avg_cost_per_1k(self) -> Optional[float]:
        """Blended input/output price, or None if either side is unknown.

        A rough 1:3 input:output ratio (typical chat workload) is used so a
        single scalar can feed the cost term of the fit score.
        """
        ci, co = self.cost_per_1k_input, self.cost_per_1k_output
        if ci is None or co is None:
            return None
        return (ci + 3.0 * co) / 4.0


class ProbeResult(BaseModel):
    """Outcome of running one benchmark task against one model.

    ``success=False`` means the call itself failed (network/auth/model error);
    ``error`` then holds a human-readable reason and ``quality_score`` stays 0.
    """

    latency_ms: float = 0.0
    success: bool = False
    quality_score: float = 0.0  # 0..1, from the judge model
    tokens_out: int = 0
    error: Optional[str] = None


class ModelAssessment(BaseModel):
    """Everything we know about one candidate after an assessment run."""

    metadata: ModelMetadata
    # Keyed by TaskType.value (JSON-friendly string keys).
    probes: Dict[str, ProbeResult] = Field(default_factory=dict)
    fit_scores: Dict[str, float] = Field(default_factory=dict)
    assessed_at: Optional[str] = None  # ISO-8601 UTC timestamp

    @property
    def key(self) -> str:
        """Khoá định danh của model được chấm điểm (provider + model id)."""
        return self.metadata.key

    def fit_for(self, task_type: TaskType) -> float:
        """Fit score for ``task_type`` (0.0 if this model was never scored for it)."""
        return float(self.fit_scores.get(task_type.value, 0.0))


# --------------------------------------------------------------------------- #
# Switch decision + pending confirmation
# --------------------------------------------------------------------------- #
class SwitchDecision(BaseModel):
    """The verdict of comparing the current model against the selector's best.

    ``should_switch`` is False when routing is Off, when the best candidate IS
    the current model, or when the score gain is below ``min_score_gain``.
    """

    should_switch: bool
    from_model: Optional[str] = None  # candidate key, or None if nothing active yet
    to_model: Optional[str] = None
    from_score: float = 0.0
    to_score: float = 0.0
    score_gain: float = 0.0
    reason: str = ""
    mode: SwitchMode = SwitchMode.OFF
    task_type: Optional[str] = None


class PendingSwitch(BaseModel):
    """A Manual-mode switch proposal held until the user confirms/rejects.

    Stored in-memory with a TTL; ``result`` caches the executed task output so
    a repeated confirm of the same ``request_id`` is idempotent (returns the
    cached result instead of running the task twice).
    """

    request_id: str
    task_payload: Dict[str, Any] = Field(default_factory=dict)
    decision: SwitchDecision
    created_at: float  # epoch seconds (monotonic wall clock at creation)
    expires_at: float
    status: SwitchStatus = SwitchStatus.PENDING
    result: Optional[Dict[str, Any]] = None  # cached task result once executed


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def candidate_key(provider: str, model_id: str) -> str:
    """The canonical ``provider/model_id`` string used as a dict key."""
    return f"{provider}/{model_id}"


def split_key(key: str) -> tuple[str, str]:
    """Inverse of :func:`candidate_key`. Splits on the first ``/`` only, so a
    model id that itself contains ``/`` (some gateways use ``org/model``) is
    preserved intact."""
    provider, _, model_id = key.partition("/")
    return provider, model_id


__all__ = [
    "TaskType",
    "Policy",
    "SwitchMode",
    "SwitchStatus",
    "ModelMetadata",
    "ProbeResult",
    "ModelAssessment",
    "SwitchDecision",
    "PendingSwitch",
    "candidate_key",
    "split_key",
]