201 lines
7.0 KiB
Python
201 lines
7.0 KiB
Python
"""
|
||
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 读取助手 --------
|
||
_DOTENV_CACHE: Dict[str, str] | None = None
|
||
|
||
|
||
def _read_dotenv() -> Dict[str, str]:
|
||
global _DOTENV_CACHE
|
||
if _DOTENV_CACHE is not None:
|
||
return _DOTENV_CACHE
|
||
path = Path(__file__).resolve().parents[1] / ".env"
|
||
result: Dict[str, str] = {}
|
||
if path.exists():
|
||
try:
|
||
for raw_line in path.read_text(encoding="utf-8").splitlines():
|
||
line = raw_line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
key, value = line.split("=", 1)
|
||
key = key.strip()
|
||
value = value.strip().strip('"').strip("'")
|
||
if key:
|
||
result[key] = value
|
||
except Exception:
|
||
pass
|
||
_DOTENV_CACHE = result
|
||
return result
|
||
|
||
|
||
def _env(key: str, default: str | None = None) -> str | None:
|
||
if key in os.environ:
|
||
return os.environ[key]
|
||
return _read_dotenv().get(key, 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": "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 = _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://bss.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(),
|
||
}
|