iomgaa/PolyGateway
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases 9 Wiki Activity

fix: stop passing est_tokens off as measured usage

Browse Source 
usage 帧缺失/非法时不再拿 est_tokens(最坏情形上界)当实测值,chat 与
embedding 两处兜底改记 0 并标 unavailable;打捞覆盖加 measured 前置条件,
避免 0/0 被洗成 estimated 而算出假的 0.0。embedding 全批合并扩三态(任一批
不可得 → 整体不可得),_total_cost 遇不可得批整体记 NULL。
This commit is contained in:
iomgaa
2026-07-30 10:39:32 -04:00
parent 42e429eb58
commit 195454d2e3
6 changed files with 252 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/polygateway/embedding.py
+16 -2
View File
@@ -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
src/polygateway/transports/openai_compat.py
+25 -7
View File
@@ -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,
src/polygateway/types.py
+1 -1
View File
@@ -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]
tests/unit/test_embedding.py
+27 -2
View File
@@ -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):
tests/unit/test_openai_compat.py
+83 -10
View File
@@ -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):
tests/unit/test_usage_source_domain.py
+100 -3
View File
@@ -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()