RRRosmontis
34089d0042
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
248 lines
11 KiB
Python
248 lines
11 KiB
Python
# deepseek/deepseek_balance.py — DeepSeek 余额监控
|
||
# 文件夹化模块：配置（API key / 阈值 / 轮询间隔 / 开关）全部从本目录 setting.json 读取。
|
||
# 命令（管理员）: .deepseek help / balance / threshold <int> / polling <int> / alert <on|off>
|
||
# 监控: 定时查询余额，低于阈值时向【所有已注册通道】的管理员私聊告警。
|
||
import asyncio
|
||
import json
|
||
import os
|
||
|
||
import aiohttp
|
||
|
||
import identity
|
||
from channel import push_to_user, send_message
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
SETTING_FILE = os.path.join(BASE_DIR, "setting.json")
|
||
|
||
# ---------------- 配置（从 setting.json 读取） ----------------
|
||
API_KEY = ""
|
||
API_BASE = "https://api.deepseek.com"
|
||
THRESHOLD = 3
|
||
POLLING_MINUTES = 3
|
||
ALERT_ENABLED = True
|
||
# 启动宽限（秒）：首次余额查询前先等 TG 通道 register / NapCat 反向连接完成，
|
||
# 否则首次告警会因通道尚未注册而投递失败（启动竞态）。
|
||
MONITOR_STARTUP_DELAY = 30
|
||
|
||
|
||
def _load_setting() -> dict:
|
||
try:
|
||
with open(SETTING_FILE, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
except (OSError, json.JSONDecodeError):
|
||
data = {}
|
||
return data if isinstance(data, dict) else {}
|
||
|
||
|
||
def _save_setting(data: dict):
|
||
os.makedirs(os.path.dirname(SETTING_FILE), exist_ok=True)
|
||
with open(SETTING_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _persist():
|
||
"""写回全部配置（保留 api_key / api_base 等字段）"""
|
||
data = _load_setting()
|
||
data.update({
|
||
"api_key": API_KEY,
|
||
"api_base": API_BASE,
|
||
"threshold": THRESHOLD,
|
||
"polling": POLLING_MINUTES,
|
||
"alert_enabled": ALERT_ENABLED,
|
||
})
|
||
_save_setting(data)
|
||
|
||
|
||
# 从持久化文件恢复上次设置（跨重启保留）
|
||
_s = _load_setting()
|
||
API_KEY = str(_s.get("api_key", ""))
|
||
API_BASE = str(_s.get("api_base", API_BASE)).rstrip("/")
|
||
THRESHOLD = int(_s.get("threshold", THRESHOLD))
|
||
POLLING_MINUTES = int(_s.get("polling", POLLING_MINUTES))
|
||
ALERT_ENABLED = bool(_s.get("alert_enabled", ALERT_ENABLED))
|
||
|
||
BALANCE_URL = f"{API_BASE}/user/balance"
|
||
|
||
_alerted = False
|
||
# 本次低余额事件中已成功投递的目标令牌 {"<channel>:<target_id>"}。
|
||
# 逐目标补投的依据：只重发不在其中的目标；余额恢复时清空。
|
||
_alert_delivered = set()
|
||
|
||
|
||
async def get_balance():
|
||
if not API_KEY:
|
||
print("[balance] setting.json 缺少 api_key，跳过查询")
|
||
return None
|
||
headers = {
|
||
"Authorization": f"Bearer {API_KEY}",
|
||
"Accept": "application/json"
|
||
}
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(BALANCE_URL, headers=headers, timeout=30) as resp:
|
||
if resp.status == 200:
|
||
data = await resp.json()
|
||
for info in data.get("balance_infos") or []:
|
||
if info.get("currency") == "CNY":
|
||
return float(info["total_balance"])
|
||
return None
|
||
text = await resp.text()
|
||
print(f"[DEBUG][balance] 查询余额失败，状态码 {resp.status}，响应体: {text[:200]}")
|
||
return None
|
||
except Exception as e:
|
||
print(f"[DEBUG][balance] 查询余额异常: {e}")
|
||
return None
|
||
|
||
|
||
async def check_balance_once():
|
||
global _alerted, _alert_delivered
|
||
if not ALERT_ENABLED:
|
||
return
|
||
balance = await get_balance()
|
||
if balance is None:
|
||
print("[DEBUG][balance] 本轮查询失败，跳过")
|
||
return
|
||
print(f"[DEBUG][balance] 轮询查询成功，当前余额 {balance} 元，阈值 {THRESHOLD} 元")
|
||
if balance < THRESHOLD:
|
||
if ALERT_ENABLED and not _alerted:
|
||
msg = (f"注意到 DeepSeek Platform 余额已低于阈值 （ {THRESHOLD} 元 ），请判别业务影响。\n"
|
||
f"当前余额 {balance} 元")
|
||
# 逐目标补投：按管理员身份扇出，本轮只对「尚未成功」的目标重发，
|
||
# 已成功的目标绝不重复推送（避免 TG 重复轰炸）。
|
||
pending_total, sent_now = 0, 0
|
||
for uid in identity.admin_uids():
|
||
tokens = {"%s:%s" % t: t for t in identity.push_targets(uid)}
|
||
# 目标被解绑/移除即视为消失：不在 tokens 里的历史令牌不再阻塞结案
|
||
todo = [t for token, t in tokens.items() if token not in _alert_delivered]
|
||
pending_total += len(todo)
|
||
if not todo:
|
||
continue
|
||
results = await push_to_user(uid, msg, tag="DeepSeek 余额告警", targets=todo)
|
||
for (channel_name, target_id), ok in results.items():
|
||
token = "%s:%s" % (channel_name, target_id)
|
||
sent_now += 1 if ok else 0
|
||
if ok:
|
||
_alert_delivered.add(token)
|
||
else:
|
||
print(f"[DEBUG][balance] 目标 {token} 本轮投递失败，下一轮补投")
|
||
# 结案判定：当前所有管理员的目标都已成功即结案；
|
||
# 目标集合为空（无管理员身份）时也正常结案，避免每轮空转刷「未结案」日志
|
||
remaining = 0
|
||
existing = 0
|
||
for uid in identity.admin_uids():
|
||
for t in identity.push_targets(uid):
|
||
existing += 1
|
||
if "%s:%s" % t not in _alert_delivered:
|
||
remaining += 1
|
||
if remaining == 0:
|
||
_alerted = True
|
||
if existing > 0:
|
||
print(f"[DEBUG][balance] 余额预警已结案（本事件补投 {sent_now} 个目标，"
|
||
f"累计成功 {existing} 个），当前余额 {balance} 元，阈值 {THRESHOLD} 元")
|
||
else:
|
||
print(f"[DEBUG][balance] 余额低于阈值但无管理员推送目标，直接结案，"
|
||
f"当前余额 {balance} 元，阈值 {THRESHOLD} 元")
|
||
else:
|
||
print(f"[DEBUG][balance] 余额预警未结案：本轮补投 {sent_now} 个成功 / "
|
||
f"{pending_total} 个待投，剩余 {remaining} 个目标未成功（共 {existing} 个），"
|
||
f"下一轮只补投未成功目标")
|
||
return
|
||
_alerted = True
|
||
else:
|
||
# 余额恢复：复位告警与逐目标投递记录
|
||
_alerted = False
|
||
_alert_delivered.clear()
|
||
|
||
|
||
async def monitor_loop():
|
||
print(f"[DEBUG][balance] 余额监控已启动，每 {POLLING_MINUTES} 分钟查询一次，阈值 {THRESHOLD} 元")
|
||
# 启动宽限：等 TG 通道 register 与 NapCat 反向连接完成后再做首次查询，
|
||
# 避免重启后的首次真实告警投递到空通道（那也是启动竞态吞告警的根因）。
|
||
await asyncio.sleep(MONITOR_STARTUP_DELAY)
|
||
while True:
|
||
try:
|
||
await check_balance_once()
|
||
except Exception as e:
|
||
print(f"[DEBUG][balance] 轮询异常: {e}")
|
||
await asyncio.sleep(POLLING_MINUTES * 60)
|
||
|
||
|
||
HELP_TEXT = (
|
||
"DeepSeek 余额监控指令\n"
|
||
"• .deepseek help 显示本帮助\n"
|
||
"• .deepseek balance 查询当前余额\n"
|
||
"• .deepseek threshold <int> 设置余额阈值（当前: {threshold} 元）\n"
|
||
"• .deepseek polling <int> 设置查询间隔（当前: {polling} 分钟）\n"
|
||
"• .deepseek alert <on|off> 开关余额监控（当前: {alert}）\n"
|
||
)
|
||
|
||
|
||
async def handle_command(channel, msg_type, target_id, raw_text):
|
||
global THRESHOLD, POLLING_MINUTES, ALERT_ENABLED, _alerted, _alert_delivered
|
||
parts = raw_text.split(maxsplit=2)
|
||
subcmd = parts[1].lower() if len(parts) >= 2 else "help"
|
||
if subcmd == "help":
|
||
await send_message(channel, msg_type, target_id, HELP_TEXT.format(
|
||
threshold=THRESHOLD, polling=POLLING_MINUTES,
|
||
alert="开" if ALERT_ENABLED else "关"))
|
||
return
|
||
if subcmd == "balance":
|
||
balance = await get_balance()
|
||
if balance is None:
|
||
await send_message(channel, msg_type, target_id, "余额查询失败，请稍后重试")
|
||
else:
|
||
await send_message(channel, msg_type, target_id, f"DeepSeek 账户当前余额：{balance:.2f} 元")
|
||
return
|
||
if subcmd == "threshold":
|
||
if len(parts) < 3:
|
||
await send_message(channel, msg_type, target_id, "用法: .deepseek threshold <int>，例如 .deepseek threshold 3")
|
||
return
|
||
try:
|
||
value = int(parts[2].strip())
|
||
except ValueError:
|
||
await send_message(channel, msg_type, target_id, "阈值必须为整数，例如 .deepseek threshold 3")
|
||
return
|
||
if value < 0:
|
||
await send_message(channel, msg_type, target_id, "阈值不能为负数")
|
||
return
|
||
THRESHOLD = value
|
||
# 阈值变更构成新的告警判定基准：复位结案标志与投递记录，
|
||
# 否则旧事件的 _alerted/_alert_delivered 会压制新阈值下的告警
|
||
_alerted = False
|
||
_alert_delivered.clear()
|
||
_persist()
|
||
await send_message(channel, msg_type, target_id, f"已设置余额阈值为 {THRESHOLD} 元")
|
||
return
|
||
if subcmd == "polling":
|
||
if len(parts) < 3:
|
||
await send_message(channel, msg_type, target_id, "用法: .deepseek polling <int>，例如 .deepseek polling 3")
|
||
return
|
||
try:
|
||
value = int(parts[2].strip())
|
||
except ValueError:
|
||
await send_message(channel, msg_type, target_id, "查询间隔必须为整数分钟，例如 .deepseek polling 3")
|
||
return
|
||
if value < 1:
|
||
await send_message(channel, msg_type, target_id, "查询间隔至少为 1 分钟")
|
||
return
|
||
POLLING_MINUTES = value
|
||
_persist()
|
||
await send_message(channel, msg_type, target_id, f"已设置查询间隔为 {POLLING_MINUTES} 分钟")
|
||
return
|
||
if subcmd == "alert":
|
||
if len(parts) < 3 or parts[2].strip().lower() not in ("on", "off"):
|
||
await send_message(channel, msg_type, target_id, "用法: .deepseek alert <on|off>")
|
||
return
|
||
ALERT_ENABLED = parts[2].strip().lower() == "on"
|
||
# 开关切换即开始新事件：投递记录一并清空，
|
||
# 否则上次事件残留的成功令牌会让重启后的告警被整体跳过（静默吞告警）
|
||
_alerted = False
|
||
_alert_delivered.clear()
|
||
_persist()
|
||
if ALERT_ENABLED:
|
||
await send_message(channel, msg_type, target_id, "已开启余额监控")
|
||
else:
|
||
await send_message(channel, msg_type, target_id, "已关闭余额监控（暂停轮询）")
|
||
return
|
||
await send_message(channel, msg_type, target_id, "未知子指令，使用 .deepseek help 查看指令范围")
|