AgentLab/backend/app/providers/usage.py

"""模型用量计量：助手建图（含挑表、接着写、自查修正）和复核这两段花了多少（细案 01 做法 3）。

**为什么要。** 运行按 llm.end 事件记用量，问数据一轮里另外两块——助手建图、复核——一直不记：这一轮的花费是低报的，
评测也算不出按角色的成本。在每个调用点手写累计（生成、挑表、退回小额度重来、复核各一处）迟早漏掉一处，所以用 LangChain
的回调在模型这一层统一截：和 langchain_core.callbacks.get_usage_metadata_callback 同一个机制（一个登记过的上下文变量，
模型每次调用时把它挂成回调），按模型累计 input / output tokens 和调用次数。

**为什么不直接用 get_usage_metadata_callback。** 它是一个 with 块，块里设上下文变量、出块时按 token 复原。生成过程是一个
async generator：with 块跨着 yield 的话，变量在两次 yield 之间留在调用方的上下文里——调用方这时去发起运行，运行的任务
会继承它，运行里的调用就被重复计一遍；调用方中途不迭代了、在别的任务里收尾，复原还会报「token 不属于这个上下文」。
这里只在每一步 __anext__ 期间设上变量、交出操作之前复原（metered）：生成过程里起的任务（心跳那一层每一步一个任务）
在创建时拷走上下文，照样计到；调用方自己的代码一行都不在计量里。

**金额。** 按 providers/catalog.estimate_cost 算，和运行的 llm.end 一个口径。估不出的调用记 unpriced_calls：模型不在价格
目录里，或者用量没报全而目录价又不是 0。估不出的不当成 0 元，调用方据此写「≥」或「金额无法估算」。汇总的 cost_usd 永远是
估得出的那部分（同运行的 usage，全系统一个口径），有 unpriced_calls 时它只是下限：写金额、算平均、做估算的地方都要连着
unpriced_calls 一起看，一分都没估出来的写估不出、不写 0（评测的开跑估算、每题平均花费以前就漏了这一条，见 eval/runner.estimate、
eval/stats.summarize）。用量没报全有两种：
回复没带用量（有的网关流式不给）；读到一半就不读了——用量在流的末尾才报全（Anthropic 的输出 token 在最后一条，OpenAI
兼容系开了 stream_usage 也在最后）。生成过程读到 done、look、ask 那一行之后照样把流读完再关（api/copilot_gen._drain），
流末尾那一段用量收得到；只有模型在那一行之后还往下写超过 DRAIN_CHARS 字的闲话时才当场关流（不为闲话付钱），连接断了也一样，
这一次只拿得到已经报了的部分：照样计次数、计已报的 token，记成估不出，不悄悄少算。

**没做完的调用**（unfinished_calls）。被取消（挑表超时被 asyncio.wait_for 取消、调用方停了）、超时、流读到一半被关掉的调用
都记一次调用、标「未完成」：token 只计已经报了的（一段都没收到的记 0，当作不知道，记成估不出），调用次数照实。以前挑表超时
的那一次一点都不计，问数据一轮、评测按角色的调用次数就少一次。LangChain 的 ainvoke 被取消时不一定回调 on_llm_error
（agenerate 卡在 gather 上被取消，后面的回调走不到）：开始了、到汇总时还没有结束回调的调用（_pending 里剩下的）一并算未完成。
请求刚发出去就被拒了的（鉴权错、参数不对，一段都没收到）不算一次调用：模型没干活。

模型名按发起调用时的参数取（invocation_params 的 model / model_name，就是接入里配的模型 id，价格目录按它查）；调用参数里
没有的（以前的演示模型就没报，评测汇总里助手用量的模型因此记成 unknown），看 LangChain 照模型对象给回调的 metadata 里的
ls_model_name；再取不到看回复里的 model_name，都没有才记 UNKNOWN_MODEL。OpenAI 兼容、Anthropic 的模型在调用参数里就带着
model，演示模型（providers/mock_model）也把 model_name 报进调用参数。
"""
from __future__ import annotations

import threading
from collections.abc import AsyncIterator, Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.tracers.context import register_configure_hook

#: 取不到模型名的调用记在这个名下
UNKNOWN_MODEL = "unknown"
#: 每个模型累计的几项（unreported 用量没报全的次数、unfinished 没做完的次数）
_COUNTERS = ("input_tokens", "output_tokens", "calls", "unreported", "unfinished")


