6c6c48cf71
All checks were successful
continuous-integration/drone/push Build is passing
- 新增告警模块 (alerts): 告警规则配置与触发 - 新增成本管理模块 (cost): 成本统计与分析 - 新增配额模块 (quota): 配额管理与限制 - 新增微信模块 (wechat): 微信相关功能接口 - 新增缓存服务 (cache): Redis 缓存封装 - 新增请求日志中间件 (request_logger) - 新增异常处理和链路追踪中间件 - 更新 dashboard 前端展示 - 更新 SDK stats_client 功能
265 lines
7.9 KiB
Python
265 lines
7.9 KiB
Python
"""配额管理路由"""
|
|
from typing import Optional, Dict, Any
|
|
from datetime import date
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import desc
|
|
|
|
from ..database import get_db
|
|
from ..models.tenant import Subscription
|
|
from ..services.quota import QuotaService
|
|
from .auth import get_current_user, require_operator
|
|
from ..models.user import User
|
|
|
|
router = APIRouter(prefix="/quota", tags=["配额管理"])
|
|
|
|
|
|
# ============= Schemas =============
|
|
|
|
class QuotaConfigUpdate(BaseModel):
|
|
daily_calls: int = 0
|
|
daily_tokens: int = 0
|
|
monthly_calls: int = 0
|
|
monthly_tokens: int = 0
|
|
monthly_cost: float = 0
|
|
concurrent_calls: int = 0
|
|
|
|
|
|
class SubscriptionCreate(BaseModel):
|
|
tenant_id: str
|
|
app_code: str
|
|
start_date: Optional[str] = None
|
|
end_date: Optional[str] = None
|
|
quota: QuotaConfigUpdate
|
|
|
|
|
|
class SubscriptionUpdate(BaseModel):
|
|
start_date: Optional[str] = None
|
|
end_date: Optional[str] = None
|
|
quota: Optional[QuotaConfigUpdate] = None
|
|
status: Optional[str] = None
|
|
|
|
|
|
# ============= Quota Check API =============
|
|
|
|
@router.get("/check")
|
|
async def check_quota(
|
|
tenant_id: str = Query(..., alias="tid"),
|
|
app_code: str = Query(..., alias="aid"),
|
|
estimated_tokens: int = Query(0),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""检查配额是否足够
|
|
|
|
用于调用前检查，返回是否允许继续调用
|
|
"""
|
|
service = QuotaService(db)
|
|
result = service.check_quota(tenant_id, app_code, estimated_tokens)
|
|
|
|
return {
|
|
"allowed": result.allowed,
|
|
"reason": result.reason,
|
|
"quota_type": result.quota_type,
|
|
"limit": result.limit,
|
|
"used": result.used,
|
|
"remaining": result.remaining
|
|
}
|
|
|
|
|
|
@router.get("/summary")
|
|
async def get_quota_summary(
|
|
tenant_id: str = Query(...),
|
|
app_code: str = Query(...),
|
|
user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""获取配额使用汇总"""
|
|
service = QuotaService(db)
|
|
return service.get_quota_summary(tenant_id, app_code)
|
|
|
|
|
|
@router.get("/usage")
|
|
async def get_quota_usage(
|
|
tenant_id: str = Query(...),
|
|
app_code: str = Query(...),
|
|
user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""获取配额使用情况"""
|
|
service = QuotaService(db)
|
|
usage = service.get_usage(tenant_id, app_code)
|
|
|
|
return {
|
|
"daily_calls": usage.daily_calls,
|
|
"daily_tokens": usage.daily_tokens,
|
|
"monthly_calls": usage.monthly_calls,
|
|
"monthly_tokens": usage.monthly_tokens,
|
|
"monthly_cost": round(usage.monthly_cost, 2)
|
|
}
|
|
|
|
|
|
# ============= Subscription API =============
|
|
|
|
@router.get("/subscriptions")
|
|
async def list_subscriptions(
|
|
page: int = Query(1, ge=1),
|
|
size: int = Query(20, ge=1, le=100),
|
|
tenant_id: Optional[str] = None,
|
|
app_code: Optional[str] = None,
|
|
status: Optional[str] = None,
|
|
user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""获取订阅列表"""
|
|
query = db.query(Subscription)
|
|
|
|
if tenant_id:
|
|
query = query.filter(Subscription.tenant_id == tenant_id)
|
|
if app_code:
|
|
query = query.filter(Subscription.app_code == app_code)
|
|
if status:
|
|
query = query.filter(Subscription.status == status)
|
|
|
|
total = query.count()
|
|
items = query.order_by(desc(Subscription.created_at)).offset((page - 1) * size).limit(size).all()
|
|
|
|
return {
|
|
"total": total,
|
|
"page": page,
|
|
"size": size,
|
|
"items": [format_subscription(s) for s in items]
|
|
}
|
|
|
|
|
|
@router.get("/subscriptions/{subscription_id}")
|
|
async def get_subscription(
|
|
subscription_id: int,
|
|
user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""获取订阅详情"""
|
|
subscription = db.query(Subscription).filter(Subscription.id == subscription_id).first()
|
|
if not subscription:
|
|
raise HTTPException(status_code=404, detail="订阅不存在")
|
|
return format_subscription(subscription)
|
|
|
|
|
|
@router.post("/subscriptions")
|
|
async def create_subscription(
|
|
data: SubscriptionCreate,
|
|
user: User = Depends(require_operator),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""创建订阅"""
|
|
# 检查是否已存在
|
|
existing = db.query(Subscription).filter(
|
|
Subscription.tenant_id == data.tenant_id,
|
|
Subscription.app_code == data.app_code,
|
|
Subscription.status == 'active'
|
|
).first()
|
|
|
|
if existing:
|
|
raise HTTPException(status_code=400, detail="该租户应用已有活跃订阅")
|
|
|
|
subscription = Subscription(
|
|
tenant_id=data.tenant_id,
|
|
app_code=data.app_code,
|
|
start_date=data.start_date or date.today(),
|
|
end_date=data.end_date,
|
|
quota=data.quota.model_dump() if data.quota else {},
|
|
status='active'
|
|
)
|
|
db.add(subscription)
|
|
db.commit()
|
|
db.refresh(subscription)
|
|
|
|
return {"success": True, "id": subscription.id}
|
|
|
|
|
|
@router.put("/subscriptions/{subscription_id}")
|
|
async def update_subscription(
|
|
subscription_id: int,
|
|
data: SubscriptionUpdate,
|
|
user: User = Depends(require_operator),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""更新订阅"""
|
|
subscription = db.query(Subscription).filter(Subscription.id == subscription_id).first()
|
|
if not subscription:
|
|
raise HTTPException(status_code=404, detail="订阅不存在")
|
|
|
|
if data.start_date:
|
|
subscription.start_date = data.start_date
|
|
if data.end_date:
|
|
subscription.end_date = data.end_date
|
|
if data.quota:
|
|
subscription.quota = data.quota.model_dump()
|
|
if data.status:
|
|
subscription.status = data.status
|
|
|
|
db.commit()
|
|
|
|
# 清除缓存
|
|
service = QuotaService(db)
|
|
cache_key = f"quota:config:{subscription.tenant_id}:{subscription.app_code}"
|
|
service._cache.delete(cache_key)
|
|
|
|
return {"success": True}
|
|
|
|
|
|
@router.delete("/subscriptions/{subscription_id}")
|
|
async def delete_subscription(
|
|
subscription_id: int,
|
|
user: User = Depends(require_operator),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""删除订阅"""
|
|
subscription = db.query(Subscription).filter(Subscription.id == subscription_id).first()
|
|
if not subscription:
|
|
raise HTTPException(status_code=404, detail="订阅不存在")
|
|
|
|
db.delete(subscription)
|
|
db.commit()
|
|
|
|
return {"success": True}
|
|
|
|
|
|
@router.put("/subscriptions/{subscription_id}/quota")
|
|
async def update_quota(
|
|
subscription_id: int,
|
|
data: QuotaConfigUpdate,
|
|
user: User = Depends(require_operator),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
"""更新订阅配额"""
|
|
subscription = db.query(Subscription).filter(Subscription.id == subscription_id).first()
|
|
if not subscription:
|
|
raise HTTPException(status_code=404, detail="订阅不存在")
|
|
|
|
subscription.quota = data.model_dump()
|
|
db.commit()
|
|
|
|
# 清除缓存
|
|
service = QuotaService(db)
|
|
cache_key = f"quota:config:{subscription.tenant_id}:{subscription.app_code}"
|
|
service._cache.delete(cache_key)
|
|
|
|
return {"success": True}
|
|
|
|
|
|
# ============= Helper Functions =============
|
|
|
|
def format_subscription(subscription: Subscription) -> dict:
|
|
return {
|
|
"id": subscription.id,
|
|
"tenant_id": subscription.tenant_id,
|
|
"app_code": subscription.app_code,
|
|
"start_date": str(subscription.start_date) if subscription.start_date else None,
|
|
"end_date": str(subscription.end_date) if subscription.end_date else None,
|
|
"quota": subscription.quota or {},
|
|
"status": subscription.status,
|
|
"created_at": subscription.created_at,
|
|
"updated_at": subscription.updated_at
|
|
}
|