iomgaa
e06cd8e8b7
Twenty-five columns and not one of them answered "which tier was this?", so the question the whole issue exists to settle - does a higher tier buy anything - had no way to group its data. The three emit entry points deliberately disagree, the way sampling already does. A successful attempt records what the transport actually sent: with EFFORT_FALLBACK=nearest a request for medium goes out as low, and recomputing here would file the row under a tier that never left the process. A failed attempt has no response to read, so it falls back to the requested tier - which is exactly right for the tier errors that are rejected before any HTTP happens, because the rejected tier is the signal. Cache hits and terminal failures have no chosen source at all, so a source-level tier is not a thing they could report. emit_attempt now demands to be told whether the path reasons at all. Embedding and OCR share the emitter but never send reasoning parameters; without the flag a source that mistakenly carries ENABLE_THINKING would hang a tier on a call that could not possibly have run at one. The value lands as a plain str. StrEnum is a str subclass and asyncpg promises nothing about encoding subclasses, and a telemetry write that fails is only a warning - Postgres would just quietly lose the column. NULL means nobody declared a tier, which is not the same statement as 'none', and the two must never be folded together.
232 lines
8.4 KiB
Python
232 lines
8.4 KiB
Python
"""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 TestCachedInputTier:
|
||
"""issue #3: 供应商 prompt cache 命中部分按更低单价计费,不配则不猜折扣。"""
|
||
|
||
_CACHED = PricingTable(
|
||
{"m": ModelPrice(input_per_1m=10.0, output_per_1m=20.0, cached_input_per_1m=2.0)}
|
||
)
|
||
_PLAIN = PricingTable({"m": ModelPrice(input_per_1m=10.0, output_per_1m=20.0)})
|
||
|
||
def test_hit_is_billed_at_the_cached_rate(self):
|
||
# 1M prompt 中 600k 命中: 400k×10 + 600k×2 = 4.0 + 1.2
|
||
assert self._CACHED.cost("m", 1_000_000, 0, 600_000) == pytest.approx(5.2)
|
||
|
||
def test_without_the_tier_the_result_is_unchanged(self):
|
||
"""未配缓存档 = 退化为现状全额计价,绝不按经验折扣率猜(P5)。"""
|
||
full = self._PLAIN.cost("m", 1_000_000, 0)
|
||
assert self._PLAIN.cost("m", 1_000_000, 0, 600_000) == full == pytest.approx(10.0)
|
||
|
||
@pytest.mark.parametrize("cached", [None, 0])
|
||
def test_no_hit_is_billed_in_full(self, cached):
|
||
assert self._CACHED.cost("m", 1_000_000, 0, cached) == pytest.approx(10.0)
|
||
|
||
def test_negative_cached_is_billed_in_full(self):
|
||
"""负数命中数不得抬高成本: cost() 是公共方法,外部输入须校验后使用(P5)。"""
|
||
assert self._CACHED.cost("m", 1_000_000, 0, -500_000) == pytest.approx(10.0)
|
||
|
||
def test_cached_over_prompt_is_clamped_and_never_negative(self):
|
||
"""网关口径异常时按输入总数夹取: 全部按缓存价,不得算出负成本。"""
|
||
clamped = self._CACHED.cost("m", 1_000_000, 0, 5_000_000)
|
||
assert clamped == pytest.approx(2.0) and clamped >= 0
|
||
|
||
def test_legacy_three_arg_call_still_works(self):
|
||
"""embedding.py 的三参调用形态必须零改动可用。"""
|
||
assert self._CACHED.cost("m", 1_000_000, 0) == pytest.approx(10.0)
|
||
|
||
def test_from_file_accepts_and_validates_the_tier(self, tmp_path):
|
||
path = tmp_path / "p.json"
|
||
path.write_text(
|
||
json.dumps(
|
||
{"m": {"input_per_1m": 10.0, "output_per_1m": 20.0, "cached_input_per_1m": 2.0}}
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
assert PricingTable.from_file(path).cost("m", 1_000_000, 0, 1_000_000) == pytest.approx(2.0)
|
||
|
||
@pytest.mark.parametrize("bad", [-1.0, "x"])
|
||
def test_from_file_rejects_a_bad_tier(self, tmp_path, bad):
|
||
path = tmp_path / "bad.json"
|
||
path.write_text(
|
||
json.dumps(
|
||
{"m": {"input_per_1m": 1.0, "output_per_1m": 2.0, "cached_input_per_1m": bad}}
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
with pytest.raises(ValueError, match="cached_input_per_1m"):
|
||
PricingTable.from_file(path)
|
||
|
||
def test_legacy_price_file_without_the_tier_still_loads(self, tmp_path):
|
||
path = tmp_path / "old.json"
|
||
path.write_text(
|
||
json.dumps({"m": {"input_per_1m": 1.0, "output_per_1m": 2.0}}), encoding="utf-8"
|
||
)
|
||
assert PricingTable.from_file(path).cost("m", 1_000_000, 0, 500_000) == pytest.approx(1.0)
|
||
|
||
def test_negative_tier_rejected_on_construction(self):
|
||
with pytest.raises(ValueError):
|
||
ModelPrice(input_per_1m=1.0, output_per_1m=1.0, cached_input_per_1m=-0.1)
|
||
|
||
|
||
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, text_cap=None)
|
||
await emitter.emit_attempt(
|
||
request=_REQ,
|
||
source=_source(),
|
||
call_id="c",
|
||
latency_ms=1,
|
||
response=_resp(),
|
||
error=None,
|
||
reasoning_applies=True,
|
||
)
|
||
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, text_cap=None)
|
||
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, text_cap=None)
|
||
await emitter.emit_attempt(
|
||
request=_REQ,
|
||
source=_source(),
|
||
call_id="c",
|
||
latency_ms=1,
|
||
response=None,
|
||
error="TransientError: boom",
|
||
reasoning_applies=True,
|
||
)
|
||
assert rec.rows[0]["cost"] is None
|
||
|
||
async def test_unknown_model_none_without_blocking(self):
|
||
rec = _MemoryRecorder()
|
||
emitter = TelemetryEmitter(rec, pricing=_TABLE, text_cap=None)
|
||
await emitter.emit_attempt(
|
||
request=_REQ,
|
||
source=_source(model="mystery"),
|
||
call_id="c",
|
||
latency_ms=1,
|
||
response=_resp(model="mystery"),
|
||
error=None,
|
||
reasoning_applies=True,
|
||
)
|
||
assert rec.rows[0]["cost"] is None
|
||
|
||
async def test_no_pricing_keeps_none(self):
|
||
"""未注入价格表 = M1 现状: cost 恒 None(回归)。"""
|
||
rec = _MemoryRecorder()
|
||
emitter = TelemetryEmitter(rec, text_cap=None)
|
||
await emitter.emit_attempt(
|
||
request=_REQ,
|
||
source=_source(),
|
||
call_id="c",
|
||
latency_ms=1,
|
||
response=_resp(),
|
||
error=None,
|
||
reasoning_applies=True,
|
||
)
|
||
assert rec.rows[0]["cost"] is None
|