"""神识客户端：灵台唯一的模型出口边界。 两种实现可互换（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()