PolyGateway/src/polygateway/middleware/telemetry.py

"""TelemetryMW + TelemetryEmitter: 遥测调用点收敛为单一 helper(铁律)。

Emitter 是全库**唯一**调用 `record_llm_call` 的地方(三项目 4 处逐字复制
15 参调用的教训)。分工: RetryMW 经 Emitter 逐次记录每次尝试;TelemetryMW
(最外层)只记尝试层看不见的事件——缓存命中、scope 级失败、取消;
RequestRejected/ResultInvalid 已被尝试层记录,最外层放行不重复记。
"""

from __future__ import annotations

import asyncio
import json
import time
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING

from loguru import logger

from polygateway.errors import (
    GatewayUnavailableError,
    GovernanceBackendError,
    SourceNotConfiguredError,
)
from polygateway.middleware.cache import digest_messages
from polygateway.types import canonical_sampling_json, merge_sampling

if TYPE_CHECKING:
    from collections.abc import Callable, Mapping
    from typing import Any

    from polygateway.ports import CallNext, TelemetryRecorder
    from polygateway.pricing import PricingTable
    from polygateway.types import ChatRequest, LLMResponse, SourceConfig


def _canonical_meta_json(meta: Mapping[str, Any]) -> str:
    """把调用方自定义维度定型为 JSON 文本(issue #11);空 dict 落字面量 `'{}'`。

    `sort_keys=True` 让同一份维度在任意两行里字节一致,可直接等值比对与去重;
    `ensure_ascii=False` 保留中文原文,避免落库成 `\\uXXXX` 串而无法肉眼审计。

    `allow_nan=False` 是**第二道闸**(主防线是 `types.validate_caller_dimensions`
    在公共入口的校验): `json.dumps` 默认把 `nan` 写成裸 `NaN` 字面量,那不是合法
    JSON。这道闸真正的价值在 **SQLite 侧**——PG 的 JSONB 本来就会拒收 `NaN`,而
    SQLite 的 `meta` 是 TEXT 列**不做任何 JSON 校验**,没有这道闸就会把 `NaN`
    这种非法 JSON 静默存进去,污染后续一切按 JSON 解析 meta 的分析。

    注意它抛出的 `ValueError` **不会外泄给调用方**: 本函数在 `_record` 的降级
    `try` 内被求值,异常会被那里的 `except Exception` 接住 → 落 warning、整行
    遥测丢弃。即入口失守时的真实结果是"警告 + 丢一行",不是"报错给调用方"。
    """
    if not meta:
        return "{}"
    return json.dumps(dict(meta), sort_keys=True, ensure_ascii=False, allow_nan=False)


def _cap_text(text: str, cap: int | None) -> str:
    """超出 cap 时头部硬切并附省略标记 `…(略 N 字)`;cap 为 None 原样返回。"""
    if cap is None or len(text) <= cap:
        return text
    return f"{text[:cap]}…(略 {len(text) - cap} 字)"


def _cap_part(part: Any, cap: int) -> Any:
    """多模态 part 的文本截断;非 `type == "text"` 的 part 原样返回同一对象。"""
    if isinstance(part, dict) and part.get("type") == "text" and isinstance(part.get("text"), str):
        return {**part, "text": _cap_text(part["text"], cap)}
    return part


def _cap_messages(messages: list[dict[str, Any]], cap: int | None) -> list[dict[str, Any]]:
    """对每条消息的文本 content 与多模态 part 中 type == "text" 的 text 逐条施加 cap。

    非字符串 content 原样放行(外部输入形状不可控,遥测路径不得因此抛错)。

    **只产出新对象,严禁就地修改**: `digest_messages` 对 content 非 list 的消息是
    原样透传**同一个 dict 对象**(`cache.py:43`),多模态里非 image_url 的 part 同理。
    就地改它会一并污染调用方持有的 messages、后续重试尝试的请求体与缓存写入的 key,
    且全程无任何报错。
    """
    if cap is None:
        return messages
    capped: list[dict[str, Any]] = []
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, str):
            capped.append({**msg, "content": _cap_text(content, cap)})
        elif isinstance(content, list):
            capped.append({**msg, "content": [_cap_part(part, cap) for part in content]})
        else:
            capped.append(msg)
    return capped


