# -*- 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]: """引擎管理员请求头。令牌未配置时不下发 Authorization,避免 `Bearer ` 空令牌被 httpx 判定非法。""" token = config.COMPUTE_ADMIN_TOKEN if not token: return {} return {"Authorization": f"Bearer {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) async def issue_user_token(username: str, name: str = "") -> dict: """为某引擎用户签发消费令牌,返回 {id, key, ...}。""" data = await _request( method="POST", path="/api/token", headers=_admin_headers(), json={"username": username, "name": name}, ) return data.get("data") or {} async def ensure_user(username: str) -> dict: """幂等建引擎用户:已存在(409/4xx)视为成功,不抛。""" try: return await create_user(username) except ComputeError: return {"username": username, "already": True} async def list_models(status: int = 1) -> list[dict]: """列出引擎 admin 模型(models 表)。返回分页 items 列表。""" data = await _request( method="GET", path=f"/api/models?status={status}&p=1&page_size=100", headers=_admin_headers(), ) return (data.get("data") or {}).get("items") or [] async def list_user_tokens(username: str) -> dict: """列出某引擎用户的消费令牌(按用户名挂到引擎用户)。""" data = await _request( method="GET", path=f"/api/token?username={username}", headers=_admin_headers(), ) return (data.get("data") or {}).get("items") or [] async def _user_key(username: str) -> str: """取该引擎用户一枚有效 PAT(用作 /api/user/self* 鉴权)。""" items = await list_user_tokens(username) for it in items or []: k = it.get("key") or it.get("token") if k: return k return "" async def user_balance(username: str) -> dict: """用户余额(PAT 鉴权)。""" key = await _user_key(username) if not key: return {"quota": 0, "used_quota": 0, "balance": 0} data = await _request(method="GET", path="/api/user/self", headers={"Authorization": f"Bearer {key}"}) return data.get("data") or {} async def user_usage(username: str) -> dict: """本月按模型用量/费用(PAT 鉴权)。""" key = await _user_key(username) if not key: return {"items": [], "cost_quota": 0, "input_tokens": 0, "output_tokens": 0} data = await _request(method="GET", path="/api/user/self/usage", headers={"Authorization": f"Bearer {key}"}) return data.get("data") or {} async def delete_token(token_id: int) -> dict: """删除引擎令牌。""" return await _request( method="DELETE", path=f"/api/token/{token_id}", headers=_admin_headers(), ) async def adjust_user_quota(engine_user_id: int, value: int, mode: str = "add") -> dict: """调整引擎用户余额/额度(充值/扣减/覆盖)。 对应引擎 ``POST /api/user/manage``,action="add_quota": - mode="add" 充值(增加余额) - mode="subtract" 扣减 - mode="override" 直接设为该值 ``value`` 为微元(人民币实际金额,1 元 = 1_000_000,见 compute-service meter.py)。 额度耗尽时引擎 /v1 返回 403。 """ return await _request( method="POST", path="/api/user/manage", headers=_admin_headers(), json={"id": int(engine_user_id), "action": "add_quota", "value": int(value), "mode": mode}, ) async def list_engine_users(page: int = 1, page_size: int = 100, status: int = 1) -> list[dict]: """分页拉取引擎用户(GET /api/user),返回 items 列表。""" data = await _request( method="GET", path=f"/api/user?status={int(status)}&type=-1&p={int(page)}&page_size={int(page_size)}", headers=_admin_headers(), ) return (data.get("data") or {}).get("items") or [] async def get_engine_user(engine_user_id: int) -> dict | None: """按 id 取引擎用户(GET /api/user?id=),无则 None。""" data = await _request( method="GET", path=f"/api/user?id={int(engine_user_id)}&p=1&page_size=1", headers=_admin_headers(), ) items = (data.get("data") or {}).get("items") or [] return items[0] if items else None async def sync_user_mirror(db, engine_user_id: int) -> bool: """把引擎用户余额回写到平台镜像(compute_quota/compute_used_quota)。 best-effort:引擎用户不存在或平台用户未映射时静默返回 False。 """ try: engine_user = await get_engine_user(engine_user_id) username = str((engine_user or {}).get("username") or "") if not username: return False balance = await user_balance(username) platform_user = await db.users.get_by_username(username) if platform_user is None: return False await db.users.set_compute_mirror( platform_user["id"], provisioned=True, username=username, quota=int(balance.get("quota") or 0), used_quota=int(balance.get("used_quota") or 0), ) return True except Exception: # noqa: BLE001 return False async def set_user_group_by_username(username: str, group: str) -> dict: """把引擎用户划到某个用户组(企业折扣凭 User.Group 触发组倍率)。 走管理透传 ``PUT /api/user``。best-effort:失败返回 {status,body},由调用方决定是否阻断。 """ return await proxy("PUT", "/user", json_body={"username": username, "group": group}) async def grant_user_quota_by_username(username: str, value: int) -> dict: """按 username 给引擎用户充值(先经管理面解析引擎用户 id,再 add_quota)。 解析:``GET /api/user/search?username=``(new-api 返回 data.id)。best-effort。 """ res = await proxy("GET", "/user/search", params={"username": username}) body = res.get("body") or {} if not isinstance(body, dict) or not body.get("success", True): return {"ok": False, "message": "未找到引擎用户"} data = body.get("data") uid = None if isinstance(data, dict): uid = data.get("id") elif isinstance(data, list) and data: uid = data[0].get("id") if not uid: return {"ok": False, "message": "未解析到引擎用户 id"} return await adjust_user_quota(int(uid), int(value), "add") async def set_group_group_ratio(user_group: str, using_group: str, ratio: float) -> dict: """设置「用户组×模型组」倍率(= 1 - discount%),引擎计费时按此折算模型价。 走管理透传 ``PUT /api/setting``(key=group_ratio_setting)。best-effort。 """ return await proxy( "PUT", "/setting", json_body={"key": "group_ratio_setting", "value": {"group_group_ratio": {user_group: {using_group: ratio}}}}, ) async def sync_user_enabled(username: str, enabled: bool) -> dict: """按 username 同步引擎用户启用/禁用(平台用户生命周期自动同步)。""" return await _request( method="PATCH", path="/api/user/status", headers=_admin_headers(), json={"username": username, "status": 1 if enabled else 0}, ) async def sync_delete_user(username: str) -> dict: """按 username 删除引擎用户(平台用户删除时同步)。""" return await _request( method="DELETE", path=f"/api/user/by_username/{username}", headers=_admin_headers(), ) async def proxy(method: str, path: str, *, json_body: dict | None = None, params: dict | None = None) -> dict: """管理 API 透传:转发到引擎 ``/api``,返回 ``{status, body}``。 供 /admin/compute/proxy/** 使用,让 admin 端可 1:1 驱动 compute-engine(new-api) 的完整管理面 (models/channels/groups/tokens/users/logs/redemption/ratio 等),不改写每个端点契约。 保留引擎原始状态码与响应体(网关层不吞错误)。 """ if not config.COMPUTE_ADMIN_TOKEN: return {"status": 503, "body": '{"success": false, "message": "算力引擎内部令牌未配置(PINEAGENTS_INTERNAL_TOKEN)"}'} url = f"{_base()}/api{path}" try: async with httpx.AsyncClient(timeout=config.COMPUTE_TIMEOUT, follow_redirects=True) as client: resp = await client.request( method, url, json=json_body, params=params, headers=_admin_headers(), ) return {"status": resp.status_code, "body": resp.text or "{}"} except httpx.HTTPError as exc: return {"status": 502, "body": '{"error": "%s"}' % str(exc)}