agent-Specialization/modules/balance_client.py

"""
Admin balance fetchers for Kimi (Moonshot), DeepSeek, and Qwen (Aliyun BSS).

Credentials (read from environment):
  - MOONSHOT_API_KEY            : Bearer token for Kimi
  - DEEPSEEK_API_KEY            : Bearer token for DeepSeek
  - ALIYUN_ACCESS_KEY_ID        : AccessKey ID for Aliyun (Qwen billing)
  - ALIYUN_ACCESS_KEY_SECRET    : AccessKey Secret for Aliyun (Qwen billing)
  - USD_CNY_RATE                : Override USD->CNY rate (default 7.1)
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import time
import uuid
from typing import Any, Dict, Tuple
from urllib import parse, request, error


USD_CNY_RATE = float(os.environ.get("USD_CNY_RATE", "7.1"))


def _http_get(url: str, headers: Dict[str, str] | None = None, timeout: int = 8) -> Tuple[Dict[str, Any] | None, str | None]:
    """Perform a simple GET request, return (json, error_message)."""
    req = request.Request(url, headers=headers or {})
    try:
        with request.urlopen(req, timeout=timeout) as resp:
            data = resp.read()
            return json.loads(data.decode("utf-8")), None
    except Exception as exc:  # broad: network/parsing errors
        return None, str(exc)


# -------- Kimi (Moonshot) --------
def fetch_kimi_balance() -> Dict[str, Any]:
    api_key = os.environ.get("MOONSHOT_API_KEY")
    if not api_key:
        return {"success": False, "error": "MOONSHOT_API_KEY 未设置"}

    url = "https://api.moonshot.ai/v1/users/me/balance"
    payload, err = _http_get(url, headers={"Authorization": f"Bearer {api_key}"})
    if err:
        return {"success": False, "error": err}

    try:
        data = payload.get("data") or {}
        available = float(data.get("available_balance", 0))
        voucher = float(data.get("voucher_balance", 0))
        cash = float(data.get("cash_balance", 0))
        return {
            "success": True,
            "currency": "USD",
            "available": available,
            "available_cny": round(available * USD_CNY_RATE, 2),
            "voucher": voucher,
            "cash": cash,
            "rate": USD_CNY_RATE,
            "raw": payload,
        }
    except Exception as exc:  # pragma: no cover
        return {"success": False, "error": f"解析失败: {exc}", "raw": payload}


# -------- DeepSeek --------
def fetch_deepseek_balance() -> Dict[str, Any]:
    api_key = os.environ.get("DEEPSEEK_API_KEY")
    if not api_key:
        return {"success": False, "error": "DEEPSEEK_API_KEY 未设置"}

    url = "https://api.deepseek.com/user/balance"
    payload, err = _http_get(url, headers={"Authorization": f"Bearer {api_key}"})
    if err:
        return {"success": False, "error": err}

    try:
        infos = payload.get("balance_infos") or []
        primary = infos[0] if infos else {}
        total = float(primary.get("total_balance") or 0)
        granted = float(primary.get("granted_balance") or 0)
        topped_up = float(primary.get("topped_up_balance") or 0)
        return {
            "success": True,
            "currency": primary.get("currency", "CNY"),
            "available": total,
            "granted": granted,
            "topped_up": topped_up,
            "raw": payload,
        }
    except Exception as exc:  # pragma: no cover
        return {"success": False, "error": f"解析失败: {exc}", "raw": payload}


# -------- Qwen (Aliyun BSS QueryAccountBalance) --------
def _percent_encode(val: str) -> str:
    res = parse.quote(str(val), safe="~")
    return res.replace("+", "%20").replace("*", "%2A").replace("%7E", "~")


def _sign(params: Dict[str, Any], secret: str) -> str:
    sorted_params = sorted(params.items(), key=lambda x: x[0])
    canonicalized = "&".join(f"{_percent_encode(k)}={_percent_encode(v)}" for k, v in sorted_params)
    string_to_sign = f"GET&%2F&{_percent_encode(canonicalized)}"
    h = hmac.new((secret + "&").encode("utf-8"), string_to_sign.encode("utf-8"), hashlib.sha1)
    return base64.b64encode(h.digest()).decode("utf-8")


def fetch_qwen_balance() -> Dict[str, Any]:
    ak = os.environ.get("ALIYUN_ACCESS_KEY_ID")
    sk = os.environ.get("ALIYUN_ACCESS_KEY_SECRET")
    if not ak or not sk:
        return {"success": False, "error": "缺少 ALIYUN_ACCESS_KEY_ID / ALIYUN_ACCESS_KEY_SECRET"}

    params: Dict[str, Any] = {
        "Format": "JSON",
        "Version": "2017-12-14",
        "AccessKeyId": ak,
        "SignatureMethod": "HMAC-SHA1",
        "Timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "SignatureVersion": "1.0",
        "SignatureNonce": str(uuid.uuid4()),
        "Action": "QueryAccountBalance",
        "RegionId": "cn-hangzhou",
    }
    signature = _sign(params, sk)
    params["Signature"] = signature
    url = "https://bss.aliyuncs.com/?" + parse.urlencode(params)

    payload, err = _http_get(url)
    if err:
        return {"success": False, "error": err}

    try:
        data = payload.get("Data") or {}
        amount = float(data.get("AvailableAmount") or 0)
        currency = data.get("Currency", "CNY")
        return {
            "success": True,
            "currency": currency,
            "available": amount,
            "cash": float(data.get("AvailableCashAmount") or 0),
            "raw": payload,
        }
    except Exception as exc:  # pragma: no cover
        return {"success": False, "error": f"解析失败: {exc}", "raw": payload}


def fetch_all_balances() -> Dict[str, Any]:
    return {
        "rate_usd_cny": USD_CNY_RATE,
        "kimi": fetch_kimi_balance(),
        "deepseek": fetch_deepseek_balance(),
        "qwen": fetch_qwen_balance(),
    }