diff --git a/src/polygateway/pricing.py b/src/polygateway/pricing.py index e3685ed..95425b7 100644 --- a/src/polygateway/pricing.py +++ b/src/polygateway/pricing.py @@ -22,18 +22,27 @@ if TYPE_CHECKING: @dataclass(frozen=True) class ModelPrice: - """每百万 token 的输入/输出单价(币种由使用方口径统一)。""" + """每百万 token 的输入/输出单价(币种由使用方口径统一)。 + + `cached_input_per_1m` 是可选的**缓存读取单价**(issue #3): 供应商 prompt + cache 命中的那部分输入按更低单价计费。不填即不启用——库绝不按经验折扣率 + 猜一个数(P5 严禁默认值掩盖),未填时全额按 `input_per_1m` 计。 + """ input_per_1m: float output_per_1m: float + cached_input_per_1m: float | None = None def __post_init__(self) -> None: if self.input_per_1m < 0 or self.output_per_1m < 0: raise ValueError("单价不能为负") + if self.cached_input_per_1m is not None and self.cached_input_per_1m < 0: + raise ValueError("缓存读取单价不能为负") class PricingTable: - """model → 单价 的只读表;cost() 是全库唯一换算点(经 TelemetryEmitter)。""" + """model → 单价 的只读表;cost() 有两个调用点: `TelemetryEmitter`(chat 主路径) + 与 `embedding.py` 的批量换算。""" def __init__(self, prices: Mapping[str, ModelPrice]) -> None: self._prices = dict(prices) @@ -53,21 +62,61 @@ class PricingTable: 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") + cached_raw = entry.get("cached_input_per_1m") + try: + cached = None if cached_raw is None else float(cached_raw) + except (TypeError, ValueError) as exc: + raise ValueError( + f"价格表 {p} 条目 {model!r} 的 cached_input_per_1m 必须是数字: {cached_raw!r}" + ) from exc + if cached is not None and cached < 0: + raise ValueError(f"价格表 {p} 条目 {model!r} 的 cached_input_per_1m 不能为负") prices[model] = ModelPrice( input_per_1m=float(entry["input_per_1m"]), output_per_1m=float(entry["output_per_1m"]), + cached_input_per_1m=cached, ) return cls(prices) - def cost(self, model: str, prompt_tokens: int, completion_tokens: int) -> float | None: - """换算一次调用成本;未知 model 记 None 并仅首次 warning。""" + def cost( + self, + model: str, + prompt_tokens: int, + completion_tokens: int, + cached_prompt_tokens: int | None = None, + ) -> float | None: + """换算一次调用成本;未知 model 记 None 并仅首次 warning。 + + `cached_prompt_tokens` 是供应商 prompt cache 命中的输入 token 数 + (issue #3);仅当该 model 配了 `cached_input_per_1m` 时才分段计价, + 否则全额按输入价——不猜折扣率。参数带默认值: embedding 侧的三参调用 + 形态不受影响。 + """ 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 - ) + billed_input = prompt_tokens / 1_000_000 * price.input_per_1m + if price.cached_input_per_1m is not None and cached_prompt_tokens: + cached = self._clamp_cached(model, prompt_tokens, cached_prompt_tokens) + billed_input = (prompt_tokens - cached) / 1_000_000 * price.input_per_1m + ( + cached / 1_000_000 * price.cached_input_per_1m + ) + return billed_input + completion_tokens / 1_000_000 * price.output_per_1m + + def _clamp_cached(self, model: str, prompt_tokens: int, cached: int) -> int: + """命中数按输入总数夹取: 网关口径异常不得算出负成本(每 model 只警告一次)。""" + if cached <= prompt_tokens: + return cached + key = f"{model}:cached_over_prompt" + if key not in self._warned: + self._warned.add(key) + logger.warning( + "model {!r} 上报的缓存命中 {} 超过输入总数 {},按总数夹取计价", + model, + cached, + prompt_tokens, + ) + return prompt_tokens diff --git a/tests/unit/test_pricing.py b/tests/unit/test_pricing.py index 7c85d2a..19a388e 100644 --- a/tests/unit/test_pricing.py +++ b/tests/unit/test_pricing.py @@ -56,6 +56,70 @@ class TestPricingTable: 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_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 = []