core/services/usage_service.py

"""底座 ``UsageService``：LLM 调用 / Token 累计、汇总查询与阈值告警（任务 5）。

对应 design.md「8. AuditService & UsageService」与需求 12.3/12.4/12.5：

- **累计计量**（需求 12.3）：``record_call`` 记录每次 LLM 调用的提供商、模型、
  输入 / 输出 token 估计、时延与成功标志，落入内存用量账本（用量条目 schema 见
  design「Data Models」：``{ts, provider, model, in_tokens, out_tokens, latency_ms, ok}``）。
- **汇总查询**（需求 12.4）：``summary`` / ``by_provider`` / ``by_model`` 提供累计
  调用次数与 token 消耗汇总，为管理员看板供数。
- **阈值告警**（需求 12.5）：可配置 token / 调用次数 / 成本（按单价估算）阈值；
  在**接近阈值**（达到 ``warn_ratio``，默认 80%）或超过阈值时产生告警条目。

本服务为纯内存实现（无外部依赖），与 ``LLMService`` 经 ``record_call`` 接口解耦，
便于在无网络环境下单元测试。持久化（SQLite）可在后续看板任务接入，本任务聚焦
计量、汇总与告警逻辑。
"""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# 数据载体
# ---------------------------------------------------------------------------

@dataclass
class UsageRecord:
    """单次 LLM 调用的用量条目（design「Data Models」）。"""

    ts: float
    provider: str
    model: Optional[str]
    in_tokens: int
    out_tokens: int
    latency_ms: float
    ok: bool

    @property
    def total_tokens(self) -> int:
        return self.in_tokens + self.out_tokens


@dataclass
class UsageThresholds:
    """成本 / 用量阈值配置（需求 12.5）。

    任一阈值为 ``None`` 表示不启用该维度的告警。``warn_ratio`` 为「接近阈值」的
    比例（默认 0.8，即达到阈值 80% 时即开始告警）。
    """

    max_total_tokens: Optional[int] = None
    max_calls: Optional[int] = None
    max_cost: Optional[float] = None
    warn_ratio: float = 0.8

    def __post_init__(self) -> None:
        if not (0.0 < self.warn_ratio <= 1.0):
            raise ValueError("warn_ratio 必须在 (0, 1] 区间内。")


@dataclass
class UsageAlert:
    """单条阈值告警。

    - ``metric``：触发维度（``"tokens"`` / ``"calls"`` / ``"cost"``）。
    - ``level``：``"warning"``（接近阈值）或 ``"exceeded"``（已超过阈值）。
    - ``current`` / ``limit``：当前值与阈值。
    - ``message``：面向管理员的友好提示。
    """

    metric: str
    level: str
    current: float
    limit: float
    message: str


@dataclass
class UsageSummary:
    """累计用量汇总（需求 12.4）。"""

    calls: int = 0
    successful_calls: int = 0
    failed_calls: int = 0
    in_tokens: int = 0
    out_tokens: int = 0
    total_latency_ms: float = 0.0

    @property
    def total_tokens(self) -> int:
        return self.in_tokens + self.out_tokens

    @property
    def avg_latency_ms(self) -> float:
        return self.total_latency_ms / self.calls if self.calls else 0.0


