From d66299210b64daffb6ece0203ddbed94af18754d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Tue, 21 Jul 2026 00:54:10 -0400 Subject: [PATCH] feat: add pricing table and cost calculation in telemetry emitter --- src/polygateway/client.py | 7 +- src/polygateway/middleware/telemetry.py | 14 ++- src/polygateway/pricing.py | 73 +++++++++++ tests/unit/test_pricing.py | 159 ++++++++++++++++++++++++ 4 files changed, 250 insertions(+), 3 deletions(-) create mode 100644 src/polygateway/pricing.py create mode 100644 tests/unit/test_pricing.py diff --git a/src/polygateway/client.py b/src/polygateway/client.py index 6e8a10f..77ec120 100644 --- a/src/polygateway/client.py +++ b/src/polygateway/client.py @@ -22,6 +22,7 @@ from polygateway.middleware.cache import CacheMW from polygateway.middleware.retry import RetryMW from polygateway.middleware.structured import StructuredMW from polygateway.middleware.telemetry import TelemetryEmitter, TelemetryMW +from polygateway.pricing import PricingTable from polygateway.providers import get_provider from polygateway.sources import LeastInflightSelector, RoundRobinSelector, SourceCooldownMemo from polygateway.transports.openai_compat import OpenAICompatTransport @@ -68,6 +69,7 @@ class GatewayClient: backpressure: BackpressurePolicy, quota_full: str = "wait", telemetry: TelemetryRecorder | None = None, + pricing: PricingTable | None = None, cache: CacheBackend | None = None, cache_namespace: str | None = None, cache_ttl_s: int | None = None, @@ -78,7 +80,7 @@ class GatewayClient: sleep: Any = asyncio.sleep, rng: Any = random.random, ) -> None: - emitter = TelemetryEmitter(telemetry) if telemetry is not None else None + emitter = TelemetryEmitter(telemetry, pricing=pricing) if telemetry is not None else None terminal = RetryMW( scope=scope, sources=sources, @@ -207,6 +209,9 @@ class GatewayClient: backpressure=settings.backpressure, quota_full=settings.quota_full, telemetry=telemetry if telemetry is not None else _build_telemetry(settings), + pricing=PricingTable.from_file(settings.pricing_path) + if settings.pricing_path is not None + else None, cache=cache if cache is not None else _build_cache(settings), cache_namespace=settings.cache_namespace, cache_ttl_s=settings.cache_ttl_s, diff --git a/src/polygateway/middleware/telemetry.py b/src/polygateway/middleware/telemetry.py index a247b22..2022b70 100644 --- a/src/polygateway/middleware/telemetry.py +++ b/src/polygateway/middleware/telemetry.py @@ -23,14 +23,16 @@ if TYPE_CHECKING: from collections.abc import Callable from polygateway.ports import CallNext, TelemetryRecorder + from polygateway.pricing import PricingTable from polygateway.types import ChatRequest, LLMResponse, SourceConfig class TelemetryEmitter: """从请求与结果组装 18 字段并写入 recorder;一切写失败降级 warning。""" - def __init__(self, recorder: TelemetryRecorder) -> None: + def __init__(self, recorder: TelemetryRecorder, *, pricing: PricingTable | None = None) -> None: self._recorder = recorder + self._pricing = pricing async def emit_attempt( self, @@ -123,6 +125,14 @@ class TelemetryEmitter: error: str | None, ) -> None: try: + # 成本换算(M2 §6): 成功行按单价换算;缓存命中 0.0(未产生新调用); + # 失败/终态行 None;未注入价格表 = 恒 None(M1 现状) + if cache_hit: + cost: float | None = 0.0 + elif error is None and model and self._pricing is not None: + cost = self._pricing.cost(model, prompt_tokens, completion_tokens) + else: + cost = None # messages 落库前多模态摘要,与缓存 key 共用同一函数(VT R12) messages_json = json.dumps(digest_messages(request.messages), ensure_ascii=False) await self._recorder.record_llm_call( @@ -143,7 +153,7 @@ class TelemetryEmitter: max_inter_token_ms=max_inter_token_ms, cache_hit=cache_hit, error=error, - cost=None, # M1 无 pricing;M2 换算后填充 + cost=cost, ) except asyncio.CancelledError: raise diff --git a/src/polygateway/pricing.py b/src/polygateway/pricing.py new file mode 100644 index 0000000..e3685ed --- /dev/null +++ b/src/polygateway/pricing.py @@ -0,0 +1,73 @@ +"""模型单价表与成本换算(M2 设计 §6;参考仓零先例,从头设计)。 + +**零内置单价**: 实验室走中转网关,计费非官方牌价;库内硬编码单价表 +必然过时并掩盖真实成本(P5 严禁默认值掩盖错误)。价格一律由使用方 +提供——JSON 文件(`PGW_PRICING_PATH`)或 dict 注入;币种由使用方全表 +统一口径,库不设币种字段(18 字段冻结)。查不到的 model → cost=None +且每 model 仅首次 warning(防日志风暴),不阻塞调用。 +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from loguru import logger + +if TYPE_CHECKING: + from collections.abc import Mapping + + +@dataclass(frozen=True) +class ModelPrice: + """每百万 token 的输入/输出单价(币种由使用方口径统一)。""" + + input_per_1m: float + output_per_1m: float + + def __post_init__(self) -> None: + if self.input_per_1m < 0 or self.output_per_1m < 0: + raise ValueError("单价不能为负") + + +class PricingTable: + """model → 单价 的只读表;cost() 是全库唯一换算点(经 TelemetryEmitter)。""" + + def __init__(self, prices: Mapping[str, ModelPrice]) -> None: + self._prices = dict(prices) + self._warned: set[str] = set() + + @classmethod + def from_file(cls, path: Path | str) -> PricingTable: + """装配 fail-loud: 文件缺失/JSON 坏/条目缺字段/负单价一律 ValueError。""" + p = Path(path) + try: + data = json.loads(p.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"价格表 {p} 读取失败: {exc}") from exc + if not isinstance(data, dict): + raise ValueError(f"价格表 {p} 顶层必须是对象: model -> 单价") + prices: dict[str, ModelPrice] = {} + for model, entry in data.items(): + if not isinstance(entry, dict) or not {"input_per_1m", "output_per_1m"} <= set(entry): + raise ValueError(f"价格表 {p} 条目 {model!r} 须含 input_per_1m 与 output_per_1m") + prices[model] = ModelPrice( + input_per_1m=float(entry["input_per_1m"]), + output_per_1m=float(entry["output_per_1m"]), + ) + return cls(prices) + + def cost(self, model: str, prompt_tokens: int, completion_tokens: int) -> float | None: + """换算一次调用成本;未知 model 记 None 并仅首次 warning。""" + price = self._prices.get(model) + if price is None: + if model not in self._warned: + self._warned.add(model) + logger.warning("pricing 表无 model {!r} 的单价,cost 记 None", model) + return None + return ( + prompt_tokens / 1_000_000 * price.input_per_1m + + completion_tokens / 1_000_000 * price.output_per_1m + ) diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py new file mode 100644 index 0000000..7c85d2a --- /dev/null +++ b/tests/unit/test_pricing.py @@ -0,0 +1,159 @@ +"""pricing.py 价格表与 TelemetryEmitter 成本换算(M2 设计 §6)。 + +零内置单价(实验室走中转网关,内置表必然过时掩盖真实成本,P5); +价格由使用方经 JSON 文件或 dict 注入;查不到 → cost=None + 每 model +仅首次 warning,不阻塞调用;计算点收敛在 Emitter(单一 helper 铁律)。 +""" + +import json + +import pytest + +from polygateway.middleware.telemetry import TelemetryEmitter +from polygateway.pricing import ModelPrice, PricingTable +from polygateway.types import ChatRequest, LLMResponse, SourceConfig + +_TABLE = PricingTable( + { + "qwen-max": ModelPrice(input_per_1m=2.4, output_per_1m=9.6), + "free-model": ModelPrice(input_per_1m=0.0, output_per_1m=0.0), + } +) + + +class TestPricingTable: + def test_cost_formula_per_million(self): + # 1M prompt + 0.5M completion: 2.4 + 4.8 + assert _TABLE.cost("qwen-max", 1_000_000, 500_000) == pytest.approx(7.2) + assert _TABLE.cost("qwen-max", 1000, 2000) == pytest.approx(0.0216) + assert _TABLE.cost("free-model", 999, 999) == 0.0 + + def test_unknown_model_returns_none(self): + assert _TABLE.cost("mystery", 100, 100) is None + + def test_from_file_roundtrip(self, tmp_path): + path = tmp_path / "prices.json" + path.write_text( + json.dumps({"m1": {"input_per_1m": 1.0, "output_per_1m": 2.0}}), encoding="utf-8" + ) + table = PricingTable.from_file(path) + assert table.cost("m1", 1_000_000, 1_000_000) == pytest.approx(3.0) + + def test_from_file_missing_or_bad_fails_loudly(self, tmp_path): + with pytest.raises(ValueError, match="价格表"): + PricingTable.from_file(tmp_path / "nope.json") + bad = tmp_path / "bad.json" + bad.write_text("not json", encoding="utf-8") + with pytest.raises(ValueError, match="价格表"): + PricingTable.from_file(bad) + malformed = tmp_path / "malformed.json" + malformed.write_text(json.dumps({"m": {"input_per_1m": 1.0}}), encoding="utf-8") + with pytest.raises(ValueError, match="output_per_1m"): + PricingTable.from_file(malformed) + + def test_negative_price_rejected(self): + with pytest.raises(ValueError): + ModelPrice(input_per_1m=-1.0, output_per_1m=0.0) + + +class _MemoryRecorder: + def __init__(self): + self.rows = [] + + async def record_llm_call(self, **fields): + self.rows.append(fields) + + +_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}]) + + +def _resp(**overrides): + base = { + "content": "ok", + "thinking": "", + "model": "qwen-max", + "provider": "qwen", + "prompt_tokens": 1_000_000, + "completion_tokens": 500_000, + "latency_ms": 30, + "ttft_ms": None, + "max_inter_token_ms": None, + "cache_hit": False, + "call_id": "cid-1", + "source_name": "s1", + "usage_source": "measured", + } + base.update(overrides) + return LLMResponse(**base) + + +def _source(model="qwen-max"): + return SourceConfig( + name="s1", + provider="qwen", + base_url="https://gw.example/v1", + api_key="sk", + model=model, + timeout_s=10.0, + ) + + +class TestEmitterCost: + async def test_success_row_costed(self): + rec = _MemoryRecorder() + emitter = TelemetryEmitter(rec, pricing=_TABLE) + await emitter.emit_attempt( + request=_REQ, + source=_source(), + call_id="c", + latency_ms=1, + response=_resp(), + error=None, + ) + assert rec.rows[0]["cost"] == pytest.approx(7.2) + + async def test_cache_hit_row_costs_zero(self): + rec = _MemoryRecorder() + emitter = TelemetryEmitter(rec, pricing=_TABLE) + await emitter.emit_cache_hit(request=_REQ, response=_resp(cache_hit=True)) + assert rec.rows[0]["cost"] == 0.0 + + async def test_failure_row_cost_none(self): + rec = _MemoryRecorder() + emitter = TelemetryEmitter(rec, pricing=_TABLE) + await emitter.emit_attempt( + request=_REQ, + source=_source(), + call_id="c", + latency_ms=1, + response=None, + error="TransientError: boom", + ) + assert rec.rows[0]["cost"] is None + + async def test_unknown_model_none_without_blocking(self): + rec = _MemoryRecorder() + emitter = TelemetryEmitter(rec, pricing=_TABLE) + await emitter.emit_attempt( + request=_REQ, + source=_source(model="mystery"), + call_id="c", + latency_ms=1, + response=_resp(model="mystery"), + error=None, + ) + assert rec.rows[0]["cost"] is None + + async def test_no_pricing_keeps_none(self): + """未注入价格表 = M1 现状: cost 恒 None(回归)。""" + rec = _MemoryRecorder() + emitter = TelemetryEmitter(rec) + await emitter.emit_attempt( + request=_REQ, + source=_source(), + call_id="c", + latency_ms=1, + response=_resp(), + error=None, + ) + assert rec.rows[0]["cost"] is None