src/id/llm/usage.py

"""Token ledger: append-only per-call usage log + aggregation for `id costs`.

Each LLM call appends one JSON line to ``runtime/<session>/usage.jsonl``
(Section 4). World-generation calls (no session) are logged under the world's
own ``usage.jsonl`` so generation cost is attributable too.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

from ..models import UsageRecord


class UsageLedger:
    """Append-only writer bound to a single ``usage.jsonl`` path."""

    def __init__(self, path: Path) -> None:
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def record(self, rec: UsageRecord) -> None:
        with self.path.open("a", encoding="utf-8") as fh:
            fh.write(rec.model_dump_json() + "\n")


@dataclass
class Totals:
    prompt: int = 0
    completion: int = 0
    total: int = 0
    calls: int = 0

    def add(self, rec: UsageRecord) -> None:
        self.prompt += rec.prompt_tokens
        self.completion += rec.completion_tokens
        self.total += rec.total_tokens
        self.calls += 1


@dataclass
class CostReport:
    by_task: dict[str, Totals] = field(default_factory=dict)
    by_model: dict[str, Totals] = field(default_factory=dict)
    grand: Totals = field(default_factory=Totals)
    sources: list[str] = field(default_factory=list)


def _iter_records(path: Path) -> list[UsageRecord]:
    if not path.exists():
        return []
    out: list[UsageRecord] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        out.append(UsageRecord.model_validate_json(line))
    return out


def aggregate(paths: list[Path]) -> CostReport:
    """Aggregate one or more usage.jsonl files into per-task/per-model totals."""
    report = CostReport()
    for path in paths:
        recs = _iter_records(path)
        if recs:
            report.sources.append(str(path))
        for rec in recs:
            report.by_task.setdefault(rec.task or "?", Totals()).add(rec)
            report.by_model.setdefault(rec.model or "?", Totals()).add(rec)
            report.grand.add(rec)
    return report


def estimate_cost(
    report: CostReport, prices: dict[str, dict[str, float]]
) -> dict[str, float]:
    """Per-model USD estimate using a model->{prompt,completion}/1k table."""
    out: dict[str, float] = {}
    for model, totals in report.by_model.items():
        if model not in prices:
            continue
        p = prices[model]
        out[model] = (
            totals.prompt / 1000.0 * p["prompt"]
            + totals.completion / 1000.0 * p["completion"]
        )
    return out