diff --git a/src/polygateway/embedding.py b/src/polygateway/embedding.py index 1bc7fd1..804d61d 100644 --- a/src/polygateway/embedding.py +++ b/src/polygateway/embedding.py @@ -385,14 +385,21 @@ class EmbeddingClient: vectors = [_l2_normalize(v) for v in vectors] first = outcomes[0] prompt_tokens = sum(o.result.prompt_tokens for o in outcomes) - estimated = any(o.result.usage_source == "estimated" for o in outcomes) + # 三态合并优先级(解耦设计 §3.2 #10): 任一批不可得 → 整体不可得 + sources = {o.result.usage_source for o in outcomes} + if "unavailable" in sources: + merged_source = "unavailable" + elif "estimated" in sources: + merged_source = "estimated" + else: + merged_source = "measured" return EmbeddingResponse( vectors=vectors, dim=first.result.dim, model=first.source.model, provider=first.source.provider, prompt_tokens=prompt_tokens, - usage_source="estimated" if estimated else "measured", + usage_source=merged_source, latency_ms=sum(o.latency_ms for o in outcomes), call_id=first.call_id, source_name=first.source.name, @@ -400,8 +407,15 @@ class EmbeddingClient: ) def _total_cost(self, outcomes: list[_BatchOutcome]) -> float | None: + """全批成本;任一批用量不可得则整体记 NULL(解耦设计 §3.2 #11)。 + + 逐批求和会把不可得的批当 0 计入,给出一个偏低却看似有效的金额—— + 与"宁可算不出成本,也不算错成本"的不变式相悖。 + """ if self._pricing is None: return None + if any(o.result.usage_source == "unavailable" for o in outcomes): + return None costs = [self._pricing.cost(o.source.model, o.result.prompt_tokens, 0) for o in outcomes] known = [c for c in costs if c is not None] return sum(known) if known else None diff --git a/src/polygateway/transports/openai_compat.py b/src/polygateway/transports/openai_compat.py index e927bbe..52993df 100644 --- a/src/polygateway/transports/openai_compat.py +++ b/src/polygateway/transports/openai_compat.py @@ -139,11 +139,31 @@ def _strip_think(content: str) -> tuple[str, str]: def _resolve_usage(usage: dict[str, Any], source: SourceConfig) -> tuple[int, int, str]: - """usage 帧读取;缺失/非法按 est_tokens 保守兜底并标 estimated(CHS invokers.py:241)。""" + """usage 帧读取;缺失/非法记 0/0 并标 unavailable(est_tokens 解耦设计 §3.2 #3)。 + + 不再拿 `est_tokens` 兜底: 它按 CHS 定义是"最坏情形上界",拿上界当实测值 + 只会系统性高估账单;宁可把用量记成显式的"不可得"(cost 随之为 NULL), + 让缺口可被统计,也不编一个看似有效的数字。 + """ prompt, completion = usage.get("prompt_tokens"), usage.get("completion_tokens") if isinstance(prompt, int) and isinstance(completion, int) and prompt + completion > 0: return prompt, completion, "measured" - return 0, source.est_tokens, "estimated" + return 0, 0, "unavailable" + + +def _resolve_stream_usage( + sink: dict[str, Any], salvaged: bool, source: SourceConfig +) -> tuple[int, int, str]: + """流式用量口径: 打捞路径把 measured 降级为 estimated,unavailable 原样保留。 + + 前置条件不可省(解耦设计 §3.2 #4): usage 帧本就缺失时 `0/0` 会被洗成 + `estimated`,进而按 token 换算出一个假的 `0.0` 成本。 + """ + prompt, completion, usage_source = _resolve_usage(sink.get("usage") or {}, source) + if salvaged and usage_source == "measured": + # 收到 usage 帧但流被截断: 数字真实、可信度降级(M1 设计 §6) + usage_source = "estimated" + return prompt, completion, usage_source def _extract_vectors( @@ -169,11 +189,11 @@ def _extract_vectors( def _resolve_embedding_usage(data: dict[str, Any], source: SourceConfig) -> tuple[int, str]: - """usage 读取;缺失/非法按 est_tokens 保守兜底并标 estimated(与 chat 同口径)。""" + """usage 读取;缺失/非法记 0 并标 unavailable(与 chat 同口径,设计 §3.2 #3)。""" prompt = (data.get("usage") or {}).get("prompt_tokens") if isinstance(prompt, int) and prompt > 0: return prompt, "measured" - return source.est_tokens, "estimated" + return 0, "unavailable" def _parse_embedding_payload( @@ -331,9 +351,7 @@ class OpenAICompatTransport: salvaged = self._check_done(sink, content_parts, thinking_parts, source) content, thinking = self._finalize_text(content_parts, thinking_parts, profile) self._reject_empty_completion(content, source) - prompt, completion, usage_source = _resolve_usage(sink.get("usage") or {}, source) - if salvaged: - usage_source = "estimated" # 打捞路径强制 estimated(设计 §6) + prompt, completion, usage_source = _resolve_stream_usage(sink, salvaged, source) return TransportResult( content=content, thinking=thinking, diff --git a/src/polygateway/types.py b/src/polygateway/types.py index 295c0da..adb4375 100644 --- a/src/polygateway/types.py +++ b/src/polygateway/types.py @@ -284,7 +284,7 @@ class EmbeddingTransportResult: vectors: list[list[float]] dim: int prompt_tokens: int - usage_source: str # measured | estimated + usage_source: str # measured | estimated | unavailable raw: dict[str, Any] diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py index 2a6e78e..d6c9f95 100644 --- a/tests/unit/test_embedding.py +++ b/tests/unit/test_embedding.py @@ -102,12 +102,14 @@ class TestEmbedTransport: assert result.dim == 2 assert result.prompt_tokens == 5 and result.usage_source == "measured" - async def test_missing_usage_falls_back_estimated(self): + async def test_missing_usage_is_unavailable(self): + """usage 缺失不再退到 `est_tokens`(夹具填 7),与 chat 同口径记 0 + unavailable。""" + def handler(request): return httpx.Response(200, json=_ok_body([[1.0]])) result = await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c") - assert result.prompt_tokens == 7 and result.usage_source == "estimated" # est_tokens + assert result.prompt_tokens == 0 and result.usage_source == "unavailable" @pytest.mark.parametrize( ("status", "exc_type"), @@ -161,6 +163,7 @@ from polygateway.backends.memory.breaker import InMemoryGate # noqa: E402 from polygateway.backends.memory.limiter import InMemoryLimiter # noqa: E402 from polygateway.config import EmbeddingSettings # noqa: E402 from polygateway.embedding import EmbeddingClient # noqa: E402 +from polygateway.pricing import ModelPrice, PricingTable # noqa: E402 from polygateway.sources import RoundRobinSelector # noqa: E402 from polygateway.types import ( # noqa: E402 BackpressurePolicy, @@ -260,6 +263,28 @@ class TestEmbedBatching: assert resp.usage_source == "estimated" # 任一批 estimated 则整体 estimated assert resp.prompt_tokens == 2 + 9 + async def test_unavailable_batch_dominates_and_voids_cost(self): + """三态合并优先级(设计 §3.2 #10/#11): 任一批不可得 → 整体不可得且 cost NULL。 + + 改前二值合并只看 `estimated`,measured+unavailable 会误标 measured; + `_total_cost` 逐批求和还会给出一个偏低却看似有效的金额。 + """ + estimated = EmbeddingTransportResult( + vectors=[[1.0], [1.0]], dim=1, prompt_tokens=9, usage_source="estimated", raw={} + ) + unavailable = EmbeddingTransportResult( + vectors=[[1.0], [1.0]], dim=1, prompt_tokens=0, usage_source="unavailable", raw={} + ) + client, _ = _embed_client( + [_src()], + ["ok", estimated, unavailable], + batch_size=2, + pricing=PricingTable({"embed-1": ModelPrice(input_per_1m=1.0, output_per_1m=0.0)}), + ) + resp = await client.embed(["a", "b", "c", "d", "e", "f"]) + assert resp.usage_source == "unavailable" # unavailable 压过 estimated 与 measured + assert resp.cost is None + class TestEmbedPostProcess: async def test_normalize_l2(self): diff --git a/tests/unit/test_openai_compat.py b/tests/unit/test_openai_compat.py index 0f16eda..3058437 100644 --- a/tests/unit/test_openai_compat.py +++ b/tests/unit/test_openai_compat.py @@ -13,12 +13,14 @@ from polygateway.errors import ( SourceDeadError, TransientError, ) +from polygateway.middleware.telemetry import TelemetryEmitter +from polygateway.pricing import ModelPrice, PricingTable from polygateway.transports.openai_compat import ( OpenAICompatTransport, _iter_sse_deltas, _sse_data_payload, ) -from polygateway.types import SourceConfig +from polygateway.types import ChatRequest, LLMResponse, SourceConfig def _source(**overrides): @@ -71,6 +73,53 @@ async def _complete(transport, source, *, stream=True, overlay=None): ) +# 单价刻意取"输出贵于输入"的真实形态: est_tokens 兜底把整估值塞进 completion +# 时,虚高才显形(设计 §1 的 26 倍算例即此单价)。 +_PRICING = PricingTable({"qwen-max": ModelPrice(input_per_1m=1.0, output_per_1m=8.0)}) + + +class _MemoryRecorder: + def __init__(self): + self.rows = [] + + async def record_llm_call(self, **fields): + self.rows.append(fields) + + +async def _recorded_cost(result, source): + """把 transport 产物走一遍真实计费路径,返回落库的 cost。 + + `unavailable` → cost=None 的判定在 `TelemetryEmitter` 里(设计 §3.2 #5), + 直接调 `PricingTable.cost` 对 `0/0` 只会得到 `0.0`——那正是本组用例要防的 + 假金额,故断言必须穿过 emitter 而不是单测 pricing。 + """ + recorder = _MemoryRecorder() + response = LLMResponse( + content=result.content, + thinking=result.thinking, + model=source.model, + provider=source.provider, + prompt_tokens=result.prompt_tokens, + completion_tokens=result.completion_tokens, + latency_ms=1, + ttft_ms=result.ttft_ms, + max_inter_token_ms=result.max_inter_token_ms, + cache_hit=False, + call_id="cid-1", + source_name=source.name, + usage_source=result.usage_source, + ) + await TelemetryEmitter(recorder, pricing=_PRICING).emit_attempt( + request=ChatRequest(messages=[{"role": "user", "content": "hi"}]), + source=source, + call_id="cid-1", + latency_ms=1, + response=response, + error=None, + ) + return recorder.rows[0]["cost"] + + class TestSsePureFunctions: def test_data_payload_filters_noise(self): assert _sse_data_payload("") is None @@ -129,13 +178,21 @@ class TestStreamHappyPath: assert result.content == "answer" assert result.thinking == "hmm" - async def test_usage_missing_falls_back_to_est(self): + async def test_usage_missing_is_unavailable_with_null_cost(self): + """usage 帧缺失 → 0/0 + unavailable + cost NULL(设计 §3.2 #3)。 + + 改前拿 `est_tokens` 当实测并整估值塞 completion,同一条调用记成 + `0/4000` → cost 0.032(设计 §1 的 26 倍虚高)。 + """ + def handler(request): return _sse_stream(_chunk(content="ok")) - result = await _complete(_transport_for(handler), _source(tpm=1000, est_tokens=333)) - assert result.usage_source == "estimated" - assert result.prompt_tokens == 0 and result.completion_tokens == 333 + source = _source(tpm=1000, est_tokens=4000) + result = await _complete(_transport_for(handler), source) + assert result.usage_source == "unavailable" + assert result.prompt_tokens == 0 and result.completion_tokens == 0 + assert await _recorded_cost(result, source) is None class TestMissingDoneSemantics: @@ -146,12 +203,28 @@ class TestMissingDoneSemantics: with pytest.raises(TransientError, match="missing_done|truncated"): await _complete(_transport_for(self._no_done_handler), _source()) - async def test_salvage_policy_keeps_content_as_estimated(self): - result = await _complete( - _transport_for(self._no_done_handler), _source(missing_done="salvage") - ) + async def test_salvage_with_usage_frame_degrades_to_estimated(self): + """打捞且收到 usage 帧: 数字真实、可信度降级 → estimated 且照常计费。""" + source = _source(missing_done="salvage", tpm=1000, est_tokens=4000) + result = await _complete(_transport_for(self._no_done_handler), source) assert result.content == "partial" - assert result.usage_source == "estimated" # 打捞路径强制 estimated + assert result.usage_source == "estimated" + assert result.prompt_tokens == 11 and result.completion_tokens == 7 + assert await _recorded_cost(result, source) == pytest.approx( + 11 / 1_000_000 * 1.0 + 7 / 1_000_000 * 8.0 + ) + + async def test_salvage_without_usage_frame_stays_unavailable(self): + """打捞且 usage 帧缺失: 0/0 不得被洗成 estimated,否则算出假的 0.0(设计 §3.2 #4)。""" + + def handler(request): + return _sse_stream(_chunk(content="partial"), done=False) + + source = _source(missing_done="salvage", tpm=1000, est_tokens=4000) + result = await _complete(_transport_for(handler), source) + assert result.content == "partial" + assert result.usage_source == "unavailable" + assert await _recorded_cost(result, source) is None async def test_early_eof_always_transient_even_under_salvage(self): def handler(request): diff --git a/tests/unit/test_usage_source_domain.py b/tests/unit/test_usage_source_domain.py index 38fc371..366aee6 100644 --- a/tests/unit/test_usage_source_domain.py +++ b/tests/unit/test_usage_source_domain.py @@ -9,20 +9,28 @@ telemetry 三层,放进最内层内核的类型测试会让它反向依赖具体实现。 覆盖的生产点(设计 §3.2 逐处改动表的字面量产出方): -`_resolve_usage`、`_resolve_embedding_usage`、`EmbeddingClient._merge`、 -`TelemetryEmitter.emit_attempt/emit_cache_hit/emit_terminal_failure`。 +`_resolve_usage`、`_resolve_embedding_usage`、`_resolve_stream_usage`(打捞覆盖)、 +`EmbeddingClient._merge`、`EmbeddingClient.embed` 空输入短路、 +`OcrClient._emit`、`TelemetryEmitter.emit_attempt/emit_cache_hit/emit_terminal_failure`。 """ import itertools +import json +import httpx import pytest from polygateway.backends.memory.breaker import InMemoryGate from polygateway.backends.memory.limiter import InMemoryLimiter from polygateway.embedding import EmbeddingClient, _BatchOutcome from polygateway.middleware.telemetry import TelemetryEmitter +from polygateway.ocr import OcrClient from polygateway.sources import RoundRobinSelector -from polygateway.transports.openai_compat import _resolve_embedding_usage, _resolve_usage +from polygateway.transports.openai_compat import ( + OpenAICompatTransport, + _resolve_embedding_usage, + _resolve_usage, +) from polygateway.types import ( USAGE_SOURCES, BackpressurePolicy, @@ -31,6 +39,7 @@ from polygateway.types import ( EmbeddingTransportResult, GlobalLimits, LLMResponse, + OcrTextTransportResult, RetryPolicy, SourceConfig, ) @@ -89,6 +98,94 @@ def test_resolve_embedding_usage_stays_in_domain(data): assert _resolve_embedding_usage(data, _src())[1] in USAGE_SOURCES +def _sse(*frames, done): + """构造 SSE 响应;done=False 触发打捞路径(`_complete_stream` 的覆盖分支)。""" + text = "".join(f"data: {json.dumps(f)}\n\n" for f in frames) + ( + "data: [DONE]\n\n" if done else "" + ) + return httpx.Response(200, content=text.encode(), headers={"content-type": "text/event-stream"}) + + +@pytest.mark.parametrize("usage", [{"prompt_tokens": 11, "completion_tokens": 7}, None]) +async def test_salvage_override_stays_in_domain(usage): + """打捞覆盖(`openai_compat._complete_stream`)是第三个字面量产出方。""" + frames = [{"choices": [{"delta": {"content": "partial"}}]}] + if usage is not None: + frames.append({"choices": [], "usage": usage}) + transport = OpenAICompatTransport( + client_factory=lambda src: httpx.AsyncClient( + base_url=src.base_url, + transport=httpx.MockTransport(lambda request: _sse(*frames, done=False)), + ) + ) + source = SourceConfig( + name="s1", + provider="qwen", + base_url="https://gw.example/v1", + api_key="sk", + model="m", + timeout_s=10.0, + est_tokens=4000, + missing_done="salvage", + ) + result = await transport.complete( + messages=[{"role": "user", "content": "hi"}], + source=source, + stream=True, + overlay={}, + call_id="cid", + ) + assert result.usage_source in USAGE_SOURCES + + +class _ScriptedOcrTransport: + async def recognize_text(self, *, image, source, call_id): + return OcrTextTransportResult(text="LINE-1", raw={"task_type": "text"}) + + async def parse_layout(self, *, image, source, call_id): + raise NotImplementedError + + async def check_health(self, *, source): + raise NotImplementedError + + +async def test_ocr_emit_stays_in_domain(): + """`OcrClient._emit` 的字面量(ocr.py:411)同样纳入封闭性断言。 + + 值取 `measured` 是设计 §3.3 的裁决(OCR 的 0 token 属事实);此处只断言 + 落在三态内,精确取值的防回归钉在 `test_ocr_client.py`。 + """ + source = SourceConfig( + name="m1", + provider="monkey", + base_url="http://gw.example", + api_key="none", + model="monkey-ocr", + timeout_s=10.0, + ) + recorder = _MemoryRecorder() + client = OcrClient( + scope="ocr", + sources=[source], + selector=RoundRobinSelector(), + limiter=InMemoryLimiter( + scope="ocr", + sources={source.name: source}, + global_limits=GlobalLimits(max_concurrency=0, rpm=0, tpm=0), + lease_ttl_s=100.0, + ), + breaker=InMemoryGate( + config=BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0) + ), + transport=_ScriptedOcrTransport(), + retry=RetryPolicy(max_attempts=1, backoff_base_s=0.001, backoff_max_s=0.01), + backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.001), + telemetry=recorder, + ) + await client.recognize_text(b"jpg") + assert recorder.rows[0]["usage_source"] in USAGE_SOURCES + + def _merge_client(): """构造仅用于调用 `_merge` 的最小 EmbeddingClient(不发起任何调用)。""" source = _src()