def _abandoned(error: BaseException) -> bool:
    """调用是半路放弃的（被取消、超时），不是一发出去就被拒：模型可能已经在干活、已经计费。"""
    import asyncio

    if isinstance(error, (asyncio.CancelledError, TimeoutError)):
        return True
    return any(cls.__name__.endswith(("Timeout", "TimeoutError")) for cls in type(error).__mro__)


class Meter(BaseCallbackHandler):
    """一次计量：按模型累计用量。线程安全（同步模型的回调可能在线程里来）。"""

    #: 同步模型在线程里调用时也要回调
    run_inline = True

    def __init__(self) -> None:
        super().__init__()
        self._lock = threading.Lock()
        self._pending: dict[UUID, str] = {}
        self._models: dict[str, dict[str, int]] = {}

    def _row(self, model: str) -> dict[str, int]:
        """这个模型的累计（调用方持锁）。unreported：用量没报全的次数；unfinished：没做完的次数（见模块开头）。"""
        return self._models.setdefault(model, dict.fromkeys(_COUNTERS, 0))

    # ------------------------------------------------------------------ 回调

    def _started(self, run_id: UUID, serialized: dict[str, Any] | None, kwargs: dict[str, Any]) -> None:
        params = kwargs.get("invocation_params") or {}
        model = params.get("model") or params.get("model_name") or params.get("model_id")
        if not model:
            # LangChain 照模型对象给回调的 metadata 里带着 ls_model_name（调用时传的 model，否则对象上的 model /
            # model_name）：调用参数里不报模型名的模型类也认得出
            model = (kwargs.get("metadata") or {}).get("ls_model_name")
        if not model and isinstance(serialized, dict):
            model = (serialized.get("kwargs") or {}).get("model") or (serialized.get("kwargs") or {}).get("model_name")
        with self._lock:
            self._pending[run_id] = str(model) if model else ""

    def on_chat_model_start(self, serialized: dict[str, Any], messages: Any, *, run_id: UUID,
                            **kwargs: Any) -> None:
        self._started(run_id, serialized, kwargs)

    def on_llm_start(self, serialized: dict[str, Any], prompts: Any, *, run_id: UUID, **kwargs: Any) -> None:
        self._started(run_id, serialized, kwargs)

    def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
        message = None
        try:
            message = response.generations[0][0].message
        except (AttributeError, IndexError, TypeError):
            pass
        usage = getattr(message, "usage_metadata", None) or None
        meta = getattr(message, "response_metadata", None) or {}
        with self._lock:
            model = self._pending.pop(run_id, "") or str(meta.get("model_name") or meta.get("model") or "") \
                or UNKNOWN_MODEL
            row = self._row(model)
            row["calls"] += 1
            if usage:
                row["input_tokens"] += int(usage.get("input_tokens") or 0)
                row["output_tokens"] += int(usage.get("output_tokens") or 0)
            else:
                row["unreported"] += 1

    def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
        """调用没做完：流读到一半停了（生成过程在最后一行之后读到太多闲话就关流，或者连接断了）、被取消、超时。LangChain
        照样交来已经收到的那几段。有收到东西的、被取消或超时的照计一次、记已报的 token，算用量没报全、未完成；一段都没收到
        的报错（请求刚发出去就被拒了）不计。"""
        from langchain_core.outputs import ChatGenerationChunk

        response = kwargs.get("response")
        partial = None
        try:
            first = response.generations[0][0] if response is not None and response.generations[0] else None
            partial = first.message if isinstance(first, ChatGenerationChunk) else None
        except (AttributeError, IndexError, TypeError):
            partial = None
        with self._lock:
            model = self._pending.pop(run_id, "")
            if partial is None and not _abandoned(error):
                return
            meta = getattr(partial, "response_metadata", None) or {}
            model = model or str(meta.get("model_name") or meta.get("model") or "") or UNKNOWN_MODEL
            row = self._row(model)
            row["calls"] += 1
            row["unreported"] += 1
            row["unfinished"] += 1
            reported = getattr(partial, "usage_metadata", None) or {}
            row["input_tokens"] += int(reported.get("input_tokens") or 0)
            row["output_tokens"] += int(reported.get("output_tokens") or 0)

    def _rows(self) -> dict[str, dict[str, int]]:
        """累计的一份拷贝，开始了还没结束回调的调用算成没做完的（调用方持锁）。"""
        rows = {k: dict(v) for k, v in self._models.items()}
        for model in self._pending.values():
            row = rows.setdefault(model or UNKNOWN_MODEL, dict.fromkeys(_COUNTERS, 0))
            row["calls"] += 1
            row["unreported"] += 1
            row["unfinished"] += 1
        return rows

    # ------------------------------------------------------------------ 合并、汇总

    def merge(self, other: Meter) -> None:
        """把另一次计量并进来（嵌套的计量出块时并进外层）。"""
        with other._lock:
            rows = other._rows()
        with self._lock:
            for model, row in rows.items():
                mine = self._row(model)
                for key, value in row.items():
                    mine[key] = mine.get(key, 0) + value

    def summary(self) -> dict[str, Any]:
        """{by_model: {模型: {input_tokens, output_tokens, calls, cost_usd, unpriced_calls, unfinished_calls}}, input_tokens,
        output_tokens, calls, cost_usd, unpriced_calls, unfinished_calls}。cost_usd 只是估得出的那部分；unfinished_calls 次
        没做完（token 只计了已报的）。"""
        from app.providers.catalog import estimate_cost, find_pricing

        with self._lock:
            rows = self._rows()
        by_model: dict[str, dict[str, Any]] = {}
        for model, row in sorted(rows.items()):
            pricing = find_pricing(model)
            free = pricing is not None and not any(pricing.values())
            if pricing is None:
                unpriced = row["calls"]
            else:
                unpriced = 0 if free else row["unreported"]
            by_model[model] = {"input_tokens": row["input_tokens"], "output_tokens": row["output_tokens"],
                               "calls": row["calls"],
                               "cost_usd": round(estimate_cost(model, row["input_tokens"], row["output_tokens"]), 6),
                               "unpriced_calls": unpriced, "unfinished_calls": row["unfinished"]}
        total = {key: sum(r[key] for r in by_model.values())
                 for key in ("input_tokens", "output_tokens", "calls", "unpriced_calls", "unfinished_calls")}
        return {"by_model": by_model, **total, "cost_usd": round(sum(r["cost_usd"] for r in by_model.values()), 6)}


