"""预算管理 API""" from datetime import datetime from typing import Optional from fastapi import APIRouter, Depends, Query from fastapi.responses import JSONResponse from sqlalchemy import and_, desc, select from sqlalchemy.ext.asyncio import AsyncSession from ..core.response import ErrorCode, error_response, success_response from ..database import get_db from ..dependencies import get_tenant_id, get_current_user, require_roles from ..models import Budget, BudgetAlert, User from ..schemas.budget import ( BudgetCheckRequest, BudgetCreateRequest, BudgetOverrideRequest, BudgetUpdateRequest, ) from ..services.budget_enforcement_service import BudgetEnforcementService router = APIRouter(prefix="/budget", tags=["budget"]) def _period_window(now: datetime, period: str) -> tuple[datetime, datetime]: p = (period or "monthly").lower() if p == "yearly": return datetime(now.year, 1, 1), datetime(now.year + 1, 1, 1) if p == "quarterly": q = (now.month - 1) // 3 start_month = q * 3 + 1 start = datetime(now.year, start_month, 1) end = datetime(now.year + 1, 1, 1) if start_month == 10 else datetime(now.year, start_month + 3, 1) return start, end start = datetime(now.year, now.month, 1) end = datetime(now.year + 1, 1, 1) if now.month == 12 else datetime(now.year, now.month + 1, 1) return start, end @router.get("/") async def list_budgets( budget_type: Optional[str] = Query(default=None, pattern="^(tenant|project)$"), project_id: Optional[int] = Query(default=None), current_user: User = Depends(get_current_user), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): _ = current_user query = select(Budget).where(Budget.tenant_id == tenant_id) if budget_type: query = query.where(Budget.budget_type == budget_type) if project_id is not None: query = query.where(Budget.project_id == project_id) query = query.order_by(desc(Budget.created_at)) result = await db.execute(query) rows = result.scalars().all() return success_response( data=[ { "id": b.id, "budget_type": b.budget_type, "project_id": b.project_id, "period": b.period, "budget_amount_usd": float(b.budget_amount_usd), "alert_threshold_percent": float(b.alert_threshold_percent), "hard_limit_enabled": bool(b.hard_limit_enabled), "current_usage_usd": float(b.current_usage_usd), "status": b.status, "current_period_start": b.current_period_start.isoformat() if b.current_period_start else None, "current_period_end": b.current_period_end.isoformat() if b.current_period_end else None, } for b in rows ] ) @router.post("/") async def create_budget( request: BudgetCreateRequest, current_user: User = Depends(require_roles("BOSS", "ADMIN")), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): _ = current_user now = datetime.utcnow() period_start, period_end = _period_window(now, request.period) new_budget = Budget( tenant_id=tenant_id, budget_type=request.budget_type, project_id=request.project_id, period=request.period, budget_amount_usd=request.budget_amount_usd, alert_threshold_percent=request.alert_threshold_percent, hard_limit_enabled=request.hard_limit_enabled, current_usage_usd=0.0, current_period_start=period_start, current_period_end=period_end, status="active", override_history={"records": []}, ) db.add(new_budget) await db.commit() await db.refresh(new_budget) return success_response( data={ "id": new_budget.id, "budget_type": new_budget.budget_type, "project_id": new_budget.project_id, "period": new_budget.period, "budget_amount_usd": float(new_budget.budget_amount_usd), "alert_threshold_percent": float(new_budget.alert_threshold_percent), "hard_limit_enabled": bool(new_budget.hard_limit_enabled), "status": new_budget.status, } ) @router.patch("/{budget_id}") async def update_budget( budget_id: int, request: BudgetUpdateRequest, current_user: User = Depends(require_roles("BOSS", "ADMIN")), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): _ = current_user result = await db.execute( select(Budget).where(and_(Budget.id == budget_id, Budget.tenant_id == tenant_id)) ) budget = result.scalar_one_or_none() if not budget: return JSONResponse(status_code=404, content=error_response(ErrorCode.NOT_FOUND, "预算不存在")) if request.period is not None: budget.period = request.period if request.budget_amount_usd is not None: budget.budget_amount_usd = request.budget_amount_usd if request.alert_threshold_percent is not None: budget.alert_threshold_percent = request.alert_threshold_percent if request.hard_limit_enabled is not None: budget.hard_limit_enabled = request.hard_limit_enabled if request.status is not None: budget.status = request.status await db.commit() await db.refresh(budget) return success_response( data={ "id": budget.id, "period": budget.period, "budget_amount_usd": float(budget.budget_amount_usd), "alert_threshold_percent": float(budget.alert_threshold_percent), "hard_limit_enabled": bool(budget.hard_limit_enabled), "status": budget.status, } ) @router.post("/{budget_id}/override") async def override_budget( budget_id: int, request: BudgetOverrideRequest, current_user: User = Depends(require_roles("BOSS", "ADMIN")), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): service = BudgetEnforcementService(db) try: data = await service.override_budget_limit( budget_id=budget_id, tenant_id=tenant_id, operator_user_id=int(current_user.id), note=request.note, duration_hours=request.duration_hours, ) return success_response(data=data) except ValueError as exc: return JSONResponse(status_code=404, content=error_response(ErrorCode.NOT_FOUND, str(exc))) @router.get("/alerts") async def list_budget_alerts( budget_id: Optional[int] = Query(default=None), unread_only: bool = Query(default=False), limit: int = Query(default=100, ge=1, le=500), current_user: User = Depends(get_current_user), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): _ = current_user query = select(BudgetAlert).where(BudgetAlert.tenant_id == tenant_id) if budget_id is not None: query = query.where(BudgetAlert.budget_id == budget_id) if unread_only: query = query.where(BudgetAlert.is_read.is_(False)) query = query.order_by(desc(BudgetAlert.created_at)).limit(limit) result = await db.execute(query) rows = result.scalars().all() return success_response( data=[ { "id": a.id, "budget_id": a.budget_id, "alert_type": a.alert_type, "usage_usd": float(a.usage_usd), "usage_percent": float(a.usage_percent), "message": a.message, "is_read": bool(a.is_read), "action_taken": a.action_taken, "status": a.status, "created_at": a.created_at.isoformat() if a.created_at else None, } for a in rows ] ) @router.post("/check") async def check_budget( request: BudgetCheckRequest, current_user: User = Depends(get_current_user), tenant_id: int = Depends(get_tenant_id), db: AsyncSession = Depends(get_db), ): _ = current_user service = BudgetEnforcementService(db) result = await service.check_budget_before_ai_call( tenant_id=tenant_id, estimated_cost_usd=request.estimated_cost_usd, project_id=request.project_id, ) return success_response(data=result)