PolyGateway/src/polygateway/embedding.py

"""EmbeddingClient: 治理化 embedding 调用(M2 设计 §7,方案 G2)。

独立精简治理循环,**复用**库的算法件: `RateLimiter`/`ProviderGate` 端口与
两种后端、错误四分类、`backoff_delay` 退避公式、`SourceCooldownMemo`、
`TelemetryEmitter`(遥测单一 helper 铁律)。选源/等待循环与 RetryMW 同构
——这是设计 §7.1 已声明的有限重复(chat 循环含流式/结构化/缓存分支,
强行合一才是复制);行为口径(stall 双条件、记账降级、取消穿透)与 chat
完全一致。

对参考实现的已声明裁决(设计 §7.3): async httpx;返回 list[list[float]]
(核心不依赖 numpy);normalize 开关(VT 语义,防除零 max(norm,1e-12));
必填 batch_size 批间串行;GovDoc 自研退避与 on_usage 回调、VT 同步接口
均**有意放弃**。
"""

from __future__ import annotations

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

from loguru import logger

from polygateway.client import _aclose_component, _telemetry_status_of
from polygateway.config import EmbeddingSettings
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.types import (
    ChatRequest,
    EmbeddingResponse,
    LLMResponse,
    TelemetryStatus,
    _CallContext,
    strip_unsupported_extra_body,
    validate_caller_dimensions,
)

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

    from polygateway.ports import (
        EmbeddingTransport,
        GateDecision,
        Permit,
        ProviderGate,
        RateLimiter,
        SourceSelector,
        TelemetryRecorder,
    )
    from polygateway.pricing import PricingTable
    from polygateway.types import (
        BackpressurePolicy,
        EmbeddingTransportResult,
        RetryPolicy,
        SourceConfig,
    )

_TELEMETRY_TEXT_CAP = 200  # 遥测行每条 text 截断长度(原文不整段入库,VT R12)


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


@dataclass(frozen=True)
class _BatchOutcome:
    result: EmbeddingTransportResult
    source: SourceConfig
    call_id: str
    latency_ms: int