@dataclass(frozen=True)
class _AttemptUsage:
    """一次尝试的用量视图;默认值即"失败尝试"档(无用量可言,记 0 并标 unavailable)。

    存在的理由是把 `emit_attempt` 里逐字段重复的 `X if response else Y` 收敛为
    一处判定——十处三元把该方法推到圈复杂度 C,而它们表达的是同一件事。
    """

    response_text: str = ""
    thinking: str = ""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    usage_source: str = "unavailable"
    ttft_ms: float | None = None
    max_inter_token_ms: float | None = None
    cached_prompt_tokens: int | None = None
    model_reported: str | None = None
    reasoning_tokens: int | None = None

    @classmethod
    def of(cls, response: LLMResponse | None) -> _AttemptUsage:
        """从响应取用量;`None`(失败尝试)返回全默认视图。"""
        if response is None:
            return cls()
        return cls(
            response_text=response.content,
            thinking=response.thinking,
            prompt_tokens=response.prompt_tokens,
            completion_tokens=response.completion_tokens,
            usage_source=response.usage_source,
            ttft_ms=response.ttft_ms,
            max_inter_token_ms=response.max_inter_token_ms,
            cached_prompt_tokens=response.cached_prompt_tokens,
            model_reported=response.model_reported,
            reasoning_tokens=response.reasoning_tokens,
        )


