diff --git a/src/polygateway/types.py b/src/polygateway/types.py index 4d83b16..295c0da 100644 --- a/src/polygateway/types.py +++ b/src/polygateway/types.py @@ -9,6 +9,12 @@ from typing import Any _MISSING_DONE_DOMAIN = frozenset({"retry", "salvage"}) +USAGE_SOURCES = frozenset({"measured", "estimated", "unavailable"}) +"""usage_source 值域;仅约束库内生产侧取值,不在 frozen dataclass 上做运行时校验。""" + +_EST_TOKENS_QUOTA_DIVISOR = 60 +"""未显式配置时的预扣量除数: 假定一次调用约占一秒钟的 TPM 配额份额。""" + @dataclass(frozen=True) class LLMResponse: @@ -107,6 +113,14 @@ class SourceConfig: self._validate_gates() self._validate_watchdog() + def effective_est_tokens(self) -> int: + """TPM 入场预扣量: 显式配置优先,否则按 tpm 派生(设计 §2.2)。""" + if self.est_tokens > 0: + return self.est_tokens + if self.tpm > 0: + return max(1, self.tpm // _EST_TOKENS_QUOTA_DIVISOR) + return 0 + def _validate_identity(self) -> None: for attr in ("name", "provider", "base_url", "api_key", "model"): if not getattr(self, attr).strip(): diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py index d211b53..6b97ef7 100644 --- a/tests/unit/test_types.py +++ b/tests/unit/test_types.py @@ -1,10 +1,12 @@ """types.py 冻结签名的行为测试(M1 设计 §2)。""" import dataclasses +import inspect import pytest from polygateway.types import ( + USAGE_SOURCES, BackpressurePolicy, BreakerConfig, ChatRequest, @@ -31,6 +33,18 @@ def _make_source(**overrides): return SourceConfig(**base) +def _source_without_est(tpm): + """构造"tpm > 0 且 est_tokens 未填"的源(该组合当前尚不能直接构造)。 + + `_validate_gates` 现仍强制 `tpm > 0 ⇒ est_tokens > 0`,解绑要到 T4 才做; + 此处先按合法组合构造、再绕开构造期校验写入 tpm,只为在 T1 阶段提前锁定 + 派生逻辑本身。T4 删除该约束后,本函数可整体退化为 `_make_source(tpm=...)`。 + """ + src = _make_source(tpm=0) + object.__setattr__(src, "tpm", tpm) + return src + + class TestLLMResponse: def test_eleven_legacy_fields_positional(self): """三项目 fake 的 11 参位置构造必须零改动成立(迁移兼容硬约束)。""" @@ -117,6 +131,70 @@ class TestSourceConfig: _make_source(missing_done="ignore") +class TestEffectiveEstTokens: + """TPM 入场预扣量的派生(est_tokens 解耦设计 §2.2)。""" + + def test_derives_from_tpm_scale_free(self): + """派生量随配额同比缩放: 两种配额规模的在途上限同为 60 个调用。""" + assert _source_without_est(tpm=6000).effective_est_tokens() == 100 + assert _source_without_est(tpm=600000).effective_est_tokens() == 10000 + + def test_derived_floor_is_one(self): + """极小配额下派生量不得塌到 0——0 预扣等于 TPM 闸不设防(设计 §2.2)。""" + assert _source_without_est(tpm=30).effective_est_tokens() == 1 + + def test_zero_when_tpm_gate_disabled(self): + """tpm=0 即 TPM 闸未启用,无需预扣。""" + assert _make_source().effective_est_tokens() == 0 + + def test_explicit_value_wins(self): + """显式配置是调优覆盖,优先于派生。""" + assert _make_source(tpm=6000, est_tokens=4000).effective_est_tokens() == 4000 + + def test_is_pure_sync_function(self): + """纯方法: 非协程、可重复调用、不改动自身字段(设计 §5 并发前提)。""" + assert not inspect.iscoroutinefunction(SourceConfig.effective_est_tokens) + src = _source_without_est(tpm=6000) + assert src.effective_est_tokens() == src.effective_est_tokens() == 100 + assert src.est_tokens == 0 # 派生不回写字段 + + +class TestUsageSourceDomain: + """`usage_source` 三态值域常量(设计 §3.1)。""" + + def test_domain_is_exactly_three_values(self): + assert set(USAGE_SOURCES) == {"measured", "estimated", "unavailable"} + assert isinstance(USAGE_SOURCES, frozenset) # 不可变: 调用方无法就地扩张值域 + + @pytest.mark.parametrize( + "build", + [ + lambda v: LLMResponse( + "c", "t", "m", "p", 1, 2, 3, None, None, False, "cid", usage_source=v + ), + lambda v: Usage(prompt_tokens=1, completion_tokens=2, usage_source=v), + lambda v: TransportResult( + content="c", + thinking="", + prompt_tokens=1, + completion_tokens=2, + usage_source=v, + ttft_ms=None, + max_inter_token_ms=None, + raw={}, + ), + ], + ) + def test_no_runtime_validation_on_public_dataclasses(self, build): + """越界值构造**不得**抛异常(锁定设计 §3.1 的落点裁决)。 + + 这些是运行时构造点(如 retry.py:418),裸 `ValueError` 不属 errors.py + 四分类、`RetryMW` 不捕它,会直接逃出 `chat()`——故值域只约束生产侧, + 不落在公共 frozen dataclass 的 `__post_init__` 上。 + """ + assert build("garbage").usage_source == "garbage" + + class TestResilienceConfigs: def test_retry_policy_validation(self): assert RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0) diff --git a/tests/unit/test_usage_source_domain.py b/tests/unit/test_usage_source_domain.py new file mode 100644 index 0000000..38fc371 --- /dev/null +++ b/tests/unit/test_usage_source_domain.py @@ -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