PolyGateway/src/polygateway/types.py

"""核心冻结类型(M1 设计 §2;最内层,禁止 import 任何实现)。

`LLMResponse` 前 11 个字段与三参考项目逐字保序——它们的测试按位置构造
fake,字段顺序即公共承诺;新增字段只增不删且必带默认值。
"""

from dataclasses import dataclass, field
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:
    """一次治理调用的统一响应(与三项目超集兼容,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
    call_id: str
    # —— 库新增(只增不删,必带默认值;迁移兼容硬约束)——
    source_name: str = ""
    cost: float | None = None
    usage_source: str = "measured"
    structured_data: Any | None = None


@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)


@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]


@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

    def __post_init__(self) -> None:
        self._validate_identity()
        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():
                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")


@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