agent-Specialization/modules/balance_client.py

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

Credentials (优先读取仓库根目录 .env，其次环境变量):
  - 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 pathlib import Path
from typing import Any, Dict, Tuple
from urllib import parse, request

# -------- .env 读取助手（简化版，已由 config/__init__.py 统一注入） --------


def _env(key: str, default: str | None = None) -> str | None:
    """从环境变量获取配置。"""
    if key in os.environ:
        return os.environ[key]
    return default


USD_CNY_RATE = float(_env("USD_CNY_RATE", "7.1") or "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 = _env("MOONSHOT_API_KEY") or _env("API_KEY_KIMI") or _env("AGENT_API_KEY")
    if not api_key:
        return {"success": False, "error": "Kimi 密钥未设置（MOONSHOT_API_KEY / API_KEY_KIMI / AGENT_API_KEY）"}

    base = (
        _env("API_BASE_KIMI")
        or _env("MOONSHOT_BASE_URL")
        or _env("AGENT_API_BASE_URL")
        or "https://api.moonshot.ai/v1"
    ).rstrip("/")
    url = f"{base}/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": payload.get("currency") or "CNY",
            "available": available,
            "voucher": voucher,
            "cash": cash,
            "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 = _env("DEEPSEEK_API_KEY") or _env("API_KEY_DEEPSEEK")
    if not api_key:
        return {"success": False, "error": "DeepSeek 密钥未设置（DEEPSEEK_API_KEY / API_KEY_DEEPSEEK）"}

    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 = _env("ALIYUN_ACCESS_KEY_ID")
    sk = _env("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

    # 按阿里云规范组装最终查询字符串（不可用 urlencode 的 quote_plus）
    query = "&".join(f"{_percent_encode(k)}={_percent_encode(v)}" for k, v in params.items())
    url = "https://business.aliyuncs.com/?" + query

    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(),
    }