class EmbeddingClient:
    """治理化 embedding 入口;与 GatewayClient 共享后端实例即共享全局闸。"""

    def __init__(
        self,
        *,
        scope: str,
        sources: list[SourceConfig],
        selector: SourceSelector,
        limiter: RateLimiter,
        breaker: ProviderGate,
        transport: EmbeddingTransport,
        retry: RetryPolicy,
        backpressure: BackpressurePolicy,
        quota_full: str = "wait",
        circuit_open: str = "fail_fast",
        telemetry: TelemetryRecorder | None = None,
        pricing: PricingTable | None = None,
        text_cap: int | None = None,
        batch_size: int,
        normalize: bool = False,
        expected_dim: 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:
        if batch_size < 1:
            raise ValueError("batch_size 必须 ≥ 1")
        if expected_dim is not None and expected_dim < 1:
            raise ValueError("expected_dim 必须 ≥ 1")
        # 入口即校: 装配错误当场报,不等到第一次调用才炸
        self._call_deadline_s = ensure_call_deadline(
            call_deadline_s, "EmbeddingClient(call_deadline_s=...)"
        )
        self._scope = scope
        # embed payload 硬编码 {model, input},带 extra_body 的源必须先剥离,
        # 否则遥测会记录一个从未发出的采样参数(issue #4 决策 G)
        self._sources = strip_unsupported_extra_body(list(sources), path="embedding")
        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, pricing=pricing, 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._pricing = pricing
        self._batch_size = batch_size
        self._normalize = normalize
        self._expected_dim = expected_dim
        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

    async def embed(
        self,
        texts: list[str],
        *,
        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,
    ) -> EmbeddingResponse:
        """一次治理 embedding 调用: 按 batch_size 切批,批间串行,全批合并返回。

        `tenant_id` 与 `meta` 是调用方自定义维度,只进遥测(issue #11);它们属于
        本次调用而非某一批,故每批的遥测行都带同一份维度。

        `call_deadline_s` 是本次调用的墙钟硬边界(issue #22): `None` = 继承装配值。
        **整次调用共享一份**——N 批串行跑在同一条期限内,不按批数放大 N 倍。
        `texts == []` 的早返回在期限之外(零尝试,无等待可治)。
        """
        if not isinstance(texts, list) or any(not isinstance(t, str) for t in texts):
            raise TypeError("texts 必须是 list[str](显式优于隐式,不收单条 str)")
        # 必须在切批之前校验: 洋葱/链路内的一切失败都被遥测层降级成 warning
        # (库铁律「遥测写失败降级不冒泡」),校验放下游等于没有校验——非法维度
        # 会变成静默丢失的遥测行,而调用照常发出(issue #11 §4.2)
        dimension_tenant_id, dimensions = validate_caller_dimensions(
            tenant_id, meta, origin="embed(tenant_id=..., meta=...)"
        )
        # 期限取值与校验必须在创建 awaitable **之前**(否则遗留未 await 的协程)
        deadline = (
            self._call_deadline_s
            if call_deadline_s is None
            else ensure_call_deadline(call_deadline_s, "embed(call_deadline_s=...)")
        )
        # 校验均已通过 → 进入统计边界(设计 §3.5: `texts` 类型与调用方维度校验之后)
        context = _CallContext(now=self._now)
        if not texts:
            # 合法零尝试: 返回真实统计(attempts=0),且**不写任何遥测行**
            # ——与 cache_hit 不同,不要按"遥测必录"推断它有台账行(设计 §3 M2)
            return EmbeddingResponse(
                vectors=[],
                dim=0,
                model="",
                provider="",
                prompt_tokens=0,
                usage_source="measured",
                latency_ms=0,
                call_id=str(uuid.uuid4()),
                source_name="",
                call_stats=context.snapshot(),
            )
        try:
            return await with_call_deadline(
                self._embed_all(
                    texts, session_id, parent_call_id, dimension_tenant_id, dimensions, context
                ),
                deadline_s=deadline,
                scope=self._scope,
            )
        except PolyGatewayError as exc:
            await self._emit_terminal(
                texts, session_id, parent_call_id, dimension_tenant_id, dimensions, context, exc
            )
            raise
        except asyncio.CancelledError:
            # 三条链路同一口径尽力写一条(允许 0 条);取消优先,不 shield
            await self._emit_terminal(
                texts,
                session_id,
                parent_call_id,
                dimension_tenant_id,
                dimensions,
                context,
                "cancelled",
            )
            raise

    async def _embed_all(
        self,
        texts: list[str],
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
    ) -> EmbeddingResponse:
        """切批串行执行并合并;无源的 raise 必须在本方法内——否则无源终态行写不出。"""
        if not self._sources:
            raise AllSourcesExhausted(scope=self._scope, reason="no_sources", retry_after_s=0.0)
        outcomes = []
        for start in range(0, len(texts), self._batch_size):
            outcomes.append(
                await self._embed_batch(
                    texts[start : start + self._batch_size],
                    session_id,
                    parent_call_id,
                    tenant_id,
                    meta,
                    context,
                )
            )
        # 全批共享同一上下文,故分批是实现细节而非 N 次独立逻辑调用
        return dataclasses.replace(self._merge(outcomes), call_stats=context.snapshot())

    async def _emit_terminal(
        self,
        texts: list[str],
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
        error: PolyGatewayError | str,
    ) -> None:
        """终态行的请求摘要(设计 §6 M4): 计数占位 + 第一批截断文本。

        描述的是**本次调用的整体输入**但不扩大单行正文预算: 失败批的具体文本由同
        `logical_call_id` 的 attempt 行给出,终态行不保存全量原输入。
        """
        batches = math.ceil(len(texts) / self._batch_size)
        messages = [{"role": "user", "content": f"<embed texts={len(texts)} batches={batches}>"}]
        # 与逐批行同款构造(至多 `batch_size` 条、每条 200 字符)
        messages += [
            {"role": "user", "content": t[:_TELEMETRY_TEXT_CAP]} for t in texts[: self._batch_size]
        ]
        await emit_terminal_once(
            self._emitter,
            request=ChatRequest(
                messages=messages,
                session_id=session_id,
                parent_call_id=parent_call_id,
                tenant_id=tenant_id,
                meta=meta,
                call_context=context,
            ),
            context=context,
            error=error,
            operation="embed",
        )

    # —— 治理循环(与 RetryMW 同构;设计 §7.1 已声明的有限重复)——

    async def _embed_batch(
        self,
        batch: list[str],
        session_id: str | None,
        parent_call_id: str | None,
        tenant_id: str | None,
        meta: dict[str, Any],
        context: _CallContext,
    ) -> _BatchOutcome:
        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(
                    batch, *picked, reasons, session_id, parent_call_id, tenant_id, meta, context
                )
            if isinstance(outcome, _BatchOutcome):
                return outcome
            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,
        batch: list[str],
        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,
    ) -> _BatchOutcome | _FailedBatch:
        call_id = str(uuid.uuid4())
        started = self._now()
        actual = 0
        # 局部阶段变量(与 RetryMW 同口径): 该刻库是否已算出确定结算。只服务于取消
        # 分支的兜底取值,不进任何签名; 未分类异常逃逸时仍逐字走旧的全额退还。
        settlement_known = False
        # 登记在 transport 调用**之前**(同 RetryMW): 失败与取消的尝试也真的发出去了
        context.register_attempt()
        try:
            # 裸生成时间(1.3.7 H8): 计时只包 transport 调用本身(同 RetryMW 口径),
            # 与本链路 total_latency_ms 同一只注入钟;分批累加在成功分支登记
            gen_started = self._now()
            result = await self._transport.embed(texts=batch, source=source, call_id=call_id)
            gen_ms = int((self._now() - gen_started) * 1000)
            if self._expected_dim is not None and result.dim != self._expected_dim:
                raise ResultInvalidError(
                    f"{source.name} 维度 {result.dim} 不符期望 {self._expected_dim}",
                    source_name=source.name,
                    operation="embedding",
                )
            if result.usage_source == "unavailable":
                # 与 RetryMW 同口径: 用量不可得时按入场预扣量结算(设计 §3.2 #9)
                actual = source.effective_est_tokens()
            else:
                actual = result.prompt_tokens
            # 真实 usage 恰为 0 也是已知事实, 后续取消不得改写成 est
            settlement_known = True
            context.record_generation(gen_ms, accumulate=True)
            await self._record_quietly(self._breaker.record_success(entry))
            await self._record_quietly(self._quota.mark_progress())
            latency_ms = int((self._now() - started) * 1000)
            await self._emit(
                batch,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                result,
            )
            return _BatchOutcome(result, source, call_id, latency_ms)
        except (RequestRejectedError, ResultInvalidError) as exc:
            actual, settlement_known = 0, True  # 逐字保住 1.3.5 口径(本版不改这一族记账)
            await self._gate_on_terminal(exc, entry)
            await self._emit(
                batch,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                error=exc,
            )
            raise
        except asyncio.CancelledError:
            if not settlement_known:
                # 端口已开始、结算未定: 保守保留预扣(设计 §6.3 S7)
                actual = source.effective_est_tokens()
            if entry.is_probe:
                await self._record_quietly(self._breaker.release_probe(entry))
            await self._emit(
                batch,
                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 之前: 同级 except CancelledError 接不住
            # 落在本块 await 上的取消,它直穿 finally。值与 1.3.5 逐字相同,只是算得更早。
            actual = 0 if dead else source.effective_est_tokens()
            settlement_known = True
            await self._record_quietly(self._breaker.record_failure(entry, reason, dead))
            await self._emit(
                batch,
                source,
                call_id,
                started,
                session_id,
                parent_call_id,
                tenant_id,
                meta,
                context,
                error=exc,
            )
            return _FailedBatch(exc, immediate=dead)
        finally:
            await settle_and_release(permit, actual)

    # —— 辅助 ——

    async def _gate_on_terminal(self, exc: PolyGatewayError, entry: GateDecision) -> None:
        """终态异常的门控写回(与 RetryMW 同口径): 坏结果/网关健康拒绝 ≠ 坏服务
        → 记成功;网关没响应的拒绝若持探针则归还。"""
        if isinstance(exc, ResultInvalidError) or exc.status_code is not None:
            # M2.5 §3.1: 记成功但不计失败率窗口样本(坏结果/坏请求 ≠ 坏服务)
            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))

    async def _record_quietly(self, write_back: Awaitable[object]) -> None:
        """记账侧写回降级(与 RetryMW._record_quietly 同口径,设计 §10)。"""
        try:
            await write_back
        except asyncio.CancelledError:
            raise
        except (GovernanceBackendError, SourceNotConfiguredError) as exc:
            logger.warning("embedding 治理记账写回降级(不冒泡): {}", exc)

    async def _emit(
        self,
        batch: list[str],
        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: EmbeddingTransportResult | None = None,
        error: PolyGatewayError | str | None = None,
    ) -> None:
        """逐批遥测(经同一 Emitter): messages=截断 texts、向量绝不入库。"""
        if self._emitter is None:
            return
        # 这个 ChatRequest 只为复用同一个 Emitter 而现场构造(embedding 不走 chat
        # 洋葱),故调用方维度必须在这里显式填回,否则 embed 行的维度恒为空
        request = ChatRequest(
            messages=[{"role": "user", "content": t[:_TELEMETRY_TEXT_CAP]} for t in batch],
            session_id=session_id,
            parent_call_id=parent_call_id,
            tenant_id=tenant_id,
            meta=meta,
            call_context=context,
        )
        response = None
        if result is not None:
            response = LLMResponse(
                content=f"<vectors n={len(result.vectors)} dim={result.dim}>",
                thinking="",
                model=source.model,
                provider=source.provider,
                prompt_tokens=result.prompt_tokens,
                completion_tokens=0,
                latency_ms=int((self._now() - started) * 1000),
                ttft_ms=None,
                max_inter_token_ms=None,
                cache_hit=False,
                call_id=call_id,
                source_name=source.name,
                usage_source=result.usage_source,
            )
        await self._emitter.emit_attempt(
            request=request,
            source=source,
            call_id=call_id,
            latency_ms=int((self._now() - started) * 1000),
            response=response,
            # 异常对象原样下传: 状态码/底层异常类型/网关正文在 Emitter 内定型
            error=error,
            # embedding payload 硬编码 {model, input},从不带推理参数;源上即便
            # 误配了 ENABLE_THINKING,记一个档也是替这次调用声称它没做过的事
            reasoning_applies=False,
            operation="embed",
        )

    def _merge(self, outcomes: list[_BatchOutcome]) -> EmbeddingResponse:
        """全批合并(设计 §7.3): vectors 拼接、tokens/latency 求和、保守 usage_source。"""
        vectors = [v for o in outcomes for v in o.result.vectors]
        if self._normalize:
            vectors = [_l2_normalize(v) for v in vectors]
        first = outcomes[0]
        prompt_tokens = sum(o.result.prompt_tokens for o in outcomes)
        # 三态合并优先级(解耦设计 §3.2 #10): 任一批不可得 → 整体不可得
        sources = {o.result.usage_source for o in outcomes}
        if "unavailable" in sources:
            merged_source = "unavailable"
        elif "estimated" in sources:
            merged_source = "estimated"
        else:
            merged_source = "measured"
        return EmbeddingResponse(
            vectors=vectors,
            dim=first.result.dim,
            model=first.source.model,
            provider=first.source.provider,
            prompt_tokens=prompt_tokens,
            usage_source=merged_source,
            latency_ms=sum(o.latency_ms for o in outcomes),
            call_id=first.call_id,
            source_name=first.source.name,
            cost=self._total_cost(outcomes),
        )

    def _total_cost(self, outcomes: list[_BatchOutcome]) -> float | None:
        """全批成本;任一批用量不可得则整体记 NULL(解耦设计 §3.2 #11)。

        逐批求和会把不可得的批当 0 计入,给出一个偏低却看似有效的金额——
        与"宁可算不出成本,也不算错成本"的不变式相悖。
        """
        if self._pricing is None:
            return None
        if any(o.result.usage_source == "unavailable" for o in outcomes):
            return None
        costs = [self._pricing.cost(o.source.model, o.result.prompt_tokens, 0) for o in outcomes]
        known = [c for c in costs if c is not None]
        return sum(known) if known else None

    @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:
        """幂等释放**自建**资源(与 GatewayClient 对称);注入的组件一律不碰。"""
        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) -> EmbeddingClient:
        return self

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

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

    @classmethod
    def from_settings(
        cls,
        settings: EmbeddingSettings,
        *,
        limiter: RateLimiter | None = None,
        breaker: ProviderGate | None = None,
        telemetry: TelemetryRecorder | None = None,
        registry: Mapping[str, object] | None = None,
    ) -> EmbeddingClient:
        """按配置装配;显式传入的后端实例即共享(与 chat scope 共享全局闸)。"""
        from polygateway.client import (
            _build_breaker,
            _build_limiter,
            _build_selector,
            _build_telemetry,
            _mark_owned_components,
        )
        from polygateway.pricing import PricingTable
        from polygateway.transports.openai_compat import OpenAICompatTransport

        gw = settings.gateway
        sources = list(gw.sources)
        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=OpenAICompatTransport(registry=registry),
            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),
            pricing=PricingTable.from_file(gw.pricing_path)
            if gw.pricing_path is not None
            else None,
            # embed 行与 chat 行写同一张 llm_calls;漏传这一条,同表内就一半受控
            # 一半不受控(issue #12)
            text_cap=gw.telemetry_text_cap,
            batch_size=settings.batch_size,
            normalize=settings.normalize,
            expected_dim=settings.expected_dim,
            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 = "EMBED",
        *,
        limiter: RateLimiter | None = None,
        breaker: ProviderGate | None = None,
        telemetry: TelemetryRecorder | None = None,
        registry: Mapping[str, object] | None = None,
        env: Mapping[str, str] | None = None,
    ) -> EmbeddingClient:
        """从 .env/环境变量装配一个 embedding scope 的 client。"""
        return cls.from_settings(
            EmbeddingSettings.from_env(scope, env=env),
            limiter=limiter,
            breaker=breaker,
            telemetry=telemetry,
            registry=registry,
        )


def _l2_normalize(vector: list[float]) -> list[float]:
    """L2 归一化;`max(norm, 1e-12)` 防除零(VT embedding.py:167-170 语义)。"""
    norm = max(math.sqrt(sum(x * x for x in vector)), 1e-12)
    return [x / norm for x in vector]