mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 05:51:57 +00:00
ThearchyHelios
eddeded2a4
- 第一批（安全加固与生命周期止血）：隔离测试宿主机密钥与 .env 泄漏；修复 SQLite 取消查询挂死、:memory: 连接池报错与 EngineCache 世代竞态；补齐运行与评测删除时的 Checkpointer 线程、Artifact 磁盘文件与无主附件回收；加固三套沙箱路径越界、进程组清理与 Seatbelt 数据目录隔离，封堵 SSRF IPv4-mapped IPv6 元数据绕过与响应体 OOM；修复推理模型多轮 reasoning_content、结构化输出 Token 漏计与缓存折扣计费 - 第二批（数据库/集群/问数提速与首屏分包）：补齐核心表索引与 SQLite WAL 调优，列表接口按需 defer 大列，收敛启动扫描并打通跨进程 slot 槽位唤醒；优化跨源全局选表、哈希桶数字绑定、流式 CSV 导入及多方言分析算子；剥离首屏静态引入的 @xyflow/react 与全局样式（首屏主包降至 461.59 kB），补充流式 rAF 攒批、Store 精细订阅与表单无障碍绑定 - 第三批（架构演进与工程自动化闭环）：记录库超挂接阈值历史版本自动增量归档至单库多表（解除 SQLite 10 库 ATTACH 硬上限）；重构知识库为词面 BM25 与全集合归一化向量矩阵真·双路独立召回并支持游标分批重建；升级多方言感知 SQL 模板词法分析（E'...' / $$...$$ / q'[...]' / # / 嵌套注释）与 Decimal / bool / NaN 参数校验，消除 ledger 层对 api 路由层的反向依赖；拆分 3101 行 CaliberDrafts 页面为 7 个领域子模块，消除检查脚本本机路径硬编码并新增 GitHub Actions CI 门禁与三批回归测试套件
269 lines
15 KiB
Python
269 lines
15 KiB
Python
"""模型用量计量：助手建图（含挑表、接着写、自查修正）和复核这两段花了多少（细案 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", "cache_read_tokens", "cache_creation_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)
|
||
details = usage.get("input_token_details") or {}
|
||
if isinstance(details, dict):
|
||
row["cache_read_tokens"] += int(details.get("cache_read") or 0)
|
||
row["cache_creation_tokens"] += int(details.get("cache_creation") 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)
|
||
details = reported.get("input_token_details") or {}
|
||
if isinstance(details, dict):
|
||
row["cache_read_tokens"] += int(details.get("cache_read") or 0)
|
||
row["cache_creation_tokens"] += int(details.get("cache_creation") 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"]
|
||
cache_read = int(row.get("cache_read_tokens") or 0)
|
||
cache_create = int(row.get("cache_creation_tokens") or 0)
|
||
if cache_read > 0 or cache_create > 0:
|
||
try:
|
||
raw_cost = estimate_cost(
|
||
model, row["input_tokens"], row["output_tokens"],
|
||
cache_read_tokens=cache_read, cache_creation_tokens=cache_create,
|
||
)
|
||
except TypeError:
|
||
raw_cost = estimate_cost(model, row["input_tokens"], row["output_tokens"])
|
||
else:
|
||
raw_cost = estimate_cost(model, row["input_tokens"], row["output_tokens"])
|
||
by_model[model] = {"input_tokens": row["input_tokens"], "output_tokens": row["output_tokens"],
|
||
"calls": row["calls"],
|
||
"cost_usd": round(raw_cost, 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"]
|