#: 计量挂在这个上下文变量上（inheritable：模型调用时 LangChain 按它把 Meter 挂成回调，子任务拷走上下文时一并带走）
_CURRENT: ContextVar[Meter | None] = ContextVar("agentlab_usage_meter", default=None)
register_configure_hook(_CURRENT, inheritable=True)


def meter() -> Meter:
    """新的一次计量。用 metering（一段代码）或 metered（一个生成过程）把它挂上。"""
    return Meter()


def current() -> Meter | None:
    return _CURRENT.get()


@contextmanager
def metering(m: Meter) -> Iterator[Meter]:
    """这一段代码里的模型调用计进 m。只给不跨 yield 的一段用（复核那一次调用）；生成过程用 metered。

    嵌套的：出块时并进外层的那一次，外层照样看得到。复原用 set 而不是按 token reset：在别的任务里收尾也不报错。"""
    outer = _CURRENT.get()
    _CURRENT.set(m)
    try:
        yield m
    finally:
        _CURRENT.set(outer)
        if outer is not None and outer is not m:
            outer.merge(m)


async def metered(source: AsyncIterator[Any], m: Meter) -> AsyncIterator[Any]:
    """迭代一个生成过程，它的模型调用计进 m：每一步 __anext__ 期间挂上，交出操作之前复原（见模块开头）。

    开始时外层已经有一次计量的（评测把一整题的助手用量记成一份）：结束时把 m 并进去，同 metering。"""
    enclosing = _CURRENT.get()
    iterator = source.__aiter__()
    try:
        while True:
            outer = _CURRENT.get()
            _CURRENT.set(m)
            try:
                item = await iterator.__anext__()
            except StopAsyncIteration:
                return
            finally:
                _CURRENT.set(outer)
            yield item
    finally:
        if enclosing is not None and enclosing is not m:
            enclosing.merge(m)


def empty() -> dict[str, Any]:
    """一次调用都没有时的汇总（同 Meter().summary()）。"""
    return Meter().summary()


__all__ = ["Meter", "UNKNOWN_MODEL", "current", "empty", "meter", "metered", "metering"]