Initial commit: clean project structure
- Backend: FastAPI + SQLAlchemy + Celery (Python 3.11+) - Frontend: Vue 3 + TypeScript + Pinia + Tailwind - Admin Frontend: separate Vue 3 app for management - Docker Compose: 9 services orchestration - Specs: design prototypes, memory system PRD, product roadmap Cleanup performed: - Removed temporary debug scripts from backend root - Removed deprecated admin_app.py (embedded UI) - Removed duplicate docs from admin-frontend - Updated .gitignore for Vite cache and egg-info
This commit is contained in:
zhangtuo
196
backend/app/services/cost_tracker.py
Normal file
196
backend/app/services/cost_tracker.py
Normal file
@@ -0,0 +1,196 @@
|
||||
"""成本追踪服务。
|
||||
|
||||
记录 API 调用成本，支持预算控制。
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from decimal import Decimal
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.logging import get_logger
|
||||
from app.db.admin_models import CostRecord, UserBudget
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class BudgetExceededError(Exception):
|
||||
"""预算超限错误。"""
|
||||
|
||||
def __init__(self, limit_type: str, used: Decimal, limit: Decimal):
|
||||
self.limit_type = limit_type
|
||||
self.used = used
|
||||
self.limit = limit
|
||||
super().__init__(f"{limit_type} 预算已超限: {used}/{limit} USD")
|
||||
|
||||
|
||||
class CostTracker:
|
||||
"""成本追踪器。"""
|
||||
|
||||
async def record_cost(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
user_id: str,
|
||||
provider_name: str,
|
||||
capability: str,
|
||||
estimated_cost: float,
|
||||
provider_id: str | None = None,
|
||||
) -> CostRecord:
|
||||
"""记录一次 API 调用成本。"""
|
||||
record = CostRecord(
|
||||
user_id=user_id,
|
||||
provider_id=provider_id,
|
||||
provider_name=provider_name,
|
||||
capability=capability,
|
||||
estimated_cost=Decimal(str(estimated_cost)),
|
||||
)
|
||||
db.add(record)
|
||||
await db.commit()
|
||||
|
||||
logger.debug(
|
||||
"cost_recorded",
|
||||
user_id=user_id,
|
||||
provider=provider_name,
|
||||
capability=capability,
|
||||
cost=estimated_cost,
|
||||
)
|
||||
return record
|
||||
|
||||
async def get_daily_cost(self, db: AsyncSession, user_id: str) -> Decimal:
|
||||
"""获取用户今日成本。"""
|
||||
today_start = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
result = await db.execute(
|
||||
select(func.sum(CostRecord.estimated_cost)).where(
|
||||
CostRecord.user_id == user_id,
|
||||
CostRecord.timestamp >= today_start,
|
||||
)
|
||||
)
|
||||
total = result.scalar()
|
||||
return Decimal(str(total)) if total else Decimal("0")
|
||||
|
||||
async def get_monthly_cost(self, db: AsyncSession, user_id: str) -> Decimal:
|
||||
"""获取用户本月成本。"""
|
||||
now = datetime.utcnow()
|
||||
month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
result = await db.execute(
|
||||
select(func.sum(CostRecord.estimated_cost)).where(
|
||||
CostRecord.user_id == user_id,
|
||||
CostRecord.timestamp >= month_start,
|
||||
)
|
||||
)
|
||||
total = result.scalar()
|
||||
return Decimal(str(total)) if total else Decimal("0")
|
||||
|
||||
async def get_cost_by_capability(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
user_id: str,
|
||||
days: int = 30,
|
||||
) -> dict[str, Decimal]:
|
||||
"""按能力类型统计成本。"""
|
||||
since = datetime.utcnow() - timedelta(days=days)
|
||||
|
||||
result = await db.execute(
|
||||
select(CostRecord.capability, func.sum(CostRecord.estimated_cost))
|
||||
.where(CostRecord.user_id == user_id, CostRecord.timestamp >= since)
|
||||
.group_by(CostRecord.capability)
|
||||
)
|
||||
return {row[0]: Decimal(str(row[1])) for row in result.all()}
|
||||
|
||||
async def check_budget(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
user_id: str,
|
||||
estimated_cost: float,
|
||||
) -> bool:
|
||||
"""检查预算是否允许此次调用。
|
||||
|
||||
Returns:
|
||||
True 如果允许，否则抛出 BudgetExceededError
|
||||
"""
|
||||
budget = await self.get_user_budget(db, user_id)
|
||||
if not budget or not budget.enabled:
|
||||
return True
|
||||
|
||||
# 检查日预算
|
||||
daily_cost = await self.get_daily_cost(db, user_id)
|
||||
if daily_cost + Decimal(str(estimated_cost)) > budget.daily_limit_usd:
|
||||
raise BudgetExceededError("日", daily_cost, budget.daily_limit_usd)
|
||||
|
||||
# 检查月预算
|
||||
monthly_cost = await self.get_monthly_cost(db, user_id)
|
||||
if monthly_cost + Decimal(str(estimated_cost)) > budget.monthly_limit_usd:
|
||||
raise BudgetExceededError("月", monthly_cost, budget.monthly_limit_usd)
|
||||
|
||||
return True
|
||||
|
||||
async def get_user_budget(self, db: AsyncSession, user_id: str) -> UserBudget | None:
|
||||
"""获取用户预算配置。"""
|
||||
result = await db.execute(
|
||||
select(UserBudget).where(UserBudget.user_id == user_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def set_user_budget(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
user_id: str,
|
||||
daily_limit: float | None = None,
|
||||
monthly_limit: float | None = None,
|
||||
alert_threshold: float | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> UserBudget:
|
||||
"""设置用户预算。"""
|
||||
budget = await self.get_user_budget(db, user_id)
|
||||
|
||||
if budget is None:
|
||||
budget = UserBudget(user_id=user_id)
|
||||
db.add(budget)
|
||||
|
||||
if daily_limit is not None:
|
||||
budget.daily_limit_usd = Decimal(str(daily_limit))
|
||||
if monthly_limit is not None:
|
||||
budget.monthly_limit_usd = Decimal(str(monthly_limit))
|
||||
if alert_threshold is not None:
|
||||
budget.alert_threshold = Decimal(str(alert_threshold))
|
||||
if enabled is not None:
|
||||
budget.enabled = enabled
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(budget)
|
||||
return budget
|
||||
|
||||
async def get_cost_summary(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
user_id: str,
|
||||
) -> dict:
|
||||
"""获取用户成本摘要。"""
|
||||
daily = await self.get_daily_cost(db, user_id)
|
||||
monthly = await self.get_monthly_cost(db, user_id)
|
||||
by_capability = await self.get_cost_by_capability(db, user_id)
|
||||
budget = await self.get_user_budget(db, user_id)
|
||||
|
||||
return {
|
||||
"daily_cost_usd": float(daily),
|
||||
"monthly_cost_usd": float(monthly),
|
||||
"by_capability": {k: float(v) for k, v in by_capability.items()},
|
||||
"budget": {
|
||||
"daily_limit_usd": float(budget.daily_limit_usd) if budget else None,
|
||||
"monthly_limit_usd": float(budget.monthly_limit_usd) if budget else None,
|
||||
"daily_usage_percent": float(daily / budget.daily_limit_usd * 100)
|
||||
if budget and budget.daily_limit_usd
|
||||
else None,
|
||||
"monthly_usage_percent": float(monthly / budget.monthly_limit_usd * 100)
|
||||
if budget and budget.monthly_limit_usd
|
||||
else None,
|
||||
"enabled": budget.enabled if budget else False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# 全局单例
|
||||
cost_tracker = CostTracker()
|
||||
Reference in New Issue
Block a user