hunjing/backend/app/shenshi_client/client.py

"""神识客户端：灵台唯一的模型出口边界。

两种实现可互换（Protocol）：
- ``HttpShenshiClient``：经 HTTP+SSE 调用 Go 神识服务（SHENSHI_BASE_URL 配置时启用）；
- ``EchoShenshi``：进程内回声桩，逐字回流输入，供本地开发与测试在无神识时跑通端到端。

领域模块禁止绕过本边界 import Provider SDK（docs/architecture/service-boundaries.md §4）。
"""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import Any, Protocol

import httpx

from app.shenshi_client.types import ShenshiChatRequest, ShenshiDelta, ShenshiError


class ShenshiClient(Protocol):
    """神识数据面契约：流式对话补全 + 用户面（BYOK/用量，神识阁数据源）。"""

    def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]: ...

    async def get_my_usage(self, user_id: str) -> dict[str, Any]: ...

    async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]: ...

    async def register_byok_key(
        self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str
    ) -> dict[str, Any]: ...

    async def revoke_byok_key(self, user_id: str, key_id: str) -> None: ...


class EchoShenshi:
    """进程内回声桩：把最后一条用户消息按固定粒度切分回流。

    数据面只用于骨架验证（登录 → 占位 Run → SSE 端到端）；用户面为最小内存实现，
    供神识阁在无神识部署时演示与测试桩。
    """

    chunk_size = 8

    def __init__(self) -> None:
        self._byok: dict[str, list[dict[str, Any]]] = {}
        self._usage: dict[str, list[dict[str, Any]]] = {}

    async def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]:
        prompt = next((m.content for m in reversed(request.messages) if m.role == "user"), "")
        text = f"【回声桩】{prompt}"
        for i in range(0, len(text), self.chunk_size):
            yield ShenshiDelta(delta=text[i : i + self.chunk_size])
        yield ShenshiDelta(
            finish_reason="stop",
            usage_prompt_tokens=max(len(prompt) // 4, 1),
            usage_completion_tokens=max(len(text) // 4, 1),
        )
        if request.user_id:
            self._usage.setdefault(request.user_id, []).append(
                {
                    "user_id": request.user_id,
                    "model": request.model or "",
                    "endpoint": "echo",
                    "requests": 1,
                    "ok": 1,
                    "prompt_tokens": max(len(prompt) // 4, 1),
                    "completion_tokens": max(len(text) // 4, 1),
                }
            )

    async def get_my_usage(self, user_id: str) -> dict[str, Any]:
        return {"rows": self._usage.get(user_id, [])}

    async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]:
        return self._byok.get(user_id, [])

    async def register_byok_key(
        self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str
    ) -> dict[str, Any]:
        entry = {
            "id": f"byok_echo_{len(self._byok.get(user_id, [])) + 1}",
            "name": name,
            "type": type,
            "base_url": base_url,
            "model": model,
            "api_key_mask": api_key[:3] + "****",
        }
        self._byok.setdefault(user_id, []).append(entry)
        return entry

    async def revoke_byok_key(self, user_id: str, key_id: str) -> None:
        self._byok[user_id] = [k for k in self._byok.get(user_id, []) if k["id"] != key_id]


class HttpShenshiClient:
    """HTTP+SSE 神识客户端。

    认证：服务令牌经 ``Authorization: Bearer`` 透传（管理面 RBAC 不在此层）。
    流解析：``data:`` 帧为 ShenshiDelta JSON，``[DONE]`` 结束；非 2xx 映射 ShenshiError。
    """

    def __init__(self, base_url: str, service_token: str, timeout: float = 120.0) -> None:
        self._base_url = base_url.rstrip("/")
        self._headers = {"Authorization": f"Bearer {service_token}"} if service_token else {}
        self._timeout = timeout

    def _user_headers(self, user_id: str) -> dict[str, str]:
        return {**self._headers, "X-Hunjing-User-Id": user_id}

    async def _request_json(
        self, method: str, path: str, *, user_id: str, json_body: dict[str, Any] | None = None
    ) -> Any:
        async with httpx.AsyncClient(trust_env=False, timeout=self._timeout) as client:
            response = await client.request(
                method,
                f"{self._base_url}{path}",
                json=json_body,
                headers=self._user_headers(user_id),
            )
        if response.status_code >= 400:
            code = f"HJ-2{max(response.status_code - 400, 0):03d}"
            try:
                detail = response.json()
                code = detail.get("error", {}).get("code", code)
                message = detail.get("error", {}).get("message", response.text)
            except ValueError:
                message = response.text
            raise ShenshiError(code, message)
        if response.status_code == 204 or not response.content:
            return None
        return response.json()

    async def get_my_usage(self, user_id: str) -> dict[str, Any]:
        data = await self._request_json("GET", "/me/usage", user_id=user_id)
        return data or {}

    async def list_byok_keys(self, user_id: str) -> list[dict[str, Any]]:
        data = await self._request_json("GET", "/me/byok-keys", user_id=user_id)
        return (data or {}).get("keys", [])

    async def register_byok_key(
        self, user_id: str, *, name: str, type: str, base_url: str, api_key: str, model: str
    ) -> dict[str, Any]:
        data = await self._request_json(
            "POST",
            "/me/byok-keys",
            user_id=user_id,
            json_body={"name": name, "type": type, "base_url": base_url, "api_key": api_key, "model": model},
        )
        return data or {}

    async def revoke_byok_key(self, user_id: str, key_id: str) -> None:
        await self._request_json("DELETE", f"/me/byok-keys/{key_id}", user_id=user_id)

    async def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]:
        payload = request.model_dump(mode="json")
        assert payload.get("stream") is True  # 数据面当前只有流式形态
        async with httpx.AsyncClient(trust_env=False, timeout=self._timeout) as client:
            async with client.stream(
                "POST",
                f"{self._base_url}/v1/chat/completions",
                json=payload,
                headers=self._headers,
            ) as response:
                if response.status_code >= 400:
                    body = (await response.aread()).decode("utf-8", errors="replace")
                    code = f"HJ-2{max(response.status_code - 400, 0):03d}"
                    try:
                        detail = json.loads(body)
                        code = detail.get("error", {}).get("code", code)
                        body = detail.get("error", {}).get("message", body)
                    except json.JSONDecodeError:
                        pass
                    raise ShenshiError(code, body)
                async for line in response.aiter_lines():
                    if not line.startswith("data:"):
                        continue
                    data = line[len("data:") :].strip()
                    if data == "[DONE]":
                        return
                    if data:
                        yield ShenshiDelta.model_validate_json(data)


def build_shenshi_client(base_url: str, service_token: str) -> ShenshiClient:
    """按配置装配客户端：未配置神识地址时降级为回声桩（M0 骨架形态）。"""
    if base_url:
        return HttpShenshiClient(base_url, service_token)
    return EchoShenshi()