portfolio-optimiser-claude/src/portfolio_optimiser_claude/costsim.py

"""Pre-run cost simulation (method-spec §8/§10 analog; S3.6; paritetsrad 18; K6).

Before ANY spend, the operator can see an ESTIMATED upper-bound USD cost for a
(portfolio-)run — a what-if over the models configured in ``model_map.json``
(Claude models) × effort levels. Pricing is schema-validated config
(``data/pricing.example.json``): a per-Mtok USD rate per model id, each carrying
a REQUIRED source + date so a stale price is visible, never silent (§1 honesty).
A model configured in ``model_map`` with no price fails FAST ("missing price for
<id>") — the estimate is never guessed from a hardcoded rate (there is no price
literal anywhere in this module; the grep-guard test proves it).

The estimate is a deterministic UPPER BOUND, not a forecast: it bills the whole
token cap (portfolio shape × per-project cap × an effort weight) at the model's
per-Mtok rate. The example rates are Anthropic's OUTPUT price (the higher of the
two published rates), so billing the whole cap at that single rate can only
overstate, never understate — the figure is marked ESTIMAT in the output. Real
runs cost less (input is cheaper and often cached, output is small). No network,
no model call, no key: pure config arithmetic (offline invariant; bound by the
import-purity test, mirroring ``okf.py``).

Run:  uv run python -m portfolio_optimiser_claude.costsim [--projects N]
      [--token-cap T] [--pricing FILE]
"""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path

from portfolio_optimiser_claude.contracts import (
    ModelMapContract,
    PricingContract,
    _bundled_model_map,
    load_pricing,
    load_reference_projects,
)

# Coarse deterministic weighting of expected token consumption as a FRACTION of
# the cap, per Claude effort level (low -> max). NOT measured: a modeling weight
# for the ESTIMATE, where ``max`` effort is modeled as consuming the full cap
# (the true upper bound) and lower efforts proportionally less. These are effort
# weights, not prices — the grep-guard forbids only price literals here.
_EFFORT_FACTORS: dict[str, float] = {
    "low": 0.2,
    "medium": 0.4,
    "high": 0.6,
    "xhigh": 0.8,
    "max": 1.0,
}
_DEFAULT_EFFORTS: tuple[str, ...] = ("low", "medium", "high", "xhigh", "max")
_TOKENS_PER_MTOK = 1_000_000
_DEFAULT_TOKEN_CAP = 150_000  # mirrors run.py's --max-tokens default (the §8 cap)


@dataclass(frozen=True)
class EffortEstimate:
    """One (model, effort) cell of the what-if grid — tokens + upper-bound USD."""

    effort: str
    estimated_tokens: int
    cost_usd: float


@dataclass(frozen=True)
class ModelEstimate:
    """One model's row: its sourced rate + the per-effort upper-bound estimates."""

    model_id: str
    usd_per_mtok: float
    source: str
    source_date: str
    efforts: list[EffortEstimate]


@dataclass(frozen=True)
class CostEstimate:
    """The full deterministic what-if: portfolio shape × cap × (model × effort)."""

    n_projects: int
    token_cap_per_project: int
    models: list[ModelEstimate]


def _configured_model_ids(model_map: ModelMapContract) -> list[str]:
    """The DISTINCT model ids configured anywhere in the map, sorted (determinism)."""
    ids: set[str] = set()
    for mapping in model_map.profiles.values():
        ids.update(mapping.values())
    return sorted(ids)


def estimate_costs(
    model_map: ModelMapContract,
    pricing: PricingContract,
    *,
    n_projects: int,
    token_cap_per_project: int,
    efforts: tuple[str, ...] = _DEFAULT_EFFORTS,
) -> CostEstimate:
    """Deterministic upper-bound cost estimate over (model_map models × efforts).

    Every configured model REQUIRES a price — a missing one raises
    ``ValueError("missing price for <id>")`` (fail-fast, never a guessed rate).
    The per-cell figure is ``n_projects × token_cap_per_project × effort_factor``
    tokens billed at the model's per-Mtok rate — an UPPER BOUND (the whole cap at
    the sourced rate), reproducible from the inputs alone (no clock, no random).
    """
    if n_projects <= 0:
        raise ValueError(f"n_projects must be positive, got {n_projects}")
    if token_cap_per_project <= 0:
        raise ValueError(f"token_cap_per_project must be positive, got {token_cap_per_project}")
    models: list[ModelEstimate] = []
    for model_id in _configured_model_ids(model_map):
        price = pricing.prices.get(model_id)
        if price is None:
            raise ValueError(f"missing price for {model_id}")
        cells: list[EffortEstimate] = []
        for effort in efforts:
            factor = _EFFORT_FACTORS.get(effort)
            if factor is None:
                raise ValueError(f"unknown effort level {effort!r}")
            estimated_tokens = int(n_projects * token_cap_per_project * factor)
            cost = round(estimated_tokens * price.usd_per_mtok / _TOKENS_PER_MTOK, 6)
            cells.append(
                EffortEstimate(effort=effort, estimated_tokens=estimated_tokens, cost_usd=cost)
            )
        models.append(
            ModelEstimate(
                model_id=model_id,
                usd_per_mtok=price.usd_per_mtok,
                source=price.source,
                source_date=price.source_date,
                efforts=cells,
            )
        )
    return CostEstimate(
        n_projects=n_projects, token_cap_per_project=token_cap_per_project, models=models
    )


def render_estimate(estimate: CostEstimate) -> str:
    """Render the what-if as an ESTIMAT-marked table (the honesty label is load-bearing)."""
    lines = [
        f"ESTIMAT (deterministic upper bound) — cost for a run of "
        f"{estimate.n_projects} project(s), token cap "
        f"{estimate.token_cap_per_project}/project.",
        "NB: upper bound — the whole cap is billed at each model's per-Mtok rate; "
        "real runs cost less (input is cheaper and often cached, output is small).",
    ]
    for model in estimate.models:
        lines.append(
            f"model {model.model_id}  (${model.usd_per_mtok}/Mtok  "
            f"source={model.source} {model.source_date})"
        )
        for cell in model.efforts:
            lines.append(
                f"    {cell.effort:<7} ~{cell.estimated_tokens:>12} tok  ~${cell.cost_usd:.6f}"
            )
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    """The thin CLI: schema-validate pricing (§10) → estimate (offline) → print."""
    parser = argparse.ArgumentParser(
        description=(
            "Estimate the upper-bound USD cost for a (portfolio-)run BEFORE any "
            "spend — a what-if over model_map models x effort levels (offline, ESTIMAT)."
        )
    )
    parser.add_argument("--projects", type=int, default=None)
    parser.add_argument("--token-cap", type=int, default=_DEFAULT_TOKEN_CAP)
    parser.add_argument("--pricing", type=Path, default=None)
    args = parser.parse_args(argv)

    # §10: pricing is schema-validated BEFORE any estimate; a bad price fails fast.
    pricing = load_pricing(
        json.loads(args.pricing.read_text(encoding="utf-8")) if args.pricing is not None else None
    )
    model_map = ModelMapContract(**_bundled_model_map())
    n_projects = (
        args.projects
        if args.projects is not None
        else len(load_reference_projects().projects)  # portfolio shape from the config
    )
    estimate = estimate_costs(
        model_map, pricing, n_projects=n_projects, token_cap_per_project=args.token_cap
    )
    print(render_estimate(estimate))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())