PolyGateway/src/polygateway/ocr.py

"""OcrClient: 治理化 OCR 调用(M3 设计 §5,方案 A)。

独立精简治理循环,**复用**库的算法件: `RateLimiter`/`ProviderGate` 端口与
两种后端、错误四分类、`backoff_delay` 退避公式、`SourceCooldownMemo`、
`TelemetryEmitter`(遥测单一 helper 铁律)、选源器(含 OutcomeAwareSelector
喂数)。循环与 EmbeddingClient 同构——设计 §2.A 已声明的第三份有限重复
(chat 循环的 AIMD/429 免预算/流式看门狗/缓存均不适用于 OCR)。

与 embedding 循环的有意差异(设计 §5): ① settle 恒为 0(OCR 无 token
计费,失败也不按 est 保守结算);② 无批处理外循环;③ 接
OutcomeAwareSelector 喂数(多实例 LAN 服务单机可挂,健康选源正为此设计)。
"""

from __future__ import annotations

import asyncio
import random
import time
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

from loguru import logger

from polygateway.client import _aclose_component, _telemetry_status_of
from polygateway.deadline import ensure_call_deadline, with_call_deadline
from polygateway.errors import (
    AllSourcesExhausted,
    GovernanceBackendError,
    PolyGatewayError,
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    SourceNotConfiguredError,
    TransientError,
)
from polygateway.middleware.admission import SourceAdmission, settle_and_release
from polygateway.middleware.breaker import BreakerGate
from polygateway.middleware.ratelimit import QuotaGate
from polygateway.middleware.retry import StallClock, _failure_reason, backoff_delay
from polygateway.middleware.telemetry import TelemetryEmitter, emit_terminal_once
from polygateway.ports import OutcomeAwareSelector
from polygateway.types import (
    CallStats,
    ChatRequest,
    LLMResponse,
    OcrLayoutResult,
    OcrTextResult,
    TelemetryStatus,
    Usage,
    _CallContext,
    strip_unsupported_extra_body,
    validate_caller_dimensions,
)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable, Mapping

    from polygateway.config import OcrSettings
    from polygateway.ports import (
        GateDecision,
        OcrTransport,
        Permit,
        ProviderGate,
        RateLimiter,
        SourceSelector,
        TelemetryRecorder,
    )
    from polygateway.types import (
        BackpressurePolicy,
        CallOperation,
        OcrLayoutTransportResult,
        OcrTextTransportResult,
        RetryPolicy,
        SourceConfig,
    )

_RESPONSE_TEXT_CAP = 200  # 遥测行 text 截断长度(与 embedding 口径一致)

_OcrKind = Literal["text", "layout"]


@dataclass(frozen=True)
class _FailedAttempt:
    exc: PolyGatewayError
    immediate: bool


@dataclass(frozen=True)
class _AttemptOutcome:
    result: OcrTextTransportResult | OcrLayoutTransportResult
    source: SourceConfig
    call_id: str
    latency_ms: int


