SupplyAgent/api/services/llm_billing.py

"""LLM billing query service for the admin dashboard."""

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any
from zoneinfo import ZoneInfo

from supply_infra.db.repositories.llm_usage_repo import LlmUsageRepository
from supply_infra.db.session import get_session

_SHANGHAI = ZoneInfo("Asia/Shanghai")


def _today() -> datetime:
    return datetime.now(_SHANGHAI)


def _biz_dt(day: datetime) -> str:
    return day.strftime("%Y%m%d")


def _shift_biz_dt(biz_dt: str, days: int) -> str:
    base = datetime.strptime(biz_dt, "%Y%m%d").replace(tzinfo=_SHANGHAI)
    return _biz_dt(base + timedelta(days=days))


def _format_biz_dt_label(biz_dt: str) -> str:
    return f"{biz_dt[0:4]}-{biz_dt[4:6]}-{biz_dt[6:8]}"


def _serialize_event(row: Any) -> dict[str, Any]:
    return {
        "id": row.id,
        "biz_dt": row.biz_dt,
        "provider": row.provider,
        "model": row.model,
        "agent_name": row.agent_name,
        "run_id": row.run_id,
        "iteration": row.iteration,
        "prompt_tokens": row.prompt_tokens,
        "completion_tokens": row.completion_tokens,
        "total_tokens": row.total_tokens,
        "cost_usd": float(row.cost_usd) if row.cost_usd is not None else None,
        "created_at": row.created_at.isoformat(sep=" ", timespec="seconds")
        if row.created_at
        else None,
    }


def get_billing_overview(
    *,
    start_dt: str | None = None,
    end_dt: str | None = None,
    recent_limit: int = 50,
) -> dict[str, Any]:
    """Aggregate LLM cost for a date range with agent/model breakdowns."""
    today = _today()
    today_dt = _biz_dt(today)
    end = end_dt or today_dt
    start = start_dt or _shift_biz_dt(end, -29)

    if start > end:
        start, end = end, start

    yesterday_dt = _shift_biz_dt(today_dt, -1)
    week_start = _shift_biz_dt(today_dt, -6)

    with get_session() as session:
        repo = LlmUsageRepository(session)
        daily = repo.summarize_daily(start_dt=start, end_dt=end)
        by_agent = repo.breakdown_by_agent(start_dt=start, end_dt=end)
        by_model = repo.breakdown_by_model(start_dt=start, end_dt=end)
        recent = repo.list_recent_events(start_dt=start, end_dt=end, limit=recent_limit)

        # Fill missing days so charts are continuous.
        daily_map = {item["biz_dt"]: item for item in daily}
        filled_daily: list[dict[str, Any]] = []
        cursor = start
        while cursor <= end:
            item = daily_map.get(
                cursor,
                {
                    "biz_dt": cursor,
                    "call_count": 0,
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0,
                    "total_cost_usd": 0.0,
                    "missing_cost_count": 0,
                },
            )
            filled_daily.append({**item, "label": _format_biz_dt_label(cursor)})
            cursor = _shift_biz_dt(cursor, 1)

        range_cost = sum(item["total_cost_usd"] for item in filled_daily)
        range_calls = sum(item["call_count"] for item in filled_daily)
        range_tokens = sum(item["total_tokens"] for item in filled_daily)

        today_item = daily_map.get(today_dt)
        yesterday_item = daily_map.get(yesterday_dt)
        week_cost = sum(
            item["total_cost_usd"]
            for item in filled_daily
            if week_start <= item["biz_dt"] <= today_dt
        )

        return {
            "timezone": "Asia/Shanghai",
            "currency": "USD",
            "range": {
                "start_dt": start,
                "end_dt": end,
                "total_cost_usd": range_cost,
                "call_count": range_calls,
                "total_tokens": range_tokens,
            },
            "summary": {
                "today_biz_dt": today_dt,
                "today_cost_usd": float(today_item["total_cost_usd"]) if today_item else 0.0,
                "today_call_count": int(today_item["call_count"]) if today_item else 0,
                "yesterday_cost_usd": float(yesterday_item["total_cost_usd"])
                if yesterday_item
                else 0.0,
                "week_cost_usd": week_cost,
                "range_cost_usd": range_cost,
            },
            "daily": filled_daily,
            "by_agent": by_agent,
            "by_model": by_model,
            "recent_events": [_serialize_event(row) for row in recent],
        }