# -*- coding: utf-8 -*- """算力引擎(compute-engine / new-api)客户端封装。 云超服算力统一经此服务消费引擎(loopback :3000,仅服务),前端不直连引擎: 模型中继 ``/v1/*`` + 管理 API ``/api/*``(建号 / 发 PAT / 额度 / 流水)。 - 引擎仅作模型计量与额度/令牌/流水,不承担身份(身份在 server-core)。 - 管理调用带引擎管理令牌(``COMPUTE_ADMIN_TOKEN``);用户侧余额/流水用其 PAT。 路径:引擎管理路由起点见 compute-engine ``controller/user.go``、``controller/log.go``、 ``router/api-router.go`` 的 admin 段;以下默认路径按设计 §4.3 拟定,若与引擎实际 路由不符,仅需调整对应 ``_path`` 常量,不影响上层调用。 """ from __future__ import annotations import asyncio import httpx from .. import config class ComputeError(Exception): """算力引擎调用异常(网络 / 非 2xx / 引擎未就绪等)。""" def _base() -> str: return config.COMPUTE_BASE_URL.rstrip("/") def _admin_headers() -> dict[str, str]: return {"Authorization": f"Bearer {config.COMPUTE_ADMIN_TOKEN}"} async def _request(*, method: str, path: str, headers: dict[str, str] | None = None, json: dict | None = None) -> dict: """带超时/重试的引擎请求;2xx 返回 JSON,否则抛 ComputeError。""" url = f"{_base()}{path}" last: Exception | None = None for _ in range(max(1, config.COMPUTE_RETRIES + 1)): try: async with httpx.AsyncClient(timeout=config.COMPUTE_TIMEOUT) as client: resp = await client.request( method, url, json=json, headers=headers, ) if resp.status_code == 200: try: return resp.json() if resp.content else {} except Exception: # noqa: BLE001 return {} if resp.status_code < 500: raise ComputeError( f"compute-engine {resp.status_code}: {resp.text[:200]}" ) last = ComputeError(f"compute-engine {resp.status_code}: {resp.text[:200]}") except httpx.HTTPError as exc: last = ComputeError(f"compute-engine request failed: {exc}") except ComputeError as exc: last = exc if " 4" in str(exc) or " 3" in str(exc): break await asyncio.sleep(0.2) raise last or ComputeError("compute-engine unavailable") async def ping() -> dict: """引擎连通性健康检查。""" try: return await _request(method="GET", path="/api/status", headers=_admin_headers()) except ComputeError as exc: return {"status": "error", "detail": str(exc)} async def create_user(username: str) -> dict: """在引擎建号(username = 云超服 uid)。""" return await _request( method="POST", path="/api/user", headers=_admin_headers(), json={"username": username}, ) async def issue_pat(username: str) -> str: """为引擎用户签发 personal access token(下发前端注入 provider)。""" data = await _request( method="POST", path="/api/token", headers=_admin_headers(), json={"username": username}, ) return str(data.get("key") or data.get("token") or "") async def get_balance(user_pat: str) -> dict: """用户余额(需该用户 PAT)。""" headers = {"Authorization": f"Bearer {user_pat}"} return await _request(method="GET", path="/api/user/self", headers=headers) async def get_month_usage(user_pat: str) -> dict: """本月消耗(按日额度)。""" headers = {"Authorization": f"Bearer {user_pat}"} return await _request(method="GET", path="/api/data/self", headers=headers) async def get_logs(user_pat: str) -> dict: """用户 token 流水。""" headers = {"Authorization": f"Bearer {user_pat}"} return await _request(method="GET", path="/api/log/self", headers=headers)