forked from xiaohei/taiji-AI-PAD
421 lines
15 KiB
Python
421 lines
15 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
计费系统健康检查脚本
|
||
|
||
功能：
|
||
1. 检测假用量问题（配额数据与实际运行Agent不一致）
|
||
2. 检查余额不足用户
|
||
3. 测试余额不足自动停止功能
|
||
4. 生成详细健康报告
|
||
|
||
使用方法：
|
||
在mcp-server容器内运行：
|
||
|
||
# 完整检查
|
||
python check_billing_health.py
|
||
|
||
# 只检查假用量
|
||
python check_billing_health.py --check-quota-only
|
||
|
||
# 只检查余额
|
||
python check_billing_health.py --check-balance-only
|
||
|
||
# 测试自动停止（需要指定用户ID）
|
||
python check_billing_health.py --test-auto-stop <user_id>
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import argparse
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, List, Any, Optional
|
||
from decimal import Decimal
|
||
|
||
from sqlalchemy import select, func, and_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from database import get_db
|
||
from models import (
|
||
Agent,
|
||
TenantCustomAgentQuota,
|
||
AgentBillingRecord,
|
||
Balance,
|
||
User
|
||
)
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BillingHealthChecker:
|
||
"""计费系统健康检查器"""
|
||
|
||
def __init__(self):
|
||
self.report_file = Path("/app/logs/billing_health_report.json")
|
||
self.report_file.parent.mkdir(exist_ok=True)
|
||
|
||
async def check_quota_consistency(self, db: AsyncSession) -> Dict[str, Any]:
|
||
"""
|
||
检查配额一致性（假用量检测）
|
||
|
||
对比：
|
||
- 配额表中的使用量
|
||
- 实际运行的Agent数量和资源
|
||
"""
|
||
logger.info("=" * 60)
|
||
logger.info("🔍 开始检查配额一致性（假用量检测）")
|
||
logger.info("=" * 60)
|
||
|
||
# 获取所有配额记录
|
||
result = await db.execute(select(TenantCustomAgentQuota))
|
||
quotas = result.scalars().all()
|
||
|
||
inconsistencies = []
|
||
total_fake_cpu = 0.0
|
||
total_fake_memory = 0.0
|
||
total_fake_agents = 0
|
||
|
||
for quota in quotas:
|
||
tenant_id = str(quota.tenant_id)
|
||
|
||
# 查询配额使用量
|
||
quota_usage = {
|
||
'cpu_used': float(quota.cpu_used or 0),
|
||
'memory_used': float(quota.memory_used or 0),
|
||
'agent_count': quota.agent_count or 0
|
||
}
|
||
|
||
# 查询实际运行的Agent
|
||
real_agents_result = await db.execute(
|
||
select(Agent)
|
||
.where(Agent.owner_id == tenant_id)
|
||
.where(Agent.type == 'custom')
|
||
.where(Agent.status == 'active')
|
||
)
|
||
real_agents = real_agents_result.scalars().all()
|
||
|
||
# 计算实际使用量
|
||
real_usage = {
|
||
'cpu_used': sum(float(a.cpu or 0) for a in real_agents),
|
||
'memory_used': sum(float(a.memory or 0) for a in real_agents),
|
||
'agent_count': len(real_agents)
|
||
}
|
||
|
||
# 计算差异
|
||
cpu_diff = quota_usage['cpu_used'] - real_usage['cpu_used']
|
||
memory_diff = quota_usage['memory_used'] - real_usage['memory_used']
|
||
agent_diff = quota_usage['agent_count'] - real_usage['agent_count']
|
||
|
||
# 允许小误差（0.01核，0.01GB）
|
||
has_inconsistency = (
|
||
abs(cpu_diff) > 0.01 or
|
||
abs(memory_diff) > 0.01 or
|
||
agent_diff != 0
|
||
)
|
||
|
||
if has_inconsistency:
|
||
inconsistency = {
|
||
'tenant_id': tenant_id,
|
||
'quota_usage': quota_usage,
|
||
'real_usage': real_usage,
|
||
'diff': {
|
||
'cpu': cpu_diff,
|
||
'memory': memory_diff,
|
||
'agent_count': agent_diff
|
||
},
|
||
'severity': 'high' if (cpu_diff > 1 or memory_diff > 1 or agent_diff > 2) else 'medium'
|
||
}
|
||
inconsistencies.append(inconsistency)
|
||
|
||
total_fake_cpu += cpu_diff
|
||
total_fake_memory += memory_diff
|
||
total_fake_agents += agent_diff
|
||
|
||
logger.warning(
|
||
f"⚠️ 检测到配额不一致: 租户={tenant_id}\n"
|
||
f" 配额显示: CPU={quota_usage['cpu_used']:.2f}核, "
|
||
f"内存={quota_usage['memory_used']:.2f}GB, Agent={quota_usage['agent_count']}个\n"
|
||
f" 实际运行: CPU={real_usage['cpu_used']:.2f}核, "
|
||
f"内存={real_usage['memory_used']:.2f}GB, Agent={real_usage['agent_count']}个\n"
|
||
f" 假用量: CPU={cpu_diff:.2f}核, "
|
||
f"内存={memory_diff:.2f}GB, Agent={agent_diff}个"
|
||
)
|
||
|
||
summary = {
|
||
'total_tenants': len(quotas),
|
||
'inconsistent_tenants': len(inconsistencies),
|
||
'consistency_rate': (len(quotas) - len(inconsistencies)) / len(quotas) * 100 if quotas else 100,
|
||
'total_fake_cpu': total_fake_cpu,
|
||
'total_fake_memory': total_fake_memory,
|
||
'total_fake_agents': total_fake_agents,
|
||
'details': inconsistencies
|
||
}
|
||
|
||
if inconsistencies:
|
||
logger.error(
|
||
f"❌ 发现 {len(inconsistencies)} 个租户存在配额不一致问题！\n"
|
||
f" 假用量汇总: CPU={total_fake_cpu:.2f}核, "
|
||
f"内存={total_fake_memory:.2f}GB, Agent={total_fake_agents}个\n"
|
||
f" 建议运行: python fix_fake_quota.py"
|
||
)
|
||
else:
|
||
logger.info("✅ 所有租户配额数据一致，无假用量问题")
|
||
|
||
return summary
|
||
|
||
async def check_balance_status(self, db: AsyncSession) -> Dict[str, Any]:
|
||
"""
|
||
检查用户余额状态
|
||
|
||
识别：
|
||
- 余额不足的用户
|
||
- 透支用户
|
||
- 有运行Agent但余额不足的用户（风险）
|
||
"""
|
||
logger.info("=" * 60)
|
||
logger.info("💰 开始检查用户余额状态")
|
||
logger.info("=" * 60)
|
||
|
||
# 查询所有用户及其余额
|
||
result = await db.execute(
|
||
select(User, Balance)
|
||
.outerjoin(Balance, User.id == Balance.user_id)
|
||
)
|
||
users_data = result.all()
|
||
|
||
low_balance_users = []
|
||
overdraft_users = []
|
||
risky_users = [] # 有运行Agent但余额不足
|
||
|
||
for user, balance in users_data:
|
||
user_id = str(user.id)
|
||
eu_balance = float(balance.eu_balance) if balance else 0.0
|
||
credit_limit = float(user.credit_limit or 0)
|
||
available = eu_balance + credit_limit
|
||
|
||
# 检查是否有运行中的Agent
|
||
running_agents_result = await db.execute(
|
||
select(func.count(AgentBillingRecord.id))
|
||
.where(AgentBillingRecord.user_id == user_id)
|
||
.where(AgentBillingRecord.end_time == None)
|
||
)
|
||
running_agent_count = running_agents_result.scalar() or 0
|
||
|
||
user_info = {
|
||
'user_id': user_id,
|
||
'email': user.email,
|
||
'eu_balance': eu_balance,
|
||
'credit_limit': credit_limit,
|
||
'available_balance': available,
|
||
'running_agents': running_agent_count
|
||
}
|
||
|
||
# 透支（可用余额为负）
|
||
if available < 0:
|
||
overdraft_users.append(user_info)
|
||
logger.error(
|
||
f"💥 透支用户: {user.email} (ID: {user_id})\n"
|
||
f" 账户余额: {eu_balance:.2f} EU\n"
|
||
f" 授信额度: {credit_limit:.2f} EU\n"
|
||
f" 可用余额: {available:.2f} EU\n"
|
||
f" 运行Agent数: {running_agent_count}"
|
||
)
|
||
|
||
# 余额不足（低于10 EU）
|
||
elif available < 10:
|
||
low_balance_users.append(user_info)
|
||
logger.warning(
|
||
f"⚠️ 余额不足: {user.email} (ID: {user_id})\n"
|
||
f" 可用余额: {available:.2f} EU\n"
|
||
f" 运行Agent数: {running_agent_count}"
|
||
)
|
||
|
||
# 风险用户：有运行Agent但余额很低（< 5 EU）
|
||
if running_agent_count > 0 and available < 5:
|
||
risky_users.append(user_info)
|
||
logger.warning(
|
||
f"🚨 风险用户: {user.email} (ID: {user_id})\n"
|
||
f" 有 {running_agent_count} 个Agent运行但余额仅剩 {available:.2f} EU"
|
||
)
|
||
|
||
summary = {
|
||
'total_users': len(users_data),
|
||
'overdraft_users': len(overdraft_users),
|
||
'low_balance_users': len(low_balance_users),
|
||
'risky_users': len(risky_users),
|
||
'overdraft_details': overdraft_users,
|
||
'low_balance_details': low_balance_users,
|
||
'risky_users_details': risky_users
|
||
}
|
||
|
||
logger.info(
|
||
f"\n📊 余额状态汇总:\n"
|
||
f" 总用户数: {len(users_data)}\n"
|
||
f" 透支用户: {len(overdraft_users)}\n"
|
||
f" 余额不足: {len(low_balance_users)}\n"
|
||
f" 风险用户: {len(risky_users)}"
|
||
)
|
||
|
||
return summary
|
||
|
||
async def test_auto_stop(self, db: AsyncSession, user_id: str) -> Dict[str, Any]:
|
||
"""
|
||
测试余额不足自动停止功能
|
||
|
||
模拟周期计费检测到余额不足时的行为
|
||
"""
|
||
logger.info("=" * 60)
|
||
logger.info(f"🧪 测试自动停止功能: user_id={user_id}")
|
||
logger.info("=" * 60)
|
||
|
||
from app.billing import get_available_balance
|
||
from app.periodic_billing import stop_user_agents
|
||
|
||
# 检查用户余额
|
||
balance, credit_limit, available = await get_available_balance(user_id, db)
|
||
|
||
logger.info(
|
||
f"用户余额状态:\n"
|
||
f" 账户余额: {balance:.2f} EU\n"
|
||
f" 授信额度: {credit_limit:.2f} EU\n"
|
||
f" 可用余额: {available:.2f} EU"
|
||
)
|
||
|
||
# 检查运行中的Agent
|
||
running_agents_result = await db.execute(
|
||
select(AgentBillingRecord)
|
||
.where(AgentBillingRecord.user_id == user_id)
|
||
.where(AgentBillingRecord.end_time == None)
|
||
)
|
||
running_agents = running_agents_result.scalars().all()
|
||
|
||
logger.info(f"运行中的Agent数量: {len(running_agents)}")
|
||
for agent in running_agents:
|
||
logger.info(f" - {agent.agent_name} (启动时间: {agent.start_time})")
|
||
|
||
result = {
|
||
'user_id': user_id,
|
||
'balance': float(balance),
|
||
'credit_limit': float(credit_limit),
|
||
'available': float(available),
|
||
'running_agents_before': len(running_agents),
|
||
'stopped_agents': [],
|
||
'action_taken': 'none'
|
||
}
|
||
|
||
# 如果余额不足，执行停止
|
||
if available < 0:
|
||
logger.warning(f"⚠️ 可用余额为负 ({available:.2f} EU)，将停止所有Agent...")
|
||
|
||
stopped = await stop_user_agents(user_id, db)
|
||
result['stopped_agents'] = stopped
|
||
result['action_taken'] = 'stopped'
|
||
|
||
logger.info(
|
||
f"✅ 已停止 {len(stopped)} 个Agent:\n" +
|
||
"\n".join(f" - {name}" for name in stopped)
|
||
)
|
||
else:
|
||
logger.info(f"✅ 余额充足 ({available:.2f} EU)，无需停止Agent")
|
||
result['action_taken'] = 'no_action_needed'
|
||
|
||
return result
|
||
|
||
async def generate_full_report(self, db: AsyncSession) -> Dict[str, Any]:
|
||
"""生成完整的健康检查报告"""
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("📋 生成计费系统健康检查报告")
|
||
logger.info("=" * 60 + "\n")
|
||
|
||
report = {
|
||
'check_time': datetime.utcnow().isoformat(),
|
||
'quota_consistency': await self.check_quota_consistency(db),
|
||
'balance_status': await self.check_balance_status(db)
|
||
}
|
||
|
||
# 生成健康评分
|
||
quota_ok = report['quota_consistency']['inconsistent_tenants'] == 0
|
||
no_overdraft = report['balance_status']['overdraft_users'] == 0
|
||
few_risky = report['balance_status']['risky_users'] < 5
|
||
|
||
health_score = 0
|
||
if quota_ok:
|
||
health_score += 40
|
||
if no_overdraft:
|
||
health_score += 40
|
||
if few_risky:
|
||
health_score += 20
|
||
|
||
report['health_score'] = health_score
|
||
report['health_status'] = (
|
||
'healthy' if health_score >= 90 else
|
||
'warning' if health_score >= 70 else
|
||
'critical'
|
||
)
|
||
|
||
# 生成建议
|
||
recommendations = []
|
||
if not quota_ok:
|
||
recommendations.append("运行 fix_fake_quota.py 修复配额不一致问题")
|
||
if not no_overdraft:
|
||
recommendations.append(f"有 {report['balance_status']['overdraft_users']} 个用户透支，建议停止其Agent或充值")
|
||
if not few_risky:
|
||
recommendations.append(f"有 {report['balance_status']['risky_users']} 个风险用户，建议提醒充值")
|
||
|
||
report['recommendations'] = recommendations
|
||
|
||
# 保存报告
|
||
with open(self.report_file, 'w', encoding='utf-8') as f:
|
||
json.dump(report, f, indent=2, ensure_ascii=False)
|
||
|
||
logger.info(f"\n{'=' * 60}")
|
||
logger.info(f"📊 健康检查完成")
|
||
logger.info(f"{'=' * 60}")
|
||
logger.info(f"健康评分: {health_score}/100 ({report['health_status'].upper()})")
|
||
logger.info(f"报告保存: {self.report_file}")
|
||
|
||
if recommendations:
|
||
logger.info(f"\n💡 建议:")
|
||
for i, rec in enumerate(recommendations, 1):
|
||
logger.info(f" {i}. {rec}")
|
||
|
||
return report
|
||
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
parser = argparse.ArgumentParser(description='计费系统健康检查')
|
||
parser.add_argument('--check-quota-only', action='store_true', help='只检查配额一致性')
|
||
parser.add_argument('--check-balance-only', action='store_true', help='只检查余额状态')
|
||
parser.add_argument('--test-auto-stop', type=str, metavar='USER_ID', help='测试自动停止功能（指定用户ID）')
|
||
|
||
args = parser.parse_args()
|
||
|
||
checker = BillingHealthChecker()
|
||
|
||
async for db in get_db():
|
||
try:
|
||
if args.check_quota_only:
|
||
await checker.check_quota_consistency(db)
|
||
elif args.check_balance_only:
|
||
await checker.check_balance_status(db)
|
||
elif args.test_auto_stop:
|
||
await checker.test_auto_stop(db, args.test_auto_stop)
|
||
else:
|
||
await checker.generate_full_report(db)
|
||
|
||
except Exception as e:
|
||
logger.error(f"执行失败: {str(e)}", exc_info=True)
|
||
finally:
|
||
break
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|