feat: add openai-compatible transport with sse parsing and error translation

2026-07-20 06:45:55 -04:00
parent 315142ceb1
commit 454a8b5e0f
2 changed files with 559 additions and 0 deletions
@@ -0,0 +1,304 @@
+"""OpenAI 兼容 transport(D2 默认): 手写 httpx + SSE 解析 + 错误翻译。
+
+SSE 纯函数移植 VT `adapters/llm.py:51-124`;错误翻译移植 CHS
+`app/providers/invokers.py:127-227`。职责只到"一次原始调用"——重试/限流/
+缓存归中间件。看门狗活性口径: content 与 reasoning_content 增量都作为流
+元素产出,思考流天然刷新计时(CHS 迁移约束 R1)。
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from polygateway.errors import RequestRejectedError, SourceDeadError, TransientError
+from polygateway.providers import ProviderProfile, get_provider
+from polygateway.streaming import StreamLivenessTimeout, stream_with_liveness_timeouts
+from polygateway.types import SourceConfig, TransportResult
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable, Mapping
+
+_THINK_PATTERN = re.compile(r"<think>(.*?)</think>", re.DOTALL)
+
+# —— SSE 纯函数(VT llm.py 同款)——
+
+
+def _sse_data_payload(raw: str) -> str | None:
+    """提取 SSE data 行载荷;ping(: 开头)/空行/非 data 行返回 None 跳过。"""
+    line = raw.strip()
+    if not line or line.startswith(":") or not line.startswith("data:"):
+        return None
+    return line[len("data:") :].strip()
+
+
+def _sse_delta(chunk: dict[str, Any], usage_sink: dict[str, Any]) -> tuple[bool, str] | None:
+    """从 chunk 提取增量: (True, content) 或 (False, reasoning);usage 帧旁路进 sink。"""
+    if chunk.get("usage"):
+        usage_sink["usage"] = chunk["usage"]
+    choices = chunk.get("choices") or []
+    if not choices:
+        return None
+    delta = choices[0].get("delta") or {}
+    content = delta.get("content")
+    if content:
+        return (True, content)
+    reasoning = delta.get("reasoning_content")
+    if reasoning:
+        return (False, reasoning)
+    return None
+
+
+async def _iter_sse_deltas(
+    lines: AsyncIterator[str], usage_sink: dict[str, Any]
+) -> AsyncIterator[tuple[bool, str]]:
+    """逐行解析 SSE 流;[DONE] 置 sink["done"];畸形 JSON 帧 → 瞬时错误(可重试)。"""
+    async for raw in lines:
+        data = _sse_data_payload(raw)
+        if data is None:
+            continue
+        if data == "[DONE]":
+            usage_sink["done"] = True
+            return
+        try:
+            chunk = json.loads(data)
+        except json.JSONDecodeError as exc:
+            raise TransientError(f"SSE 帧畸形(malformed_json): {data[:80]!r}", operation="chat") from exc
+        delta = _sse_delta(chunk, usage_sink)
+        if delta is not None:
+            yield delta
+
+
+# —— 错误翻译(CHS invokers.py 同款)——
+
+
+def _parse_retry_after(raw: str | None) -> float | None:
+    """解析 Retry-After 头;仅支持秒数形态,HTTP-date 返回 None(CHS 同款)。"""
+    if raw is None:
+        return None
+    try:
+        seconds = float(raw.strip())
+    except ValueError:
+        return None
+    return seconds if seconds > 0 else None
+
+
+def _translate_429(source: SourceConfig, body_text: str, headers: Mapping[str, str]) -> Exception:
+    try:
+        err_type = json.loads(body_text).get("error", {}).get("type", "")
+    except (json.JSONDecodeError, AttributeError):
+        err_type = ""
+    if err_type == "insufficient_quota":
+        return SourceDeadError(
+            f"{source.name} 配额耗尽(insufficient_quota)",
+            source_name=source.name, status_code=429, operation="chat",
+        )
+    return TransientError(
+        f"{source.name} 限速: 429",
+        retry_after_s=_parse_retry_after(headers.get("retry-after")),
+        source_name=source.name, status_code=429, operation="chat",
+    )
+
+
+def _status_to_error(
+    source: SourceConfig, status: int, body_text: str, headers: Mapping[str, str]
+) -> Exception:
+    ctx: dict[str, Any] = {"source_name": source.name, "status_code": status, "operation": "chat"}
+    if status in (401, 403):
+        return SourceDeadError(f"{source.name} 凭据失效/欠费: {status}", **ctx)
+    if status == 400:
+        return RequestRejectedError(f"{source.name} 请求被拒: 400", **ctx)
+    if status == 429:
+        return _translate_429(source, body_text, headers)
+    if status >= 500:
+        return TransientError(f"{source.name} 瞬时错误: {status}", **ctx)
+    return RequestRejectedError(f"{source.name} 客户端错误: {status}", **ctx)
+
+
+def _strip_think(content: str) -> tuple[str, str]:
+    """剥离 <think> 标签(qwen 系),返回 (正文, 思考流)。VT llm.py:147-164 同款。"""
+    match = _THINK_PATTERN.search(content)
+    if match is None:
+        return content, ""
+    return _THINK_PATTERN.sub("", content).strip(), match.group(1).strip()
+
+
+def _resolve_usage(usage: dict[str, Any], source: SourceConfig) -> tuple[int, int, str]:
+    """usage 帧读取;缺失/非法按 est_tokens 保守兜底并标 estimated(CHS invokers.py:241)。"""
+    prompt, completion = usage.get("prompt_tokens"), usage.get("completion_tokens")
+    if isinstance(prompt, int) and isinstance(completion, int) and prompt + completion > 0:
+        return prompt, completion, "measured"
+    return 0, source.est_tokens, "estimated"
+
+
+def _default_client_factory(source: SourceConfig) -> httpx.AsyncClient:
+    return httpx.AsyncClient(
+        headers={"Authorization": f"Bearer {source.api_key}"},
+        timeout=httpx.Timeout(source.timeout_s),
+        trust_env=source.trust_env,
+    )
+
+
+class OpenAICompatTransport:
+    """默认 transport: 每源一个预配 httpx client,懒创建,aclose 统一释放。"""
+
+    def __init__(
+        self,
+        *,
+        registry: Mapping[str, ProviderProfile] | None = None,
+        client_factory: Callable[[SourceConfig], httpx.AsyncClient] | None = None,
+    ) -> None:
+        self._registry = registry
+        self._client_factory = client_factory or _default_client_factory
+        self._clients: dict[str, httpx.AsyncClient] = {}
+
+    def _client_for(self, source: SourceConfig) -> httpx.AsyncClient:
+        client = self._clients.get(source.name)
+        if client is None:
+            client = self._client_factory(source)
+            self._clients[source.name] = client
+        return client
+
+    def _build_payload(
+        self, *, messages: list[dict[str, Any]], source: SourceConfig,
+        profile: ProviderProfile, stream: bool, overlay: dict[str, Any],
+    ) -> dict[str, Any]:
+        payload: dict[str, Any] = {"model": source.model, "messages": messages, "stream": stream}
+        if stream:
+            payload["stream_options"] = {"include_usage": True}  # 强制 usage 帧(三项目同款)
+        if source.enable_thinking is True:
+            payload.update(profile.thinking_on)
+        elif source.enable_thinking is False:
+            payload.update(profile.thinking_off)
+        payload.update(overlay)
+        return payload
+
+    async def complete(
+        self, *, messages: list[dict[str, Any]], source: SourceConfig,
+        stream: bool, overlay: dict[str, Any], call_id: str,
+    ) -> TransportResult:
+        """一次原始调用;HTTP/线路/流式异常按 ARCH §6.2 翻译为领域错误。"""
+        profile = get_provider(source.provider, registry=self._registry)
+        payload = self._build_payload(
+            messages=messages, source=source, profile=profile, stream=stream, overlay=overlay
+        )
+        url = source.base_url.rstrip("/") + "/chat/completions"
+        client = self._client_for(source)
+        ctx: dict[str, Any] = {"source_name": source.name, "operation": "chat"}
+        try:
+            if stream:
+                return await self._complete_stream(client, url, payload, source, profile)
+            return await self._complete_once(client, url, payload, source, profile)
+        except StreamLivenessTimeout as exc:
+            raise TransientError(f"{source.name} 流活性超时({exc.kind})", **ctx) from exc
+        except httpx.TimeoutException as exc:
+            raise TransientError(f"{source.name} 超时: {exc}", **ctx) from exc
+        except httpx.TransportError as exc:
+            # VT 宽集: 覆盖断连/协议错误/读写失败(设计 §9 行 8)
+            raise TransientError(f"{source.name} 网络错误: {exc}", **ctx) from exc
+
+    async def _complete_stream(
+        self, client: httpx.AsyncClient, url: str, payload: dict[str, Any],
+        source: SourceConfig, profile: ProviderProfile,
+    ) -> TransportResult:
+        started = time.monotonic()
+        async with client.stream("POST", url, json=payload) as resp:
+            if resp.status_code != 200:
+                body = (await resp.aread()).decode("utf-8", errors="replace")
+                raise _status_to_error(source, resp.status_code, body, resp.headers)
+            sink: dict[str, Any] = {}
+            guarded = stream_with_liveness_timeouts(
+                _iter_sse_deltas(resp.aiter_lines(), sink),
+                ttft_s=source.ttft_timeout_s or source.timeout_s,
+                inter_token_s=source.inter_token_timeout_s or source.timeout_s,
+                total_s=source.timeout_s,
+            )
+            content_parts: list[str] = []
+            thinking_parts: list[str] = []
+            ttft_ms: float | None = None
+            last = started
+            max_gap = 0.0
+            async for is_content, text in guarded:
+                now = time.monotonic()
+                if ttft_ms is None:
+                    ttft_ms = (now - started) * 1000
+                else:
+                    max_gap = max(max_gap, (now - last) * 1000)
+                last = now
+                (content_parts if is_content else thinking_parts).append(text)
+        salvaged = self._check_done(sink, content_parts, thinking_parts, source)
+        content, thinking = self._finalize_text(content_parts, thinking_parts, profile)
+        prompt, completion, usage_source = _resolve_usage(sink.get("usage") or {}, source)
+        if salvaged:
+            usage_source = "estimated"  # 打捞路径强制 estimated(设计 §6)
+        return TransportResult(
+            content=content, thinking=thinking, prompt_tokens=prompt,
+            completion_tokens=completion, usage_source=usage_source,
+            ttft_ms=ttft_ms, max_inter_token_ms=(max_gap if ttft_ms is not None else None),
+            raw={"usage": sink.get("usage")},
+        )
+
+    def _check_done(
+        self, sink: dict[str, Any], content_parts: list[str],
+        thinking_parts: list[str], source: SourceConfig,
+    ) -> bool:
+        """缺 [DONE] 语义(设计 §6): 零内容恒 retry;有内容按 missing_done 策略。"""
+        if sink.get("done"):
+            return False
+        ctx: dict[str, Any] = {"source_name": source.name, "operation": "chat"}
+        if not content_parts and not thinking_parts:
+            raise TransientError(f"{source.name} SSE early_eof: 零内容断流", **ctx)
+        if source.missing_done == "retry":
+            raise TransientError(f"{source.name} SSE missing_done: 截断且无 [DONE]", **ctx)
+        return True
+
+    def _finalize_text(
+        self, content_parts: list[str], thinking_parts: list[str], profile: ProviderProfile
+    ) -> tuple[str, str]:
+        content = "".join(content_parts)
+        thinking = "".join(thinking_parts)
+        if profile.strip_think_tags and "<think>" in content:
+            content, tag_thinking = _strip_think(content)
+            if tag_thinking:
+                thinking = tag_thinking
+        return content, thinking
+
+    async def _complete_once(
+        self, client: httpx.AsyncClient, url: str, payload: dict[str, Any],
+        source: SourceConfig, profile: ProviderProfile,
+    ) -> TransportResult:
+        """非流式快路径(三项目均无,库新增): 单 JSON 响应,仅 total 超时。"""
+        resp = await client.post(url, json=payload)
+        if resp.status_code != 200:
+            raise _status_to_error(source, resp.status_code, resp.text, resp.headers)
+        try:
+            body = resp.json()
+        except json.JSONDecodeError as exc:
+            raise TransientError(
+                f"{source.name} 非流式响应非法 JSON", source_name=source.name, operation="chat"
+            ) from exc
+        choices = body.get("choices") or []
+        if not choices:
+            raise TransientError(
+                f"{source.name} 响应缺 choices", source_name=source.name, operation="chat"
+            )
+        message = choices[0].get("message") or {}
+        content, thinking = self._finalize_text(
+            [message.get("content") or ""], [message.get("reasoning_content") or ""], profile
+        )
+        prompt, completion, usage_source = _resolve_usage(body.get("usage") or {}, source)
+        return TransportResult(
+            content=content, thinking=thinking, prompt_tokens=prompt,
+            completion_tokens=completion, usage_source=usage_source,
+            ttft_ms=None, max_inter_token_ms=None, raw={"usage": body.get("usage")},
+        )
+
+    async def aclose(self) -> None:
+        """幂等关闭全部源 client。"""
+        clients, self._clients = self._clients, {}
+        for client in clients.values():
+            await client.aclose()