forked from xiaohei/taiji-AI-PAD
392 lines
12 KiB
Python
392 lines
12 KiB
Python
"""
|
||
资源使用监控模块
|
||
采集和统计用户资源使用情况
|
||
|
||
注意：计费数据已迁移到新表：
|
||
- AgentBillingRecord: Agent 运行时计费
|
||
- ModelBillingRecord: 模型调用计费（LiteLLM）
|
||
旧的 BillingRecord 表已废弃，不再使用。
|
||
"""
|
||
|
||
from datetime import datetime, timedelta
|
||
from decimal import Decimal
|
||
from typing import Dict, List, Optional, Tuple
|
||
from sqlalchemy import select, func, and_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from models import ResourceUsage, User, Agent, Execution, AgentBillingRecord, ModelBillingRecord
|
||
|
||
|
||
async def record_resource_usage(
|
||
user_id: str,
|
||
agent_id: Optional[str],
|
||
cpu_seconds: float,
|
||
memory_mb_seconds: float,
|
||
network_bytes: int,
|
||
storage_bytes: int,
|
||
api_calls: int,
|
||
period_start: datetime,
|
||
period_end: datetime,
|
||
granularity: str,
|
||
db: AsyncSession
|
||
) -> ResourceUsage:
|
||
"""
|
||
记录资源使用情况
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
agent_id: Agent ID
|
||
cpu_seconds: CPU使用秒数
|
||
memory_mb_seconds: 内存使用 (MB*秒)
|
||
network_bytes: 网络流量字节
|
||
storage_bytes: 存储使用字节
|
||
api_calls: API调用次数
|
||
period_start: 周期开始时间
|
||
period_end: 周期结束时间
|
||
granularity: 粒度 (hourly, daily, monthly)
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
ResourceUsage记录
|
||
"""
|
||
# 检查是否已有该周期的记录
|
||
result = await db.execute(
|
||
select(ResourceUsage).where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.agent_id == agent_id if agent_id else True,
|
||
ResourceUsage.period_start == period_start,
|
||
ResourceUsage.granularity == granularity,
|
||
)
|
||
)
|
||
)
|
||
existing = result.scalar_one_or_none()
|
||
|
||
if existing:
|
||
# 更新现有记录
|
||
existing.cpu_seconds += cpu_seconds
|
||
existing.memory_mb_seconds += memory_mb_seconds
|
||
existing.network_bytes += network_bytes
|
||
existing.storage_bytes += storage_bytes
|
||
existing.api_calls += api_calls
|
||
usage = existing
|
||
else:
|
||
# 创建新记录
|
||
usage = ResourceUsage(
|
||
user_id=user_id,
|
||
agent_id=agent_id,
|
||
cpu_seconds=cpu_seconds,
|
||
memory_mb_seconds=memory_mb_seconds,
|
||
network_bytes=network_bytes,
|
||
storage_bytes=storage_bytes,
|
||
api_calls=api_calls,
|
||
period_start=period_start,
|
||
period_end=period_end,
|
||
granularity=granularity,
|
||
)
|
||
db.add(usage)
|
||
|
||
await db.commit()
|
||
await db.refresh(usage)
|
||
|
||
return usage
|
||
|
||
|
||
async def get_user_resource_summary(
|
||
user_id: str,
|
||
start_date: datetime,
|
||
end_date: datetime,
|
||
db: AsyncSession
|
||
) -> Dict:
|
||
"""
|
||
获取用户资源使用汇总
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
资源使用汇总
|
||
"""
|
||
result = await db.execute(
|
||
select(
|
||
func.sum(ResourceUsage.cpu_seconds).label("total_cpu_seconds"),
|
||
func.sum(ResourceUsage.memory_mb_seconds).label("total_memory_mb_seconds"),
|
||
func.sum(ResourceUsage.network_bytes).label("total_network_bytes"),
|
||
func.sum(ResourceUsage.storage_bytes).label("total_storage_bytes"),
|
||
func.sum(ResourceUsage.api_calls).label("total_api_calls"),
|
||
)
|
||
.where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.period_start >= start_date,
|
||
ResourceUsage.period_end <= end_date,
|
||
)
|
||
)
|
||
)
|
||
|
||
row = result.first()
|
||
|
||
return {
|
||
"totalCpuSeconds": float(row.total_cpu_seconds or 0),
|
||
"totalMemoryMbSeconds": float(row.total_memory_mb_seconds or 0),
|
||
"totalNetworkBytes": int(row.total_network_bytes or 0),
|
||
"totalStorageBytes": int(row.total_storage_bytes or 0),
|
||
"totalApiCalls": int(row.total_api_calls or 0),
|
||
"startDate": start_date.isoformat(),
|
||
"endDate": end_date.isoformat(),
|
||
}
|
||
|
||
|
||
async def get_resource_trends(
|
||
user_id: str,
|
||
period: str,
|
||
granularity: str,
|
||
db: AsyncSession
|
||
) -> List[Dict]:
|
||
"""
|
||
获取资源使用趋势
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
period: 时间范围 (7d, 30d, 90d)
|
||
granularity: 粒度 (hourly, daily)
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
资源使用趋势列表
|
||
"""
|
||
days_map = {"7d": 7, "30d": 30, "90d": 90}
|
||
days = days_map.get(period, 7)
|
||
start_date = datetime.utcnow() - timedelta(days=days)
|
||
|
||
result = await db.execute(
|
||
select(ResourceUsage)
|
||
.where(
|
||
and_(
|
||
ResourceUsage.user_id == user_id,
|
||
ResourceUsage.period_start >= start_date,
|
||
ResourceUsage.granularity == granularity,
|
||
)
|
||
)
|
||
.order_by(ResourceUsage.period_start)
|
||
)
|
||
|
||
usages = result.scalars().all()
|
||
|
||
return [
|
||
{
|
||
"periodStart": usage.period_start.isoformat(),
|
||
"periodEnd": usage.period_end.isoformat(),
|
||
"cpuSeconds": float(usage.cpu_seconds),
|
||
"memoryMbSeconds": float(usage.memory_mb_seconds),
|
||
"networkBytes": int(usage.network_bytes),
|
||
"apiCalls": int(usage.api_calls),
|
||
}
|
||
for usage in usages
|
||
]
|
||
|
||
|
||
async def get_agent_resource_stats(
|
||
agent_id: str,
|
||
start_date: datetime,
|
||
end_date: datetime,
|
||
db: AsyncSession
|
||
) -> Dict:
|
||
"""
|
||
获取Agent资源统计
|
||
|
||
Args:
|
||
agent_id: Agent ID
|
||
start_date: 开始日期
|
||
end_date: 结束日期
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
Agent资源统计
|
||
"""
|
||
# 执行次数和时间统计
|
||
exec_result = await db.execute(
|
||
select(
|
||
func.count(Execution.id).label("total_executions"),
|
||
func.avg(Execution.execution_time).label("avg_execution_time"),
|
||
func.sum(Execution.eu_consumed).label("total_eu"),
|
||
)
|
||
.where(
|
||
and_(
|
||
Execution.agent_id == agent_id,
|
||
Execution.started_at >= start_date,
|
||
Execution.started_at <= end_date,
|
||
)
|
||
)
|
||
)
|
||
|
||
exec_row = exec_result.first()
|
||
|
||
# 成功率
|
||
success_result = await db.execute(
|
||
select(func.count(Execution.id))
|
||
.where(
|
||
and_(
|
||
Execution.agent_id == agent_id,
|
||
Execution.started_at >= start_date,
|
||
Execution.status == "completed",
|
||
)
|
||
)
|
||
)
|
||
success_count = success_result.scalar() or 0
|
||
total_count = exec_row.total_executions or 1
|
||
|
||
return {
|
||
"agentId": agent_id,
|
||
"totalExecutions": int(exec_row.total_executions or 0),
|
||
"avgExecutionTime": float(exec_row.avg_execution_time or 0),
|
||
"totalEuConsumed": float(exec_row.total_eu or 0),
|
||
"successRate": round(success_count / total_count * 100, 2) if total_count > 0 else 0,
|
||
"startDate": start_date.isoformat(),
|
||
"endDate": end_date.isoformat(),
|
||
}
|
||
|
||
|
||
async def get_platform_resource_overview(db: AsyncSession) -> Dict:
|
||
"""
|
||
获取平台资源概览（管理员视图）
|
||
|
||
数据来源：
|
||
- AgentBillingRecord: Agent 运行时计费
|
||
- ModelBillingRecord: 模型调用计费（LiteLLM）
|
||
|
||
Args:
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
平台资源概览
|
||
"""
|
||
now = datetime.utcnow()
|
||
today_start = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||
month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||
|
||
# 今日调用次数（Agent + Model）
|
||
today_agent_calls = await db.execute(
|
||
select(func.count(AgentBillingRecord.id))
|
||
.where(AgentBillingRecord.start_time >= today_start)
|
||
)
|
||
today_model_calls = await db.execute(
|
||
select(func.count(ModelBillingRecord.id))
|
||
.where(ModelBillingRecord.created_at >= today_start)
|
||
)
|
||
today_calls_total = (today_agent_calls.scalar() or 0) + (today_model_calls.scalar() or 0)
|
||
|
||
# 本月调用次数（Agent + Model）
|
||
month_agent_calls = await db.execute(
|
||
select(func.count(AgentBillingRecord.id))
|
||
.where(AgentBillingRecord.start_time >= month_start)
|
||
)
|
||
month_model_calls = await db.execute(
|
||
select(func.count(ModelBillingRecord.id))
|
||
.where(ModelBillingRecord.created_at >= month_start)
|
||
)
|
||
month_calls_total = (month_agent_calls.scalar() or 0) + (month_model_calls.scalar() or 0)
|
||
|
||
# 活跃用户数（今日有 Agent 或 Model 调用的用户）
|
||
active_agent_users = await db.execute(
|
||
select(func.count(func.distinct(AgentBillingRecord.user_id)))
|
||
.where(AgentBillingRecord.start_time >= today_start)
|
||
)
|
||
active_model_users = await db.execute(
|
||
select(func.count(func.distinct(ModelBillingRecord.tenant_id)))
|
||
.where(ModelBillingRecord.created_at >= today_start)
|
||
)
|
||
# 注意：这里简单相加可能有重复，但作为概览统计可以接受
|
||
active_users_total = max(active_agent_users.scalar() or 0, active_model_users.scalar() or 0)
|
||
|
||
# 活跃Agent数
|
||
active_agents = await db.execute(
|
||
select(func.count(Agent.id))
|
||
.where(Agent.status == "active")
|
||
)
|
||
|
||
# 总EU消耗（Agent + Model）
|
||
agent_eu = await db.execute(
|
||
select(func.sum(AgentBillingRecord.eu_consumed))
|
||
.where(AgentBillingRecord.start_time >= month_start)
|
||
)
|
||
model_eu = await db.execute(
|
||
select(func.sum(ModelBillingRecord.eu_consumed))
|
||
.where(ModelBillingRecord.created_at >= month_start)
|
||
)
|
||
total_eu = (agent_eu.scalar() or 0) + (model_eu.scalar() or 0)
|
||
|
||
return {
|
||
"todayCalls": today_calls_total,
|
||
"monthCalls": month_calls_total,
|
||
"activeUsersToday": active_users_total,
|
||
"activeAgents": active_agents.scalar() or 0,
|
||
"monthTotalEu": int(total_eu),
|
||
"timestamp": now.isoformat(),
|
||
}
|
||
|
||
|
||
async def aggregate_hourly_usage(db: AsyncSession) -> int:
|
||
"""
|
||
聚合小时级资源使用数据（定时任务调用）
|
||
|
||
Args:
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
聚合的记录数
|
||
"""
|
||
# 获取上一小时的时间范围
|
||
now = datetime.utcnow()
|
||
period_end = now.replace(minute=0, second=0, microsecond=0)
|
||
period_start = period_end - timedelta(hours=1)
|
||
|
||
# 从执行记录聚合
|
||
result = await db.execute(
|
||
select(
|
||
Execution.agent_id,
|
||
func.count(Execution.id).label("api_calls"),
|
||
func.sum(Execution.execution_time).label("total_time"),
|
||
func.sum(Execution.cpu_usage).label("cpu_seconds"),
|
||
func.sum(Execution.memory_usage).label("memory_usage"),
|
||
)
|
||
.where(
|
||
and_(
|
||
Execution.started_at >= period_start,
|
||
Execution.started_at < period_end,
|
||
)
|
||
)
|
||
.group_by(Execution.agent_id)
|
||
)
|
||
|
||
count = 0
|
||
for row in result.all():
|
||
if row.agent_id:
|
||
# 获取Agent所有者
|
||
agent_result = await db.execute(
|
||
select(Agent.owner_id).where(Agent.id == row.agent_id)
|
||
)
|
||
owner = agent_result.scalar()
|
||
|
||
if owner:
|
||
await record_resource_usage(
|
||
user_id=str(owner),
|
||
agent_id=str(row.agent_id),
|
||
cpu_seconds=float(row.cpu_seconds or 0),
|
||
memory_mb_seconds=float(row.memory_usage or 0) * (float(row.total_time or 0) / 1000),
|
||
network_bytes=0,
|
||
storage_bytes=0,
|
||
api_calls=int(row.api_calls or 0),
|
||
period_start=period_start,
|
||
period_end=period_end,
|
||
granularity="hourly",
|
||
db=db,
|
||
)
|
||
count += 1
|
||
|
||
return count
|
||
|