feat: add pricing table and cost calculation in telemetry emitter
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user