server-core/app/services/compute_client.py

# -*- 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 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) 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)}