feat: derive TPM reservation and pin the usage_source domain
Task 1 of the est_tokens decoupling: capability only, no call site touched, so library behaviour is unchanged word for word. SourceConfig.effective_est_tokens() returns the explicit est_tokens when set, otherwise tpm // 60 floored at 1, otherwise 0 when the TPM gate is off. The divisor is scale free: any quota size yields the same in-flight ceiling of roughly sixty calls, which is what makes the default explainable where a fixed constant was not. USAGE_SOURCES lands with the two assertions the design asks for, not as a dead constant. test_usage_source_domain.py drives every production point -- _resolve_usage, _resolve_embedding_usage, _merge and the three TelemetryEmitter.emit_* helpers -- and asserts the output stays inside the domain; it is a separate file because the assertion spans transports, embedding and telemetry, and the innermost kernel test should not depend on implementations. The second assertion pins the opposite ruling: constructing LLMResponse with an out-of-domain value must not raise, since a bare ValueError at a runtime construction point falls outside the four error categories and would escape chat(). tpm > 0 with est_tokens = 0 is still rejected until Task 4, so the derivation tests build the future-legal shape through a helper that bypasses the constraint; the helper collapses back to _make_source once the constraint is gone.
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
"""`usage_source` 值域封闭: 库内所有生产点的产出恒落在 `USAGE_SOURCES` 内。
|
||||
|
||||
设计 §3.1 裁定值域**只约束生产侧**——公共 frozen dataclass 不加运行时校验
|
||||
(裸 `ValueError` 不属四分类,会逃出 `chat()`;该裁决的锁定断言在
|
||||
`test_types.py::TestUsageSourceDomain`)。因此封闭性只能由"逐个驱动生产点、
|
||||
断言其产出在三态内"来保证,本文件即该断言的载体。
|
||||
|
||||
独立成文件而非并入 `test_types.py`: 断言横跨 transports / embedding /
|
||||
telemetry 三层,放进最内层内核的类型测试会让它反向依赖具体实现。
|
||||
|
||||
覆盖的生产点(设计 §3.2 逐处改动表的字面量产出方):
|
||||
`_resolve_usage`、`_resolve_embedding_usage`、`EmbeddingClient._merge`、
|
||||
`TelemetryEmitter.emit_attempt/emit_cache_hit/emit_terminal_failure`。
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
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.sources import RoundRobinSelector
|
||||
from polygateway.transports.openai_compat import _resolve_embedding_usage, _resolve_usage
|
||||
from polygateway.types import (
|
||||
USAGE_SOURCES,
|
||||
BackpressurePolicy,
|
||||
BreakerConfig,
|
||||
ChatRequest,
|
||||
EmbeddingTransportResult,
|
||||
GlobalLimits,
|
||||
LLMResponse,
|
||||
RetryPolicy,
|
||||
SourceConfig,
|
||||
)
|
||||
|
||||
_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}])
|
||||
_DOMAIN = sorted(USAGE_SOURCES)
|
||||
|
||||
|
||||
def _src():
|
||||
return SourceConfig(
|
||||
name="s1",
|
||||
provider="p",
|
||||
base_url="https://gw.example/v1",
|
||||
api_key="sk",
|
||||
model="m",
|
||||
timeout_s=10.0,
|
||||
est_tokens=4000, # 兜底口径的历史来源: 生产点不得因它落到三态之外
|
||||
)
|
||||
|
||||
|
||||
class _MemoryRecorder:
|
||||
def __init__(self):
|
||||
self.rows = []
|
||||
|
||||
async def record_llm_call(self, **fields):
|
||||
self.rows.append(fields)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage",
|
||||
[
|
||||
{"prompt_tokens": 12, "completion_tokens": 34}, # 完整可信
|
||||
{}, # 整帧缺失
|
||||
{"prompt_tokens": 0, "completion_tokens": 0}, # 全 0(和不为正)
|
||||
{"prompt_tokens": "12", "completion_tokens": 34}, # 类型非法
|
||||
{"prompt_tokens": None, "completion_tokens": None},
|
||||
{"prompt_tokens": 12}, # 半帧
|
||||
],
|
||||
)
|
||||
def test_resolve_usage_stays_in_domain(usage):
|
||||
assert _resolve_usage(usage, _src())[2] in USAGE_SOURCES
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[
|
||||
{"usage": {"prompt_tokens": 12}},
|
||||
{},
|
||||
{"usage": None},
|
||||
{"usage": {}},
|
||||
{"usage": {"prompt_tokens": 0}},
|
||||
{"usage": {"prompt_tokens": "12"}},
|
||||
],
|
||||
)
|
||||
def test_resolve_embedding_usage_stays_in_domain(data):
|
||||
assert _resolve_embedding_usage(data, _src())[1] in USAGE_SOURCES
|
||||
|
||||
|
||||
def _merge_client():
|
||||
"""构造仅用于调用 `_merge` 的最小 EmbeddingClient(不发起任何调用)。"""
|
||||
source = _src()
|
||||
return EmbeddingClient(
|
||||
scope="embed",
|
||||
sources=[source],
|
||||
selector=RoundRobinSelector(),
|
||||
limiter=InMemoryLimiter(
|
||||
scope="embed",
|
||||
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=object(),
|
||||
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),
|
||||
batch_size=2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("first", "second"), list(itertools.product(_DOMAIN, repeat=2)))
|
||||
def test_merge_stays_in_domain(first, second):
|
||||
"""任意两批 usage_source 组合(含尚无生产者的 unavailable)合并后仍在三态内。"""
|
||||
source = _src()
|
||||
outcomes = [
|
||||
_BatchOutcome(
|
||||
result=EmbeddingTransportResult(
|
||||
vectors=[[1.0]], dim=1, prompt_tokens=1, usage_source=value, raw={}
|
||||
),
|
||||
source=source,
|
||||
call_id="c",
|
||||
latency_ms=1,
|
||||
)
|
||||
for value in (first, second)
|
||||
]
|
||||
assert _merge_client()._merge(outcomes).usage_source in USAGE_SOURCES
|
||||
|
||||
|
||||
async def test_empty_input_short_circuit_stays_in_domain():
|
||||
"""空输入短路自造响应(embedding.py:151),不经 transport 也须落在三态内。"""
|
||||
resp = await _merge_client().embed([])
|
||||
assert resp.usage_source in USAGE_SOURCES
|
||||
|
||||
|
||||
def _resp(usage_source):
|
||||
return LLMResponse(
|
||||
content="ok",
|
||||
thinking="",
|
||||
model="m",
|
||||
provider="p",
|
||||
prompt_tokens=1,
|
||||
completion_tokens=2,
|
||||
latency_ms=30,
|
||||
ttft_ms=None,
|
||||
max_inter_token_ms=None,
|
||||
cache_hit=False,
|
||||
call_id="cid",
|
||||
source_name="s1",
|
||||
usage_source=usage_source,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("emitted", _DOMAIN)
|
||||
async def test_emit_attempt_success_stays_in_domain(emitted):
|
||||
recorder = _MemoryRecorder()
|
||||
await TelemetryEmitter(recorder).emit_attempt(
|
||||
request=_REQ,
|
||||
source=_src(),
|
||||
call_id="cid",
|
||||
latency_ms=10,
|
||||
response=_resp(emitted),
|
||||
error=None,
|
||||
)
|
||||
assert recorder.rows[0]["usage_source"] in USAGE_SOURCES
|
||||
|
||||
|
||||
async def test_emit_attempt_failed_attempt_stays_in_domain():
|
||||
"""失败尝试无 response,`usage_source` 取 emitter 自己的字面量。"""
|
||||
recorder = _MemoryRecorder()
|
||||
await TelemetryEmitter(recorder).emit_attempt(
|
||||
request=_REQ,
|
||||
source=_src(),
|
||||
call_id="cid",
|
||||
latency_ms=10,
|
||||
response=None,
|
||||
error="boom",
|
||||
)
|
||||
assert recorder.rows[0]["usage_source"] in USAGE_SOURCES
|
||||
|
||||
|
||||
@pytest.mark.parametrize("emitted", _DOMAIN)
|
||||
async def test_emit_cache_hit_stays_in_domain(emitted):
|
||||
recorder = _MemoryRecorder()
|
||||
await TelemetryEmitter(recorder).emit_cache_hit(request=_REQ, response=_resp(emitted))
|
||||
assert recorder.rows[0]["usage_source"] in USAGE_SOURCES
|
||||
|
||||
|
||||
async def test_emit_terminal_failure_stays_in_domain():
|
||||
"""终态失败无具体源,`usage_source` 同样取 emitter 字面量。"""
|
||||
recorder = _MemoryRecorder()
|
||||
await TelemetryEmitter(recorder).emit_terminal_failure(
|
||||
request=_REQ, call_id="cid", latency_ms=10, error="cancelled"
|
||||
)
|
||||
assert recorder.rows[0]["usage_source"] in USAGE_SOURCES
|
||||
Reference in New Issue
Block a user