120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
|
"""
|
|||
|
|
用量追踪中间件：每日对话次数限制与强制执行
|
|||
|
|
"""
|
|||
|
|
import logging
|
|||
|
|
from datetime import date
|
|||
|
|
from fastapi import Request
|
|||
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|||
|
|
from starlette.responses import JSONResponse
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
|
|||
|
|
from app.core.database import SessionLocal
|
|||
|
|
from app.core.security import decode_access_token
|
|||
|
|
from app.models.user import User
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# 套餐每日限额
|
|||
|
|
TIER_DAILY_LIMITS = {
|
|||
|
|
"free": 50,
|
|||
|
|
"pro": 500,
|
|||
|
|
"enterprise": 0, # 0 = 无限
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 需要计数的路径前缀
|
|||
|
|
COUNTED_PATH_PREFIXES = [
|
|||
|
|
"/api/v1/agent-chat/",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 不计数路径（导出、会话管理等非对话操作）
|
|||
|
|
EXCLUDED_PATH_PATTERNS = [
|
|||
|
|
"/sessions/",
|
|||
|
|
"/search-messages",
|
|||
|
|
"/export",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class UsageMiddleware(BaseHTTPMiddleware):
|
|||
|
|
"""用量追踪中间件：检查每日限额 → 放行 → 递增计数"""
|
|||
|
|
|
|||
|
|
async def dispatch(self, request: Request, call_next):
|
|||
|
|
path = request.url.path
|
|||
|
|
should_count = (
|
|||
|
|
request.method == "POST"
|
|||
|
|
and any(path.startswith(p) for p in COUNTED_PATH_PREFIXES)
|
|||
|
|
and not any(p in path for p in EXCLUDED_PATH_PATTERNS)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not should_count:
|
|||
|
|
return await call_next(request)
|
|||
|
|
|
|||
|
|
# 从 Authorization header 解码 user_id
|
|||
|
|
user_id = None
|
|||
|
|
auth_header = request.headers.get("Authorization", "")
|
|||
|
|
if auth_header.startswith("Bearer "):
|
|||
|
|
token = auth_header[7:]
|
|||
|
|
try:
|
|||
|
|
payload = decode_access_token(token)
|
|||
|
|
user_id = payload.get("sub")
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
if not user_id:
|
|||
|
|
return await call_next(request)
|
|||
|
|
|
|||
|
|
db: Session = SessionLocal()
|
|||
|
|
try:
|
|||
|
|
user = db.query(User).filter(User.id == user_id).first()
|
|||
|
|
if not user:
|
|||
|
|
return await call_next(request)
|
|||
|
|
|
|||
|
|
today = date.today()
|
|||
|
|
|
|||
|
|
# 日期重置
|
|||
|
|
if user.daily_usage_date and user.daily_usage_date < today:
|
|||
|
|
user.daily_usage_count = 0
|
|||
|
|
user.daily_usage_date = today
|
|||
|
|
if not user.daily_usage_date:
|
|||
|
|
user.daily_usage_date = today
|
|||
|
|
|
|||
|
|
tier = user.subscription_tier or "free"
|
|||
|
|
limit = TIER_DAILY_LIMITS.get(tier, 50)
|
|||
|
|
current = user.daily_usage_count or 0
|
|||
|
|
|
|||
|
|
# 检查是否已达限额
|
|||
|
|
if limit > 0 and current >= limit:
|
|||
|
|
db.close()
|
|||
|
|
logger.warning("用户 %s 已达每日限额 %s/%s，拒绝请求", user_id, current, limit)
|
|||
|
|
return JSONResponse(
|
|||
|
|
status_code=429,
|
|||
|
|
content={
|
|||
|
|
"detail": f"今日对话次数已用完（{current}/{limit}）。升级 Pro 版享受每日 {TIER_DAILY_LIMITS.get('pro', 500)} 次对话。",
|
|||
|
|
"daily_used": current,
|
|||
|
|
"daily_limit": limit,
|
|||
|
|
"upgrade_required": True,
|
|||
|
|
},
|
|||
|
|
headers={"X-Upgrade-Required": "true"},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 未达限额：先放行，再递增
|
|||
|
|
response = await call_next(request)
|
|||
|
|
|
|||
|
|
if response.status_code < 400:
|
|||
|
|
# 递增计数（需要重新查询，因为 request 可能耗时较长）
|
|||
|
|
db_user = db.query(User).filter(User.id == user_id).first()
|
|||
|
|
if db_user:
|
|||
|
|
db_user.daily_usage_count = (db_user.daily_usage_count or 0) + 1
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
new_count = db_user.daily_usage_count
|
|||
|
|
if limit > 0 and new_count >= limit:
|
|||
|
|
logger.info("用户 %s 已达每日限额 %s/%s", user_id, new_count, limit)
|
|||
|
|
|
|||
|
|
return response
|
|||
|
|
except Exception as e:
|
|||
|
|
db.rollback()
|
|||
|
|
logger.error("用量追踪失败: %s", e)
|
|||
|
|
return await call_next(request)
|
|||
|
|
finally:
|
|||
|
|
db.close()
|