238 lines
8.3 KiB
Python
238 lines
8.3 KiB
Python
"""预算管理 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
|
|
|
|
from ..utils.timezone import now_cn_naive
|
|
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 = now_cn_naive()
|
|
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)
|
|
|