Pine
e44fea96ca
- relay 鉴权重写：JWT→引擎 PAT 逐用户归属，拒绝匿名调用，余额<=0 预检 403 - compute_client 增 list_engine_users/get_engine_user/sync_user_mirror 等 - 运营端新增 /admin/compute/platform-users 平台用户视图(引擎真实余额/折扣/用量)与充值镜像回写 Co-Authored-By: Claude <noreply@anthropic.com>
317 lines
12 KiB
Python
317 lines
12 KiB
Python
# -*- 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<path>``，返回 ``{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)}
|