94 lines
3.8 KiB
Python
94 lines
3.8 KiB
Python
"""神识客户端：灵台唯一的模型出口边界。
|
||
|
||
两种实现可互换（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 Protocol
|
||
|
||
import httpx
|
||
|
||
from app.shenshi_client.types import ShenshiChatRequest, ShenshiDelta, ShenshiError
|
||
|
||
|
||
class ShenshiClient(Protocol):
|
||
"""神识数据面契约：流式对话补全。"""
|
||
|
||
def stream_chat(self, request: ShenshiChatRequest) -> AsyncIterator[ShenshiDelta]: ...
|
||
|
||
|
||
class EchoShenshi:
|
||
"""进程内回声桩：把最后一条用户消息按固定粒度切分回流。
|
||
|
||
只用于 M0 骨架验证（登录 → 占位 Run → SSE 端到端），不模拟任何 Provider 行为。
|
||
"""
|
||
|
||
chunk_size = 8
|
||
|
||
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),
|
||
)
|
||
|
||
|
||
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
|
||
|
||
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()
|