Files
PolyGateway/src/polygateway/types.py
T
iomgaa 89ff916bc8 feat: collect reasoning_tokens from the provider usage payload (issue #6)
Reasoning tokens are already counted inside completion_tokens, so the
cost total was never wrong -- what was missing is the attribution: how
much of a call was spent thinking rather than answering.

LLMResponse and TransportResult each gain a trailing reasoning_tokens
field, and the telemetry port grows from 21 to 22 columns with the new
column appended in both backends so fresh and migrated schemas keep the
same physical order.

None means this particular call did not report the field, not that the
source never reports it: a relay that falls back to a local tokenizer
replaces the whole usage object and drops completion_tokens_details.
Downstream checks must therefore read "in (None, 0)"; no provider was
observed reporting a literal zero.
2026-08-02 05:55:37 -04:00

427 lines
16 KiB
Python

"""核心冻结类型(M1 设计 §2;最内层,禁止 import 任何实现)。
`LLMResponse` 前 11 个字段与三参考项目逐字保序——它们的测试按位置构造
fake,字段顺序即公共承诺;新增字段只增不删且必带默认值。
"""
import dataclasses
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Any
from loguru import logger
_MISSING_DONE_DOMAIN = frozenset({"retry", "salvage"})
_PROTECTED_OVERLAY_KEYS: Mapping[str, str] = MappingProxyType(
{
"model": "会让遥测记录的 model 与实际请求分叉,成本按错单价换算",
"messages": "会同时破坏缓存 key 与遥测的 messages 口径",
"stream": "会绕过流式活性看门狗(TTFT/inter-token 超时全部失效)",
"stream_options": "会丢 usage 帧,导致成本遥测归零、TPM 闸按预扣量结算失准",
}
)
"""禁止出现在采样参数覆盖层里的键: 它们由治理层拥有,被覆盖即击穿治理。"""
USAGE_SOURCES = frozenset({"measured", "estimated", "unavailable"})
"""usage_source 值域;仅约束库内生产侧取值,不在 frozen dataclass 上做运行时校验。"""
_EST_TOKENS_QUOTA_DIVISOR = 60
"""未显式配置时的预扣量除数: 假定一次调用约占一秒钟的 TPM 配额份额。"""
def validate_request_overlay(overlay: Mapping[str, Any], *, origin: str) -> dict[str, Any]:
"""校验采样参数覆盖层并返回浅拷贝;origin 用于把错误指回配置/调用点。
两类校验缺一不可(issue #4 设计决策 B):保护键会击穿治理;不可 JSON
序列化的值会在 `CacheMW` 的降级 try **之外**抛裸 `TypeError`——那条路径
不属错误四分类、`TelemetryMW` 也不捕,结果是一行遥测都没有就崩了。
两者都在进洋葱之前收口,故抛裸 `ValueError`(调用方编程错误,不可重试)。
"""
# Phase 1: 键形态——必须先于序列化试探,否则非 str 键会因 sort_keys 的
# 比较失败被误报成"值不可序列化",把人指向错误的方向
for key in overlay:
if not isinstance(key, str):
raise ValueError(f"{origin} 的键必须是 str: {key!r}(canonical JSON 要求)")
# Phase 2: 保护键
for key, reason in _PROTECTED_OVERLAY_KEYS.items():
if key in overlay:
raise ValueError(f"{origin} 不得覆盖 {key!r}: {reason}")
# Phase 3: 值可序列化(缓存 key 与遥测列都要 json.dumps)
try:
json.dumps(dict(overlay), sort_keys=True, ensure_ascii=False)
except (TypeError, ValueError) as exc:
raise ValueError(
f"{origin} 的值必须可 JSON 序列化(如 numpy 标量请先转 float/int): {exc}"
) from exc
return dict(overlay)
def merge_sampling(extra_body: Mapping[str, Any], sampling: Mapping[str, Any]) -> dict[str, Any]:
"""合并配置级与调用级采样参数;调用级优先(issue #4 设计决策 A)。"""
return {**extra_body, **sampling}
def canonical_sampling_json(merged: Mapping[str, Any]) -> str | None:
"""缓存 key 与遥测 sampling 列共用的序列化口径;空 mapping → None。"""
if not merged:
return None
return json.dumps(dict(merged), sort_keys=True, ensure_ascii=False)
@dataclass(frozen=True)
class LLMResponse:
"""一次治理调用的统一响应(与三项目超集兼容,ARCH §5.1)。"""
content: str
thinking: str
model: str
provider: str
prompt_tokens: int
completion_tokens: int
latency_ms: int
ttft_ms: float | None
max_inter_token_ms: float | None
cache_hit: bool
"""**PolyGateway 自身响应缓存**命中(未产生网关调用);与供应商侧 prompt
cache 无关,后者见 `cached_prompt_tokens`。"""
call_id: str
# —— 库新增(只增不删,必带默认值;迁移兼容硬约束)——
source_name: str = ""
cost: float | None = None
usage_source: str = "measured"
structured_data: Any | None = None
cached_prompt_tokens: int | None = None
"""供应商 prompt cache 命中的输入 token 数(issue #3);None = 该源未上报,
与"上报了但是 0"(真实零命中)区分——两者对下游的处置不同。"""
model_reported: str | None = None
"""API 响应体里的 model 字段;None = 未上报。与 `model`(配置别名)可能
分叉——供应商把别名指向新权重时,实验复现必须认这个串。"""
reasoning_tokens: int | None = None
"""推理消耗的输出 token 数(含在 `completion_tokens` 内,故不影响成本总额,
只补归因;issue #6)。
`None` = **本次调用**未上报,**不是**"该源不上报"——中转网关在上游不返回
usage 时会用本地 tokenizer 补算并整体替换 usage 对象,把
`completion_tokens_details` 一并吃掉(findings §4c 实测同一请求 10 轮呈
6:4 双峰)。实测三家供应商在未推理时都是整个 details 缺失、无人上报 `0`,
故下游判据须为 `in (None, 0)`,写 `== 0` 的条件永远不成立。"""
@dataclass(frozen=True)
class ChatRequest:
"""洋葱内部流转的不可变请求;中间件用 dataclasses.replace 派生,禁止原地修改。"""
messages: list[dict[str, Any]]
session_id: str | None = None
parent_call_id: str | None = None
cache_salt: str | None = None
cache_namespace: str | None = None
structured: Any | None = None
stream: bool = True
overlay: dict[str, Any] = field(default_factory=dict)
sampling: Mapping[str, Any] = field(default_factory=dict)
"""调用方采样意图的快照,库内中间件**永不修改**(issue #4 设计决策 A)。
与 `overlay` 分开是因为后者会被结构化中间件注入 `response_format`,在洋葱
不同深度取值不同;缓存 key 与三个遥测入口需要一个跨层恒定的读取点,否则
同一列在不同行口径分叉。"""
@dataclass(frozen=True)
class Usage:
"""token 用量;OCR 等无计费调用填 0。"""
prompt_tokens: int
completion_tokens: int
usage_source: str = "measured"
@dataclass(frozen=True)
class SourceStats:
"""限流后端回读的单源即时指标(CHS ports.py 同款)。"""
inflight: int
rpm_used: int
tpm_used: int
@dataclass(frozen=True)
class TransportResult:
"""transport 单次原始调用的产物;治理字段由 RetryMW 补齐为 LLMResponse。"""
content: str
thinking: str
prompt_tokens: int
completion_tokens: int
usage_source: str
ttft_ms: float | None
max_inter_token_ms: float | None
raw: dict[str, Any]
# —— 可观测字段(issue #3/#6;带默认值,非 OpenAI 兼容的 transport 可不填)——
cached_prompt_tokens: int | None = None
model_reported: str | None = None
reasoning_tokens: int | None = None
@dataclass(frozen=True)
class SourceConfig:
"""单个模型源的完整配置(CHS config.py 超集;不变式在构造期报错)。
限额闸 0 表示不启用;`enable_thinking` 三态: None=不注入(模型默认)、
True=注入开启参数、False=注入关闭参数(统一 VT 与 CHS 相反的现状)。
"""
name: str
provider: str
base_url: str
api_key: str
model: str
timeout_s: float
max_concurrency: int = 0
rpm: int = 0
tpm: int = 0
est_tokens: int = 0
ttft_timeout_s: float | None = None
inter_token_timeout_s: float | None = None
enable_thinking: bool | None = None
missing_done: str = "retry"
trust_env: bool = True
extra_body: Mapping[str, Any] = field(default_factory=dict)
"""本源恒定的采样参数(如 `temperature=0`),并入请求体(issue #4)。
优先级低于调用级 overlay。注: 本字段令 `SourceConfig` 不再 hashable
(加任何 mapping 字段的固有代价,裸 dict 亦然),库内无以源作 key 的写法;
要可变副本用 `dict(source.extra_body)`,要改字段用 `dataclasses.replace`。"""
def __post_init__(self) -> None:
self._validate_identity()
self._validate_gates()
self._validate_watchdog()
self._freeze_extra_body()
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():
raise ValueError(f"SourceConfig.{attr} 不能为空")
if self.missing_done not in _MISSING_DONE_DOMAIN:
raise ValueError(
f"missing_done 必须是 {sorted(_MISSING_DONE_DOMAIN)}: {self.missing_done!r}"
)
def _validate_gates(self) -> None:
if self.timeout_s <= 0:
raise ValueError("timeout_s 必须 > 0")
for attr in ("max_concurrency", "rpm", "tpm", "est_tokens"):
if getattr(self, attr) < 0:
raise ValueError(f"SourceConfig.{attr} 不能为负(0 表示不启用)")
# 注: 不再强制 `tpm > 0 ⇒ est_tokens > 0`——预扣量由 effective_est_tokens()
# 自 tpm 派生,运维只需填供应商配额页上抄得到的 tpm(设计 §3.2 #1)
def _validate_watchdog(self) -> None:
# CHS config.py:66-82: 流式看门狗成对配置且 0 < inter < ttft < timeout_s
if (self.ttft_timeout_s is None) != (self.inter_token_timeout_s is None):
raise ValueError("ttft_timeout_s 与 inter_token_timeout_s 必须同时设置或同时缺省")
if self.ttft_timeout_s is not None and not (
0 < self.inter_token_timeout_s < self.ttft_timeout_s < self.timeout_s
):
raise ValueError("看门狗不变式要求 0 < inter_token < ttft < timeout_s")
def _freeze_extra_body(self) -> None:
"""校验后转只读视图: 装配完成的源不应再被就地改采样参数(设计决策 E)。"""
validated = validate_request_overlay(
self.extra_body, origin=f"SourceConfig({self.name}).extra_body"
)
object.__setattr__(self, "extra_body", MappingProxyType(validated))
def strip_unsupported_extra_body(sources: list[SourceConfig], *, path: str) -> list[SourceConfig]:
"""剥离非 chat 路径不消费的 `extra_body` 并 warning(issue #4 决策 G)。
剥离是必需的而非顺手清理: embedding 的 payload 硬编码 `{model, input}`、
MonkeyOCR 只发 multipart 表单,两者都不会把 `extra_body` 发出去;但遥测的
`sampling` 列会并上 `source.extra_body`,不剥离就等于**记录一个从未发出的
参数**——那是数据造假,污染的恰是事后复现的唯一依据。
选择 warning 放行而非报错: 这两条路径本无采样语义,配错的后果远轻于 chat
路径,不值得让下游整个装配起不来(2026-07-31 人类拍板)。
"""
stripped = []
for source in sources:
if source.extra_body:
logger.warning(
"{} 路径暂不支持 extra_body,源 {} 的该配置已被忽略"
"(需要 dimensions 等参数请提 issue): {}",
path,
source.name,
dict(source.extra_body),
)
source = dataclasses.replace(source, extra_body={})
stripped.append(source)
return stripped
@dataclass(frozen=True)
class RetryPolicy:
"""重试策略;max_attempts = 总尝试次数(含首次,M1 设计 §2.3 统一语义)。"""
max_attempts: int
backoff_base_s: float
backoff_max_s: float
def __post_init__(self) -> None:
if self.max_attempts < 1:
raise ValueError("max_attempts 必须 ≥ 1(含首次尝试)")
if self.backoff_base_s <= 0 or self.backoff_max_s < self.backoff_base_s:
raise ValueError("退避参数要求 0 < backoff_base_s ≤ backoff_max_s")
@dataclass(frozen=True)
class BreakerConfig:
"""熔断配置;probe_ttl_s 是半开探针租约时长(持有者死亡后自动回收)。
M2.5 双通道: fail_threshold 是连续失败通道;min_calls/fail_rate/window_s
是失败率通道(窗口样本 ≥ min_calls 且失败率 ≥ fail_rate 即开路,429 不入);
开路时长按重开次数指数递增,封顶 max_cooldown_s(设计 2026-07-21-m25)。
"""
fail_threshold: int
cooldown_s: float
probe_ttl_s: float
min_calls: int = 10
fail_rate: float = 0.6
window_s: float = 60.0
max_cooldown_s: float = 300.0
def __post_init__(self) -> None:
if self.fail_threshold < 1 or self.cooldown_s <= 0 or self.probe_ttl_s <= 0:
raise ValueError("熔断配置要求 fail_threshold ≥ 1 且 cooldown_s/probe_ttl_s > 0")
if self.min_calls < 1 or not (0.0 < self.fail_rate <= 1.0) or self.window_s <= 0:
raise ValueError("失败率通道要求 min_calls ≥ 1、0 < fail_rate ≤ 1、window_s > 0")
if self.max_cooldown_s < self.cooldown_s:
raise ValueError("max_cooldown_s 不得小于 cooldown_s(退避封顶低于初值)")
@dataclass(frozen=True)
class BackpressurePolicy:
"""背压配置;M1 仅使用 poll_interval_s,stall 判定 M2 启用。"""
stall_window_s: float
poll_interval_s: float
def __post_init__(self) -> None:
if self.stall_window_s <= 0 or self.poll_interval_s <= 0:
raise ValueError("背压参数必须 > 0")
@dataclass(frozen=True)
class GlobalLimits:
"""scope 级全局限额;0 表示该闸不启用。"""
max_concurrency: int
rpm: int
tpm: int
def __post_init__(self) -> None:
if self.max_concurrency < 0 or self.rpm < 0 or self.tpm < 0:
raise ValueError("全局限额不能为负(0 表示不启用)")
@dataclass(frozen=True)
class OcrLayoutElement:
"""版面单元(M3 设计 §3.1): 来自 `_middle.json` para_blocks 的带类型块。
type 为开放字符串(实测 table/image/text,不枚举锁死——零业务假设);
bbox 为 OCR 原生页面坐标 (x1, y1, x2, y2),几何映射留业务侧(D9)。
"""
type: str
bbox: tuple[float, float, float, float]
page_index: int
@dataclass(frozen=True)
class OcrTextResult:
"""一次治理 OCR 文本转录的统一响应(/ocr/text;M3 设计 §3.1)。
text 空串 = 合法"无文字";行过滤/去重/拼帧留业务侧(VT 迁移 §3)。
"""
text: str
source_name: str
usage: Usage # OCR 无计费: Usage(0, 0);耗时由 latency_ms 承载
latency_ms: int
call_id: str
raw: dict[str, Any]
@dataclass(frozen=True)
class OcrLayoutResult:
"""一次治理版面解析的统一响应(/parse → ZIP;M3 设计 §3.1)。
elements 空 = 合法"无元素";CHS 首表 = 首个 type=="table" 元素。
page_sizes 按 page_index 索引。
"""
elements: list[OcrLayoutElement]
page_sizes: list[tuple[float, float]]
source_name: str
usage: Usage
latency_ms: int
call_id: str
raw: dict[str, Any]
@dataclass(frozen=True)
class OcrTextTransportResult:
"""transport 单次 /ocr/text 调用产物;治理字段由 OcrClient 补齐。"""
text: str
raw: dict[str, Any]
@dataclass(frozen=True)
class OcrLayoutTransportResult:
"""transport 单次 /parse 两段调用产物;治理字段由 OcrClient 补齐。"""
elements: list[OcrLayoutElement]
page_sizes: list[tuple[float, float]]
raw: dict[str, Any]
@dataclass(frozen=True)
class EmbeddingTransportResult:
"""一次原始 embedding 调用的解析结果(M2 设计 §7.2;transport → client)。"""
vectors: list[list[float]]
dim: int
prompt_tokens: int
usage_source: str # measured | estimated | unavailable
raw: dict[str, Any]
@dataclass(frozen=True)
class EmbeddingResponse:
"""一次治理 embedding 调用的统一响应(多批合并;与输入等长保序)。"""
vectors: list[list[float]]
dim: int
model: str
provider: str
prompt_tokens: int
usage_source: str
latency_ms: int
call_id: str
source_name: str
cost: float | None = None