class TelemetryEmitter:
    """从请求与结果组装 24 字段并写入 recorder;一切写失败降级 warning。"""

    def __init__(
        self,
        recorder: TelemetryRecorder,
        *,
        pricing: PricingTable | None = None,
        text_cap: int | None,
    ) -> None:
        """`text_cap` 无默认值是有意的: 它是关键行为参数,漏传即静默改变落库正文。

        本类是库内部类,唯一构造者是三个公共 Client,必填能保证没有一处漏传。
        同理,值域校验也放在这一处: 三个 Client 的 `text_cap` 全部汇流到这里,
        `GatewaySettings` 那道只管 env 一条路,而直接构造 Client 是库承诺的另一
        条公共装配路——`text_cap=0` 会让每条正文只剩一个省略标记(P5 不得静默)。
        """
        if text_cap is not None and text_cap <= 0:
            raise ValueError(f"text_cap 必须 > 0(不截断请传 None): {text_cap}")
        self._recorder = recorder
        self._pricing = pricing
        self._text_cap = text_cap

    async def emit_attempt(
        self,
        *,
        request: ChatRequest,
        source: SourceConfig,
        call_id: str,
        latency_ms: int,
        response: LLMResponse | None,
        error: str | None,
    ) -> None:
        """逐次尝试记录(RetryMW 调用);失败尝试无用量可言,记 0 并标 unavailable。"""
        usage = _AttemptUsage.of(response)
        await self._record(
            request=request,
            call_id=call_id,
            model=source.model,
            provider=source.provider,
            source_name=source.name,
            response_text=usage.response_text,
            thinking=usage.thinking,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            usage_source=usage.usage_source,
            latency_ms=latency_ms,
            ttft_ms=usage.ttft_ms,
            max_inter_token_ms=usage.max_inter_token_ms,
            cache_hit=False,
            error=error,
            cached_prompt_tokens=usage.cached_prompt_tokens,
            model_reported=usage.model_reported,
            reasoning_tokens=usage.reasoning_tokens,
            # 唯一有"生效源"的入口,故是唯一能并上 extra_body 的(设计决策 D)
            sampling=canonical_sampling_json(merge_sampling(source.extra_body, request.sampling)),
            tenant_id=request.tenant_id,
            meta=request.meta,
        )

    async def emit_cache_hit(self, *, request: ChatRequest, response: LLMResponse) -> None:
        """缓存命中记录: cache_hit=True、latency_ms=0(VT 同款)。"""
        await self._record(
            request=request,
            call_id=response.call_id,
            model=response.model,
            provider=response.provider,
            source_name=response.source_name,
            response_text=response.content,
            thinking=response.thinking,
            prompt_tokens=response.prompt_tokens,
            completion_tokens=response.completion_tokens,
            usage_source=response.usage_source,
            latency_ms=0,
            ttft_ms=None,
            max_inter_token_ms=None,
            cache_hit=True,
            error=None,
            # 决策 B1: 与 model/prompt_tokens 同一口径,原样回放历史值。
            # 统计供应商缓存命中率必须带 WHERE cache_hit = false,否则重复计数。
            cached_prompt_tokens=response.cached_prompt_tokens,
            model_reported=response.model_reported,
            reasoning_tokens=response.reasoning_tokens,
            # 由最外层 TelemetryMW 调用,手上没有 source。缓存命中行无损:
            # sampling 已进缓存 key,能命中即意味调用级参数与历史那次逐字相同
            sampling=canonical_sampling_json(request.sampling),
            # 与上面的 model/prompt_tokens 相反,维度读 request 而非 response:
            # 维度回答的是"本次调用由谁发起",不是历史那次。读历史会把本次调用
            # 记到上一个租户头上,两边的账同时错且无任何报错(issue #11 设计 §4.3)
            tenant_id=request.tenant_id,
            meta=request.meta,
        )

    async def emit_terminal_failure(
        self, *, request: ChatRequest, call_id: str, latency_ms: int, error: str
    ) -> None:
        """scope 级失败/取消记录: 无具体源,溯源字段置空标记。"""
        await self._record(
            request=request,
            call_id=call_id,
            model="",
            provider="",
            source_name="",
            response_text="",
            thinking="",
            prompt_tokens=0,
            completion_tokens=0,
            usage_source="unavailable",
            latency_ms=latency_ms,
            ttft_ms=None,
            max_inter_token_ms=None,
            cache_hit=False,
            error=error,
            cached_prompt_tokens=None,
            model_reported=None,
            reasoning_tokens=None,
            # 无具体源,与 model/provider/source_name 置空同一先例(设计决策 D)
            sampling=canonical_sampling_json(request.sampling),
            # 源不可知,但租户归属是已知的——终态失败行恰是审计最需要的
            tenant_id=request.tenant_id,
            meta=request.meta,
        )

    async def _record(
        self,
        *,
        request: ChatRequest,
        call_id: str,
        model: str,
        provider: str,
        source_name: str,
        response_text: str,
        thinking: str,
        prompt_tokens: int,
        completion_tokens: int,
        usage_source: str,
        latency_ms: int,
        ttft_ms: float | None,
        max_inter_token_ms: float | None,
        cache_hit: bool,
        error: str | None,
        cached_prompt_tokens: int | None,
        model_reported: str | None,
        sampling: str | None,
        reasoning_tokens: int | None,
        # issue #11: 未归一化的调用方维度,归一化在本方法内收口(recorder 只落库)
        tenant_id: str | None,
        meta: Mapping[str, Any],
    ) -> None:
        try:
            # 成本换算(M2 §6): 成功行按单价换算;缓存命中 0.0(未产生新调用);
            # 失败/终态行 None;未注入价格表 = 恒 None(M1 现状)
            if cache_hit:
                cost: float | None = 0.0
            elif usage_source == "unavailable":
                # 用量不可得: 宁可算不出成本,也不算错成本(解耦设计 §3.1 不变式)。
                # 必须排在 cache_hit 之后——缓存命中未产生新调用,0.0 是事实而非未知
                cost = None
            elif error is None and model and self._pricing is not None:
                cost = self._pricing.cost(
                    model, prompt_tokens, completion_tokens, cached_prompt_tokens
                )
            else:
                cost = None
            # messages 落库前多模态摘要,与缓存 key 共用同一函数(VT R12);
            # 截断只发生在摘要之后、序列化之前的遥测分支,缓存路径不经过它(issue #12)
            messages_json = json.dumps(
                _cap_messages(digest_messages(request.messages), self._text_cap),
                ensure_ascii=False,
            )
            await self._recorder.record_llm_call(
                call_id=call_id,
                parent_call_id=request.parent_call_id,
                session_id=request.session_id,
                model=model,
                provider=provider,
                source_name=source_name,
                messages=messages_json,
                response=_cap_text(response_text, self._text_cap),
                thinking=_cap_text(thinking, self._text_cap),
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                usage_source=usage_source,
                latency_ms=latency_ms,
                ttft_ms=ttft_ms,
                max_inter_token_ms=max_inter_token_ms,
                cache_hit=cache_hit,
                error=error,
                cost=cost,
                cached_prompt_tokens=cached_prompt_tokens,
                model_reported=model_reported,
                sampling=sampling,
                reasoning_tokens=reasoning_tokens,
                # 空串是哨兵而非 NULL: NULL 的 tenant_id 在 PG 的 RLS policy 下
                # 对所有人永久不可见,空串则可用一条 SQL 审计出未归属的行
                tenant_id=tenant_id or "",
                meta=_canonical_meta_json(meta),
            )
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("遥测记录失败(降级不冒泡): {}", exc)


class TelemetryMW:
    """洋葱最外层: 观测尝试层看不见的路径,任何路径都留痕(遥测必录)。"""

    def __init__(
        self, emitter: TelemetryEmitter, now: Callable[[], float] = time.monotonic
    ) -> None:
        self._emitter = emitter
        self._now = now

    async def __call__(self, request: ChatRequest, call_next: CallNext) -> LLMResponse:
        started = self._now()
        try:
            response = await call_next(request)
        except (GatewayUnavailableError, GovernanceBackendError, SourceNotConfiguredError) as exc:
            await self._emitter.emit_terminal_failure(
                request=request,
                call_id=str(uuid.uuid4()),
                latency_ms=int((self._now() - started) * 1000),
                error=str(exc),
            )
            raise
        except asyncio.CancelledError:
            # 尽力而为: 取消也留痕(§5.1 约定④);随后立即重抛
            await self._emitter.emit_terminal_failure(
                request=request,
                call_id=str(uuid.uuid4()),
                latency_ms=int((self._now() - started) * 1000),
                error="cancelled",
            )
            raise
        if response.cache_hit:
            await self._emitter.emit_cache_hit(request=request, response=response)
        return response