spb/uqo-chat

1"""LLM usage + cost tracking, monthly budget guard."""23from __future__ import annotations45from datetime import UTC, datetime6from typing import Any78from sqlalchemy import func, select910from app.core.config import get_settings11from app.db import SessionLocal12from app.llm.router import router13from app.models import LLMUsage141516async def record_usage(task: str, course: str, usage: dict[str, Any], latency_ms: int) -> None:17    async with SessionLocal() as session:18        session.add(LLMUsage(19            generation_id=usage.get("generation_id") or None, task=task, course_code=course,20            model=usage.get("model", ""), provider=usage.get("provider", ""),21            tokens_in=int(usage.get("input_tokens", 0)), tokens_out=int(usage.get("output_tokens", 0)),22            cost_usd=float(usage.get("cost_usd", 0.0)), latency_ms=latency_ms))23        await session.commit()242526async def month_cost() -> float:27    start = datetime.now(UTC).replace(day=1, hour=0, minute=0, second=0, microsecond=0,28                                      tzinfo=None)29    async with SessionLocal() as session:30        total = await session.scalar(select(func.coalesce(func.sum(LLMUsage.cost_usd), 0.0))31                                     .where(LLMUsage.created_at >= start))32    return float(total or 0.0)333435async def budget_status(budget_override: float | None = None) -> dict[str, Any]:36    budget = budget_override if budget_override is not None else get_settings().LLM_MONTHLY_BUDGET_USD37    spent = await month_cost()38    ratio = spent / budget if budget > 0 else 0.039    router.budget_exceeded = ratio >= 1.040    return {"budget_usd": budget, "spent_usd": round(spent, 4), "ratio": round(ratio, 3),41            "warning": ratio >= 0.8, "exceeded": ratio >= 1.0}424344async def daily_costs(days: int = 30) -> list[dict[str, Any]]:45    async with SessionLocal() as session:46        day = func.date(LLMUsage.created_at)47        rows = await session.execute(48            select(day.label("day"), LLMUsage.model, func.sum(LLMUsage.cost_usd),49                   func.sum(LLMUsage.tokens_in), func.sum(LLMUsage.tokens_out), func.count())50            .group_by(day, LLMUsage.model).order_by(day.desc()).limit(days * 8))51        return [{"day": str(d), "model": m, "cost_usd": round(float(c or 0), 5),52                 "tokens_in": int(ti or 0), "tokens_out": int(to or 0), "calls": int(n)}53                for d, m, c, ti, to, n in rows]545556async def costs_by_course() -> list[dict[str, Any]]:57    async with SessionLocal() as session:58        rows = await session.execute(59            select(LLMUsage.course_code, func.sum(LLMUsage.cost_usd), func.count())60            .group_by(LLMUsage.course_code))61        return [{"course": c or "—", "cost_usd": round(float(s or 0), 4), "calls": int(n)}62                for c, s, n in rows]63