ai-agent/backend/projects/costs.py

"""Aggregate Claude conversation spend per homelab project.

A conversation is attributed to a project when the project's directory name
appears in its inferred ``projectsAuto`` list (projects it created/edited files
in) *or* the manual ``projects`` list in the action-metadata sidecar — the exact
same union the viewer uses to decide which conversations to list on a project
page (``meta.projects``). Every
attributed conversation's real per-turn usage is rolled up into a per-project
total: dollar/token totals, mean/median per conversation, and a per-token-type
($ + tokens) breakdown. The project *detail* view additionally divides the total
by the project's lines of code (see ``projects.count_loc``).
"""
from __future__ import annotations

from statistics import median

import projects as projects_mod

from conversations import rates_for, usage_cost, usage_tokens, zero_usage

# Token types we break spend down by, in display order. Each carries its own
# token count and dollar cost (priced with the conversation's model).
TYPES = ("input", "output", "cacheRead", "cacheWrite")


def _project_slugs(summary: dict, meta_all: dict,
                   fallback_for=None) -> set[str]:
    """The set of project dir-names a conversation is attributed to.

    Conversations with no project of their own roll up into the root project
    (see ``projects.attributed_projects``), exactly as in ``main._conv_meta`` —
    so the lab's own spend is a card in the gallery instead of vanishing."""
    slugs = set(summary.get("projectsAuto") or [])
    sid = summary.get("sessionId") or ""
    m = meta_all.get(sid) or {}
    for p in (m.get("projects") or []):
        if p:
            slugs.add(p)
    fallback = fallback_for(summary) if fallback_for else None
    return set(projects_mod.attributed_projects(sorted(slugs), fallback))


def _type_costs(u: dict, model: str) -> dict[str, dict]:
    """Split one usage block into per-token-type {tokens, cost}.

    Mirrors ``conversations.usage_cost``: input at the input rate, output at the
    output rate, cache reads at the model's own cache-read rate, cache writes
    priced off their weighted *units* (5m ≈ 1.25×, 1h ≈ 2×) but counted in raw
    tokens.
    """
    pi, po, pc = (r / 1e6 for r in rates_for(model))
    return {
        "input": {"tokens": u["input"], "cost": u["input"] * pi},
        "output": {"tokens": u["output"], "cost": u["output"] * po},
        "cacheRead": {"tokens": u["cacheRead"], "cost": u["cacheRead"] * pc},
        "cacheWrite": {"tokens": u["cacheWriteTokens"], "cost": u["cacheWriteUnits"] * pi},
    }


def _aggregate(convs: list[dict]) -> dict:
    """Roll a project's attributed conversations into a spend summary."""
    costs: list[float] = []
    tokens: list[int] = []
    total_cost = 0.0
    total_tokens = 0
    by_type = {t: {"tokens": 0, "cost": 0.0} for t in TYPES}

    for s in convs:
        model = s.get("model") or ""
        u = {**zero_usage(), **(s.get("usage") or {})}
        c = s.get("cost")
        if c is None:
            c = usage_cost(u, model)
        tk = s.get("tokens")
        if tk is None:
            tk = usage_tokens(u)
        costs.append(c)
        tokens.append(tk)
        total_cost += c
        total_tokens += tk
        # Split per token type at each model's own rates when the run carries
        # its per-model breakdown (a mixed Opus/Fable run priced as one model
        # would misplace the money); older summaries fall back to the primary.
        parts = [(bm.get("usage") or zero_usage(), m)
                 for m, bm in (s.get("byModel") or {}).items()] or [(u, model)]
        for pu, pm in parts:
            pu = {**zero_usage(), **pu}
            for t, v in _type_costs(pu, pm).items():
                by_type[t]["tokens"] += int(v["tokens"])
                by_type[t]["cost"] += v["cost"]

    n = len(convs)
    return {
        "conversations": n,
        "cost": total_cost,
        "tokens": total_tokens,
        "meanCost": total_cost / n if n else 0.0,
        "medianCost": float(median(costs)) if costs else 0.0,
        "meanTokens": total_tokens / n if n else 0.0,
        "medianTokens": float(median(tokens)) if tokens else 0.0,
        "byType": by_type,
    }


def costs_by_project(summaries: list[tuple[str, dict]], meta_all: dict,
                     fallback_for=None) -> dict[str, dict]:
    """{project dir-name -> spend aggregate} across all conversations.

    ``summaries`` is ``store.all_summaries()`` (``[(path, summary), …]``);
    ``meta_all`` is ``meta_store.all()``. ``fallback_for(summary)`` names the
    project a project-less conversation falls back to instead of the root (see
    ``projects.attributed_projects``) — how a work session lands on its repo.
    """
    buckets: dict[str, list[dict]] = {}
    for _path, s in summaries:
        if not s.get("messages"):
            continue
        for slug in _project_slugs(s, meta_all, fallback_for):
            buckets.setdefault(slug, []).append(s)
    return {slug: _aggregate(convs) for slug, convs in buckets.items()}