vps-manager/app/adapters/ai.py

"""AI 平台适配器（验证 Key + 余额/用量查询）

各 AI 平台 API 差异较大，此处实现 test_connection（验证 Key，附带余额）
与 get_account（余额）。余额接口可能随平台调整，填入真实凭证后以实际返回为准。
所需配置统一为 {"api_key": "..."}（Minimax 另需 group_id）。
"""

import httpx

from app.adapters.base import AccountInfo, BaseAdapter
from app.adapters.registry import register


class _AIBase(BaseAdapter):
    """AI 适配器公共逻辑（Bearer Key + JSON GET）"""

    required_config = ["api_key"]

    def _headers(self) -> dict:
        return {"Authorization": f"Bearer {self.config.get('api_key', '')}"}

    def _get(self, url: str) -> dict:
        with httpx.Client(timeout=30) as client:
            resp = client.get(url, headers=self._headers())
            resp.raise_for_status()
            return resp.json()

    def _http_error(self, e: httpx.HTTPStatusError) -> dict:
        return {"ok": False, "message": f"HTTP {e.response.status_code}：API Key 无效或权限不足"}


@register("deepseek-api")
class DeepSeekAdapter(_AIBase):
    BASE = "https://api.deepseek.com"

    def _balance(self) -> AccountInfo:
        # 官方响应为 {is_available, balance_infos: [{currency,total_balance,...}]}，
        # 没有顶层 data；同时兼容旧版 data.total_balance 形态。
        data = self._get(self.BASE + "/user/balance")
        infos = data.get("balance_infos") or []
        info = next((i for i in infos if i.get("currency") == "CNY"), infos[0] if infos else {})
        if not info and isinstance(data.get("data"), dict):
            info = data["data"]
        try:
            total = float(info.get("total_balance")) if info.get("total_balance") is not None else None
        except (TypeError, ValueError):
            total = None
        return AccountInfo(
            balance=total, currency=info.get("currency") or "CNY", raw=data
        )

    def test_connection(self) -> dict:
        try:
            acc = self._balance()
            return {"ok": True, "message": f"连接成功，余额 {acc.balance} CNY"}
        except httpx.HTTPStatusError as e:
            return self._http_error(e)
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "message": str(e)}

    def get_account(self) -> AccountInfo:
        return self._balance()


@register("moonshot-api")
class KimiAdapter(_AIBase):
    BASE = "https://api.moonshot.cn/v1"

    def _balance(self) -> AccountInfo:
        data = self._get(self.BASE + "/users/me/balance").get("data", {})
        return AccountInfo(balance=data.get("available_balance"), currency="CNY", raw=data)

    def test_connection(self) -> dict:
        try:
            acc = self._balance()
            return {"ok": True, "message": f"连接成功，可用余额 {acc.balance} CNY"}
        except httpx.HTTPStatusError as e:
            return self._http_error(e)
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "message": str(e)}

    def get_account(self) -> AccountInfo:
        return self._balance()


@register("openai-api")
class OpenAIAdapter(_AIBase):
    BASE = "https://api.openai.com/v1"

    def test_connection(self) -> dict:
        try:
            data = self._get(self.BASE + "/models")
            count = len(data.get("data", []))
            return {"ok": True, "message": f"连接成功，可用模型 {count} 个"}
        except httpx.HTTPStatusError as e:
            return self._http_error(e)
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "message": str(e)}

    # 故意不实现 get_account：OpenAI 的 /dashboard/billing/credit_grants 早已下线
    # （且需组织级权限）。旧实现把异常静默吞掉并返回 balance=None，
    # 导致前端提示「已更新」但数值永远不变；不实现则 capabilities() 中
    # get_account 为 False，刷新余额时会如实提示「暂不支持余额查询」。


@register("agnes-api")
class AgnesAdapter(_AIBase):
    """Agnes 适配器（OpenAI 兼容接口，Bearer Key）

    官方 Base URL：https://api.agnes-ai.cn/v1（国际站为 api.agnes-ai.com）
    只做 Key 有效性验证与可用模型查询；官方没有公开余额接口，
    因此不实现 get_account——否则余额刷新会显示「已更新」但数值永远不变。
    """

    BASE = "https://api.agnes-ai.cn/v1"

    def _base_url(self) -> str:
        return (self.config.get("base_url") or self.BASE).rstrip("/")

    def test_connection(self) -> dict:
        try:
            data = self._get(self._base_url() + "/models")
            count = len(data.get("data", []))
            return {"ok": True, "message": f"连接成功，可用模型 {count} 个"}
        except httpx.HTTPStatusError as e:
            return self._http_error(e)
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "message": str(e)}


@register("minimax-api")
class MinimaxAdapter(_AIBase):
    """Minimax 适配器（余额查询）

    接口：GET https://api.minimax.chat/v1/balance
    认证：Bearer api_key（group_id 仅部分旧接口需要，余额查询非必需）
    响应示例：{"balance": 123.45, "currency": "CNY", ...}
    注：Minimax 国内版端点为 api.minimaxi.com，国际版为 api.minimax.chat，
    两者 API Key 不通用；默认使用国际版端点，可通过 config["base_url"] 覆盖。
    """

    required_config = ["api_key"]
    BASE = "https://api.minimax.chat/v1"

    def _base_url(self) -> str:
        return (self.config.get("base_url") or self.BASE).rstrip("/")

    def _balance(self) -> AccountInfo:
        data = self._get(self._base_url() + "/balance")
        balance = data.get("balance")
        try:
            balance = float(balance) if balance is not None else None
        except (TypeError, ValueError):
            balance = None
        currency = data.get("currency") or "CNY"
        return AccountInfo(balance=balance, currency=currency, raw=data)

    def test_connection(self) -> dict:
        try:
            acc = self._balance()
            if acc.balance is not None:
                return {"ok": True, "message": f"连接成功，余额 {acc.balance} {acc.currency}"}
            return {"ok": True, "message": "连接成功（未返回余额字段）"}
        except httpx.HTTPStatusError as e:
            return self._http_error(e)
        except Exception as e:  # noqa: BLE001
            return {"ok": False, "message": str(e)}

    def get_account(self) -> AccountInfo:
        return self._balance()