6krrt/iteration.py

"""Pure logic for spending a tier's iteration budget.

Reframes what a tier means. It used to be only a capability floor —
"don't route below this". It is now also a budget for getting the answer
right: tier 1 buys one attempt, higher tiers buy corrective attempts after a
verification failure.

Why this is better than what it replaces. `apply_escalation` bumps the tier
*preemptively* when the classifier is unsure of its own call, so an uncertain
guess costs frontier prices before anything has gone wrong. Spending after a
check has actually failed is strictly better on both mandates: the cheap
attempt usually succeeds and costs nothing extra, and when it fails you have
evidence rather than a hunch.

The retry strategy depends on the failure, because the failures have different
causes:

- **truncated** — the answer ran out of budget. If there is a cap to raise,
  raise it on the same model. If there is not, the model's own output ceiling
  is the wall, so escalate to a candidate that can emit more. (Without that
  second case the branch was unreachable in practice: through /v1 the cap is
  either the client's, which is not ours to override, or absent.)
- **malformed** — the model produced something that would not parse. More
  tokens will not help, so this escalates to the next-best candidate.

Latency is treated as a cost. Every retry doubles time-to-answer, which in
interactive use *is* a quality loss, so interactive work gets a lower cap than
batch work regardless of tier.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, Optional

INTERACTIVE = "interactive"
BATCH = "batch"

# Multiplier applied to the token budget when retrying a truncated answer.
TRUNCATION_BUDGET_MULTIPLIER = 2


@dataclass
class RetryPlan:
    """How to spend one attempt from the budget."""

    reason: Literal["truncated", "malformed"]
    model_id: str
    max_tokens: Optional[int]
    detail: str


def attempts_allowed(
    tier: int,
    latency_tolerance: str,
    attempts_by_tier: dict[int, int],
    max_attempts_interactive: int,
) -> int:
    """How many corrective attempts this request has earned.

    Interactive requests are capped below their tier's budget: a user waiting
    on an answer pays for every retry in latency, and a slow correct answer can
    be worth less than a fast one they can judge themselves.
    """
    budget = attempts_by_tier.get(tier, 0)
    if latency_tolerance == INTERACTIVE:
        return min(budget, max_attempts_interactive)
    return budget


def plan_retry(
    verdict: str,
    current_model: str,
    current_max_tokens: Optional[int],
    model_ceiling: Optional[int],
    runners_up: list[tuple[str, Optional[int]]],
    client_capped: bool,
) -> Optional[RetryPlan]:
    """Decide how to retry a failed response, or decline to.

    ``runners_up`` is (model_id, output_ceiling) in ranked order, because a
    truncation retry needs to know which alternatives can actually hold a
    longer answer.

    Returns None when retrying cannot help, which is as important as returning
    a plan — a wasted attempt costs energy against a fixed quota.
    """
    if verdict == "truncated":
        # A client that set its own max_tokens chose this outcome. Overriding
        # it would ignore an explicit instruction, and the caller may well want
        # a short answer.
        if client_capped:
            return None

        if current_max_tokens is not None:
            bigger = current_max_tokens * TRUNCATION_BUDGET_MULTIPLIER
            if model_ceiling:
                bigger = min(bigger, model_ceiling)
            if bigger > current_max_tokens:
                return RetryPlan(
                    "truncated",
                    current_model,
                    bigger,
                    f"retrying {current_model} with {bigger} tokens",
                )

        # No cap to raise, or already at this model's limit: the model's own
        # output ceiling is the wall. Retrying it changes nothing, but a
        # candidate that can emit MORE might finish the answer.
        roomier = next(
            (
                (m, ceiling)
                for m, ceiling in runners_up
                if ceiling and (not model_ceiling or ceiling > model_ceiling)
            ),
            None,
        )
        if roomier is None:
            return None
        return RetryPlan(
            "truncated",
            roomier[0],
            current_max_tokens,
            f"escalating to {roomier[0]} ({roomier[1]} output tokens)",
        )

    if verdict == "malformed":
        # More tokens will not make unparseable output parse. Try the next
        # candidate the ranking preferred.
        if not runners_up:
            return None
        return RetryPlan(
            "malformed",
            runners_up[0][0],
            current_max_tokens,
            f"escalating to {runners_up[0][0]}",
        )

    # 'ok' and 'unverifiable' are not failures and buy no retry. Spending an
    # attempt on 'unverifiable' would burn quota on the majority of prose
    # traffic for no signal at all.
    return None