class OcrClient:
    """治理化 OCR 入口: 同时实现 OcrTextPort 与 OcrLayoutPort(D9 端口族)。

    两端点打同一服务实例池,共享同一 scope 的限流/熔断/选源状态;
    与其他 scope 共享后端实例即共享全局闸(显式传入,禁止隐式全局)。
    """

    def __init__(
        self,
        *,
        scope: str,
        sources: list[SourceConfig],
        selector: SourceSelector,
        limiter: RateLimiter,
        breaker: ProviderGate,
        transport: OcrTransport,
        retry: RetryPolicy,
        backpressure: BackpressurePolicy,
        quota_full: str = "wait",
        circuit_open: str = "fail_fast",
        telemetry: TelemetryRecorder | None = None,
        text_cap: int | None = None,
        call_deadline_s: float | None = None,
        now: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
        rng: Callable[[], float] = random.random,
    ) -> None:
        # 入口即校: 装配错误当场报,不等到第一次调用才炸
        self._call_deadline_s = ensure_call_deadline(
            call_deadline_s, "OcrClient(call_deadline_s=...)"
        )
        self._scope = scope
        # MonkeyOCR 只发 multipart 表单,带 extra_body 的源必须先剥离,否则
        # 遥测会记录一个从未发出的采样参数(issue #4 决策 G)
        self._sources = strip_unsupported_extra_body(list(sources), path="OCR")
        self._selector = selector
        self._feed_health = isinstance(selector, OutcomeAwareSelector)
        self._quota = QuotaGate(limiter, scope=self._scope)
        self._breaker = BreakerGate(breaker, scope=self._scope)
        self._transport = transport
        self._retry = retry
        self._emitter = (
            TelemetryEmitter(telemetry, scope=self._scope, text_cap=text_cap) if telemetry else None
        )
        self._telemetry = telemetry
        # 限流/熔断后端在此之外只以 QuotaGate/BreakerGate 的形态存在,自持一份
        # 引用才关得到自建的 redis 客户端(设计 §3.4)
        self._limiter_backend = limiter
        self._breaker_backend = breaker
        # 所有权默认"不拥有": `__init__` 是全量注入路径,只有工厂自建时才置 True
        self._owns_transport = False
        self._owns_telemetry = False
        self._owns_limiter = False
        self._owns_breaker = False
        self._now = now
        self._sleep = sleep
        self._rng = rng
        # 准入编排三条循环共用一份(issue #14);冷却备忘由它独占
        self._admission = SourceAdmission(
            scope=self._scope,
            sources=self._sources,
            selector=selector,
            quota=self._quota,
            breaker=self._breaker,
            backpressure=backpressure,
            quota_full=quota_full,
            circuit_open=circuit_open,
            now=now,
            sleep=sleep,
            rng=rng,
        )
        self._closed = False

    # —— 公共端口(OcrTextPort / OcrLayoutPort)——

    async def recognize_text(
        self,
        image: bytes,
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
        tenant_id: str | None = None,
        meta: Mapping[str, Any] | None = None,
        call_deadline_s: float | None = None,
    ) -> OcrTextResult:
        """一次治理文本转录(/ocr/text);text 空串 = 合法"无文字"。

        `tenant_id` 与 `meta` 是调用方自定义维度,只进遥测(issue #11)。

        `call_deadline_s` 是本次调用的墙钟硬边界(issue #22): `None` = 继承装配值。
        它治理的是**等待**——到期不等于未产出,返回时刻 = 期限 + 清理耗时。
        """
        # 必须在进链路之前校验: 链路内的一切失败都被遥测层降级成 warning
        # (库铁律「遥测写失败降级不冒泡」),校验放下游等于没有校验
        dimension_tenant_id, dimensions = validate_caller_dimensions(
            tenant_id, meta, origin="recognize_text(tenant_id=..., meta=...)"
        )
        outcome, call_stats = await self._call(
            "text",
            "recognize_text",
            image,
            session_id,
            parent_call_id,
            dimension_tenant_id,
            dimensions,
            call_deadline_s,
        )
        result = outcome.result
        return OcrTextResult(
            text=result.text,
            source_name=outcome.source.name,
            usage=Usage(0, 0),
            latency_ms=outcome.latency_ms,
            call_id=outcome.call_id,
            raw=result.raw,
            call_stats=call_stats,
        )

    async def parse_layout(
        self,
        image: bytes,
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
        tenant_id: str | None = None,
        meta: Mapping[str, Any] | None = None,
        call_deadline_s: float | None = None,
    ) -> OcrLayoutResult:
        """一次治理版面解析(/parse → ZIP);elements 空 = 合法"无元素"。

        `tenant_id` 与 `meta` 是调用方自定义维度,只进遥测(issue #11)。

        `call_deadline_s` 同 `recognize_text`(issue #22): `None` = 继承装配值。
        """
        # 校验早于链路,理由同 recognize_text;origin 标明方法名以便定位入口
        dimension_tenant_id, dimensions = validate_caller_dimensions(
            tenant_id, meta, origin="parse_layout(tenant_id=..., meta=...)"
        )
        outcome, call_stats = await self._call(
            "layout",
            "parse_layout",
            image,
            session_id,
            parent_call_id,
            dimension_tenant_id,
            dimensions,
            call_deadline_s,
        )
        result = outcome.result
        return OcrLayoutResult(
            elements=result.elements,
            page_sizes=result.page_sizes,
            source_name=outcome.source.name,
            usage=Usage(0, 0),
            latency_ms=outcome.latency_ms,
            call_id=outcome.call_id,
            raw=result.raw,
            call_stats=call_stats,
        )

    async def check_health(self) -> dict[str, bool]:
        """逐源并发健康预检(R10);探测失败=False 不上抛,取消穿透。"""
        names = [s.name for s in self._sources]
        results = await asyncio.gather(
            *(self._transport.check_health(source=s) for s in self._sources)
        )
        health = dict(zip(names, results, strict=True))
        for name, ok in health.items():
            if not ok:
                logger.warning("OCR 源 {} 健康预检未通过", name)
        return health

    # —— 治理循环(与 EmbeddingClient._embed_batch 同构;设计 §2.A)——

    async def _call(
        self,
        kind: _OcrKind,
        operation: CallOperation,
        image: bytes,
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        call_deadline_s: float | None = None,
    ) -> tuple[_AttemptOutcome, CallStats]:
        if not isinstance(image, bytes):
            raise TypeError("image 必须是 bytes(路径读取/批量拼帧留业务侧,D9)")
        if not image:
            raise ValueError("image 不能为空")
        # 期限校验与 `image` 校验同列(仍在 `_CallContext` 之前、创建 awaitable 之前)
        deadline = (
            self._call_deadline_s
            if call_deadline_s is None
            else ensure_call_deadline(call_deadline_s, f"{operation}(call_deadline_s=...)")
        )
        # M1 例外: `image` 校验在 `_call` 内而非公开方法,故上下文在该校验
        # **通过之后**创建——这样设计 §3 的"校验在统计边界外"对 OCR 才成立
        context = _CallContext(now=self._now)
        try:
            return await with_call_deadline(
                self._run(
                    kind, operation, image, session_id, parent_call_id, tenant_id, meta, context
                ),
                deadline_s=deadline,
                scope=self._scope,
            )
        except PolyGatewayError as exc:
            await self._emit_terminal(
                kind, operation, image, session_id, parent_call_id, tenant_id, meta, context, exc
            )
            raise
        except asyncio.CancelledError:
            # 三条链路同一口径尽力写一条(允许 0 条);取消优先,不 shield
            await self._emit_terminal(
                kind,
                operation,
                image,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                "cancelled",
            )
            raise

    async def _run(
        self,
        kind: _OcrKind,
        operation: CallOperation,
        image: bytes,
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
    ) -> tuple[_AttemptOutcome, CallStats]:
        """选源与重试循环。

        无源的 raise 必须在本方法内(而非循环之前的调用方): 它得被 `_call` 的
        `try` 包住,否则无源终态行根本写不出来(设计 §3.5)。
        """
        if not self._sources:
            raise AllSourcesExhausted(scope=self._scope, reason="no_sources", retry_after_s=0.0)
        fails = 0
        reasons: dict[str, str] = {}
        # 只计非生产性等待(issue #8): 真实尝试由重试预算治理,不重复烧 stall 预算
        clock = StallClock(self._now)
        while True:
            picked, gate_rejections = await self._admission.pick(reasons, {})
            if picked is None:
                await self._admission.on_no_runnable(gate_rejections, reasons, clock)
                continue
            async with clock.attempting():
                outcome = await self._attempt(
                    kind,
                    operation,
                    image,
                    *picked,
                    reasons,
                    session_id,
                    parent_call_id,
                    tenant_id,
                    meta,
                    context,
                )
            if isinstance(outcome, _AttemptOutcome):
                return outcome, context.snapshot()
            fails += 1
            if fails >= self._retry.max_attempts:
                raise AllSourcesExhausted(
                    scope=self._scope,
                    reason="retry_exhausted",
                    retry_after_s=self._retry.backoff_base_s,
                    per_source_reasons=reasons,
                ) from outcome.exc
            if not outcome.immediate:
                await self._sleep(backoff_delay(self._retry, fails, outcome.exc, self._rng))

    async def _attempt(
        self,
        kind: _OcrKind,
        operation: CallOperation,
        image: bytes,
        source: SourceConfig,
        permit: Permit,
        entry: GateDecision,
        reasons: dict[str, str],
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
    ) -> _AttemptOutcome | _FailedAttempt:
        call_id = str(uuid.uuid4())
        started = self._now()
        # 四个 emit 分支(成功/终态拒绝/取消/可重试失败)都必须带调用方维度:
        # 失败行与取消行同样需要租户归属,漏掉任一分支就会写出无归属的行
        # layout 的 POST + ZIP GET 在同一次 `_invoke` 内,故这里只登记 **1** 次
        context.register_attempt()
        try:
            result = await self._invoke(kind, image, source, call_id)
            await self._record_quietly(self._breaker.record_success(entry))
            await self._record_quietly(self._quota.mark_progress())
            self._feed_outcome(source.name, ok=True)
            latency_ms = int((self._now() - started) * 1000)
            await self._emit(
                kind,
                operation,
                image,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                result,
            )
            return _AttemptOutcome(result, source, call_id, latency_ms)
        except (RequestRejectedError, ResultInvalidError) as exc:
            await self._gate_on_terminal(exc, entry)
            await self._emit(
                kind,
                operation,
                image,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                error=exc,
            )
            raise
        except asyncio.CancelledError:
            if entry.is_probe:
                await self._record_quietly(self._breaker.release_probe(entry))
            await self._emit(
                kind,
                operation,
                image,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                error="cancelled",
            )
            raise
        except (SourceDeadError, TransientError) as exc:
            dead = isinstance(exc, SourceDeadError)
            reason = _failure_reason(exc)
            reasons[source.name] = reason
            await self._record_quietly(self._breaker.record_failure(entry, reason, dead))
            self._feed_outcome(source.name, ok=False)
            await self._emit(
                kind,
                operation,
                image,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                error=exc,
            )
            return _FailedAttempt(exc, immediate=dead)
        finally:
            # OCR 的 0 token 是**事实**而非"未知"(设计 §6.3 S6): 故取消也恰恰结 0,
            # 不引入 chat/embedding 那套 settlement_known 兜底。
            await settle_and_release(permit, 0)

    async def _invoke(
        self, kind: _OcrKind, image: bytes, source: SourceConfig, call_id: str
    ) -> OcrTextTransportResult | OcrLayoutTransportResult:
        if kind == "text":
            return await self._transport.recognize_text(image=image, source=source, call_id=call_id)
        return await self._transport.parse_layout(image=image, source=source, call_id=call_id)

    # —— 辅助(与 embedding 同口径)——

    async def _gate_on_terminal(self, exc: PolyGatewayError, entry: GateDecision) -> None:
        """终态异常门控写回: 坏结果/服务的业务拒绝(有 HTTP 响应)≠ 坏服务
        → 记成功不计窗口样本;服务没响应的本地拒绝若持探针则归还。"""
        if isinstance(exc, ResultInvalidError) or exc.status_code is not None:
            await self._record_quietly(self._breaker.record_success(entry, count_attempt=False))
        elif entry.is_probe:
            await self._record_quietly(self._breaker.release_probe(entry))

    def _feed_outcome(self, source_name: str, *, ok: bool) -> None:
        """健康喂数(M2.5 口径): 真实成败喂,ResultInvalid/健康拒绝不喂;
        选源器异常吞并降级 warning(喂数失败不得影响调用)。"""
        if not self._feed_health:
            return
        try:
            self._selector.record_outcome(source_name, ok)
        except Exception as exc:  # noqa: BLE001 — 喂数侧故障降级,不冒泡
            logger.warning("OCR 健康喂数失败(不冒泡): {}", exc)

    async def _record_quietly(self, write_back: Awaitable[object]) -> None:
        try:
            await write_back
        except asyncio.CancelledError:
            raise
        except (GovernanceBackendError, SourceNotConfiguredError) as exc:
            logger.warning("OCR 治理记账写回降级(不冒泡): {}", exc)

    def _request_for(
        self,
        kind: _OcrKind,
        image: bytes,
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
    ) -> ChatRequest:
        """OCR 行的现场 ChatRequest: 占位摘要,**图像 bytes 永不入库**。

        尝试行与终态行共用同一个构造点: 占位字面量复制成两份就会漂移。
        调用方维度与上下文必须显式填回(OCR 不走 chat 洋葱),否则 OCR 行的
        维度恒为空、`logical_call_id` 恒为 NULL。
        """
        return ChatRequest(
            messages=[{"role": "user", "content": f"<ocr:{kind} image_bytes={len(image)}>"}],
            session_id=session_id,
            parent_call_id=parent_call_id,
            tenant_id=tenant_id,
            meta=meta,
            call_context=context,
        )

    async def _emit_terminal(
        self,
        kind: _OcrKind,
        operation: CallOperation,
        image: bytes,
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
        error: PolyGatewayError | str,
    ) -> None:
        """终态行: 沿用 `<ocr:{kind} image_bytes=…>` 占位,错误文本保留类名前缀。"""
        await emit_terminal_once(
            self._emitter,
            request=self._request_for(
                kind, image, session_id, parent_call_id, tenant_id, meta, context
            ),
            context=context,
            error=error,
            operation=operation,
            # metric ocr-call-success 的注册口径是按类名归组,终态行同款保留
            class_prefixed_error=True,
        )

    async def _emit(
        self,
        kind: _OcrKind,
        operation: CallOperation,
        image: bytes,
        source: SourceConfig,
        call_id: str,
        started: float,
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
        result: OcrTextTransportResult | OcrLayoutTransportResult | None = None,
        error: PolyGatewayError | str | None = None,
    ) -> None:
        """逐尝试遥测(单一 Emitter): messages 占位摘要,图像 bytes 绝不入库。"""
        if self._emitter is None:
            return
        request = self._request_for(
            kind, image, session_id, parent_call_id, tenant_id, meta, context
        )
        latency_ms = int((self._now() - started) * 1000)
        response = None
        if result is not None:
            response = LLMResponse(
                content=self._summarize(result),
                thinking="",
                model=source.model,
                provider=source.provider,
                prompt_tokens=0,
                completion_tokens=0,
                latency_ms=latency_ms,
                ttft_ms=None,
                max_inter_token_ms=None,
                cache_hit=False,
                call_id=call_id,
                source_name=source.name,
                usage_source="measured",
            )
        # 错误文本的类名前缀现由出口的显式策略参数承担(设计 §6 I7):
        # 三处各拼一遍才是下一次漂移的种子,而取消行传的是字符串,不受前缀影响
        await self._emitter.emit_attempt(
            request=request,
            source=source,
            call_id=call_id,
            latency_ms=latency_ms,
            response=response,
            error=error,
            # OCR 走 MonkeyOCR 自有端点,没有推理参数可言(理由同 embedding)
            reasoning_applies=False,
            operation=operation,
            class_prefixed_error=True,
        )

    @staticmethod
    def _summarize(result: OcrTextTransportResult | OcrLayoutTransportResult) -> str:
        if hasattr(result, "text"):
            return result.text[:_RESPONSE_TEXT_CAP]
        return f"<elements n={len(result.elements)} pages={len(result.page_sizes)}>"

    # —— 生命周期 ——

    @property
    def telemetry_status(self) -> TelemetryStatus | None:
        """遥测后端的可写状态;无遥测或注入的 recorder 不提供状态时为 None。

        判定收敛在 `_telemetry_status_of` 一处(不是三处各自探测): 三个 client
        的 `aclose` 曾各持一份逐字复制,漂移的结果就是越权关闭(设计 §3.3/§3.4)。
        """
        return _telemetry_status_of(self._telemetry)

    async def aclose(self) -> None:
        """幂等释放**自建**资源(与 EmbeddingClient 对称);注入的组件一律不碰。"""
        if self._closed:
            return
        self._closed = True
        if self._owns_transport:
            await _aclose_component(self._transport)
        if self._owns_telemetry:
            await _aclose_component(self._telemetry)
        if self._owns_limiter:
            await _aclose_component(self._limiter_backend)
        if self._owns_breaker:
            await _aclose_component(self._breaker_backend)

    async def __aenter__(self) -> OcrClient:
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        await self.aclose()

    # —— 工厂(与 EmbeddingClient 对称)——

    @classmethod
    def from_settings(
        cls,
        settings: OcrSettings,
        *,
        limiter: RateLimiter | None = None,
        breaker: ProviderGate | None = None,
        telemetry: TelemetryRecorder | None = None,
    ) -> OcrClient:
        """按配置装配;显式传入的后端实例即共享(与其他 scope 共享全局闸)。"""
        from polygateway.client import (
            _build_breaker,
            _build_limiter,
            _build_selector,
            _build_telemetry,
            _mark_owned_components,
        )
        from polygateway.transports.monkey_ocr import MonkeyOcrTransport

        gw = settings.gateway
        sources = list(gw.sources)
        # 装配防御(D9 GLM 预留档): 配了非 monkey 源必须失败,
        # 严禁静默用 MonkeyOcrTransport 打别家端点(默认值掩盖错误)
        alien = sorted({s.provider for s in sources if s.provider != "monkey"})
        if alien:
            raise ValueError(f"OCR 装配仅支持 provider=monkey(D9 其余后端预留未实现): 发现 {alien}")
        client = cls(
            scope=gw.scope,
            sources=sources,
            selector=_build_selector(gw.selector),
            limiter=limiter if limiter is not None else _build_limiter(gw, sources),
            breaker=breaker if breaker is not None else _build_breaker(gw),
            transport=MonkeyOcrTransport(),
            retry=gw.retry,
            backpressure=gw.backpressure,
            quota_full=gw.quota_full,
            circuit_open=gw.circuit_open,
            telemetry=telemetry if telemetry is not None else _build_telemetry(gw),
            # OCR 行与 chat 行写同一张 llm_calls;漏传这一条,同表内就一半受控
            # 一半不受控(issue #12)
            text_cap=gw.telemetry_text_cap,
            call_deadline_s=gw.call_deadline_s,
        )
        _mark_owned_components(client, limiter=limiter, breaker=breaker, telemetry=telemetry)
        return client

    @classmethod
    def from_env(
        cls,
        scope: str = "OCR",
        *,
        limiter: RateLimiter | None = None,
        breaker: ProviderGate | None = None,
        telemetry: TelemetryRecorder | None = None,
        env: Mapping[str, str] | None = None,
    ) -> OcrClient:
        """从 .env/环境变量装配一个 OCR scope 的 client。"""
        from polygateway.config import OcrSettings

        return cls.from_settings(
            OcrSettings.from_env(scope, env=env),
            limiter=limiter,
            breaker=breaker,
            telemetry=telemetry,
        )