class UsageService:
    """LLM 用量累计、汇总查询与阈值告警服务。"""

    def __init__(
        self,
        thresholds: Optional[UsageThresholds] = None,
        *,
        cost_per_1k_tokens: float = 0.0,
        time_func=time.time,
    ) -> None:
        """构造服务。

        参数：
        - ``thresholds``：阈值配置；``None`` 时不做告警判定。
        - ``cost_per_1k_tokens``：每千 token 的估算单价（用于成本阈值与汇总成本估算）。
        - ``time_func``：时间源，便于确定性测试。
        """
        self._thresholds = thresholds
        self._cost_per_1k = max(cost_per_1k_tokens, 0.0)
        self._time = time_func
        self._records: list[UsageRecord] = []
        self._summary = UsageSummary()
        self._by_provider: dict[str, UsageSummary] = {}
        self._by_model: dict[str, UsageSummary] = {}
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # 计量（需求 12.3）
    # ------------------------------------------------------------------
    def record_call(
        self,
        *,
        provider: str,
        model: Optional[str] = None,
        in_tokens: int = 0,
        out_tokens: int = 0,
        latency_ms: float = 0.0,
        ok: bool = True,
    ) -> UsageRecord:
        """记录单次 LLM 调用并更新累计汇总。

        返回写入的 :class:`UsageRecord`。线程安全，便于并发计量。
        """
        in_tokens = max(int(in_tokens), 0)
        out_tokens = max(int(out_tokens), 0)
        latency_ms = max(float(latency_ms), 0.0)
        record = UsageRecord(
            ts=self._time(),
            provider=provider,
            model=model,
            in_tokens=in_tokens,
            out_tokens=out_tokens,
            latency_ms=latency_ms,
            ok=bool(ok),
        )
        with self._lock:
            self._records.append(record)
            self._apply(self._summary, record)
            self._apply(self._by_provider.setdefault(provider, UsageSummary()), record)
            model_key = model or "(default)"
            self._apply(self._by_model.setdefault(model_key, UsageSummary()), record)
        return record

    @staticmethod
    def _apply(summary: UsageSummary, record: UsageRecord) -> None:
        summary.calls += 1
        if record.ok:
            summary.successful_calls += 1
        else:
            summary.failed_calls += 1
        summary.in_tokens += record.in_tokens
        summary.out_tokens += record.out_tokens
        summary.total_latency_ms += record.latency_ms

    # ------------------------------------------------------------------
    # 汇总查询（需求 12.4）
    # ------------------------------------------------------------------
    def summary(self) -> UsageSummary:
        """返回全局累计汇总（调用次数、token 消耗、平均时延等）。"""
        with self._lock:
            return UsageSummary(**vars(self._summary))

    def by_provider(self) -> dict[str, UsageSummary]:
        """返回按提供商分组的累计汇总。"""
        with self._lock:
            return {k: UsageSummary(**vars(v)) for k, v in self._by_provider.items()}

    def by_model(self) -> dict[str, UsageSummary]:
        """返回按模型分组的累计汇总。"""
        with self._lock:
            return {k: UsageSummary(**vars(v)) for k, v in self._by_model.items()}

    def records(self) -> list[UsageRecord]:
        """返回全部原始用量条目（拷贝）。"""
        with self._lock:
            return list(self._records)

    def estimated_cost(self) -> float:
        """按 ``cost_per_1k_tokens`` 估算累计成本。"""
        with self._lock:
            return self._summary.total_tokens / 1000.0 * self._cost_per_1k

    # ------------------------------------------------------------------
    # 阈值告警（需求 12.5）
    # ------------------------------------------------------------------
    def check_alerts(self) -> list[UsageAlert]:
        """对照配置阈值返回当前告警列表（接近或超过阈值时产生）。

        无阈值配置时返回空列表。同一维度仅返回严重级别更高者
        （超过阈值优先于接近阈值）。
        """
        if self._thresholds is None:
            return []

        th = self._thresholds
        with self._lock:
            total_tokens = float(self._summary.total_tokens)
            calls = float(self._summary.calls)
            cost = self._summary.total_tokens / 1000.0 * self._cost_per_1k

        alerts: list[UsageAlert] = []
        self._maybe_alert(alerts, "tokens", total_tokens, th.max_total_tokens,
                          th.warn_ratio, unit="tokens")
        self._maybe_alert(alerts, "calls", calls, th.max_calls,
                          th.warn_ratio, unit="次调用")
        if th.max_cost is not None and self._cost_per_1k > 0:
            self._maybe_alert(alerts, "cost", cost, th.max_cost,
                              th.warn_ratio, unit="成本单位")
        return alerts

    @staticmethod
    def _maybe_alert(
        alerts: list[UsageAlert],
        metric: str,
        current: float,
        limit: Optional[float],
        warn_ratio: float,
        *,
        unit: str,
    ) -> None:
        if limit is None or limit <= 0:
            return
        if current >= limit:
            alerts.append(UsageAlert(
                metric=metric,
                level="exceeded",
                current=current,
                limit=float(limit),
                message=(
                    f"{metric} 用量已超过阈值："
                    f"当前 {current:g} {unit} ≥ 上限 {limit:g} {unit}。"
                ),
            ))
        elif current >= limit * warn_ratio:
            alerts.append(UsageAlert(
                metric=metric,
                level="warning",
                current=current,
                limit=float(limit),
                message=(
                    f"{metric} 用量接近阈值："
                    f"当前 {current:g} {unit}，已达上限 {limit:g} {unit} 的 "
                    f"{current / limit * 100:.0f}%。"
                ),
            ))

    def reset(self) -> None:
        """清空所有累计用量（便于按周期重置看板统计）。"""
        with self._lock:
            self._records.clear()
            self._summary = UsageSummary()
            self._by_provider.clear()
            self._by_model.clear()


__all__ = [
    "UsageService",
    "UsageRecord",
    "UsageSummary",
    "UsageThresholds",
    "UsageAlert",
]