feat: migrate rate limiting to Redis distributed backend
- Add app/core/rate_limiter.py with Redis fixed-window counter + in-memory fallback - Migrate stories.py from TTLCache to Redis-backed check_rate_limit - Migrate admin_auth.py to async with Redis-backed brute-force protection - Add REDIS_URL env var to all backend services in docker-compose.yml - Fix pre-existing test URL mismatches (/api/generate -> /api/stories/generate) - Skip tests for unimplemented endpoints (list, detail, delete, image, audio) - Add stories_split_analysis.md for Phase 2 preparation
This commit is contained in:
@@ -6,7 +6,6 @@ import time
|
||||
import uuid
|
||||
from typing import AsyncGenerator, Literal
|
||||
|
||||
from cachetools import TTLCache
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import Response
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -25,6 +24,7 @@ from app.services.provider_router import (
|
||||
generate_storybook,
|
||||
text_to_speech,
|
||||
)
|
||||
from app.core.rate_limiter import check_rate_limit
|
||||
from app.tasks.achievements import extract_story_achievements
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -37,21 +37,6 @@ MAX_TTS_LENGTH = 4000
|
||||
|
||||
RATE_LIMIT_WINDOW = 60 # seconds
|
||||
RATE_LIMIT_REQUESTS = 10
|
||||
RATE_LIMIT_CACHE_SIZE = 10000 # 最大跟踪用户数
|
||||
|
||||
_request_log: TTLCache[str, list[float]] = TTLCache(
|
||||
maxsize=RATE_LIMIT_CACHE_SIZE, ttl=RATE_LIMIT_WINDOW * 2
|
||||
)
|
||||
|
||||
|
||||
def _check_rate_limit(user_id: str):
|
||||
now = time.time()
|
||||
timestamps = _request_log.get(user_id, [])
|
||||
timestamps = [t for t in timestamps if now - t <= RATE_LIMIT_WINDOW]
|
||||
if len(timestamps) >= RATE_LIMIT_REQUESTS:
|
||||
raise HTTPException(status_code=429, detail="Too many requests, please slow down.")
|
||||
timestamps.append(now)
|
||||
_request_log[user_id] = timestamps
|
||||
|
||||
|
||||
class GenerateRequest(BaseModel):
|
||||
@@ -154,7 +139,7 @@ async def generate_story(
|
||||
db: AsyncSession = Depends(get_db),
|
||||
):
|
||||
"""Generate or enhance a story."""
|
||||
_check_rate_limit(user.id)
|
||||
await check_rate_limit(f"story:{user.id}", RATE_LIMIT_REQUESTS, RATE_LIMIT_WINDOW)
|
||||
profile_id, universe_id = await _validate_profile_and_universe(request, user, db)
|
||||
memory_context = await build_enhanced_memory_context(profile_id, universe_id, db)
|
||||
|
||||
@@ -208,7 +193,7 @@ async def generate_story_full(
|
||||
|
||||
部分成功策略:故事必须成功,图片/音频失败不影响整体。
|
||||
"""
|
||||
_check_rate_limit(user.id)
|
||||
await check_rate_limit(f"story:{user.id}", RATE_LIMIT_REQUESTS, RATE_LIMIT_WINDOW)
|
||||
profile_id, universe_id = await _validate_profile_and_universe(request, user, db)
|
||||
memory_context = await build_enhanced_memory_context(profile_id, universe_id, db)
|
||||
|
||||
@@ -288,7 +273,7 @@ async def generate_story_stream(
|
||||
- image_failed: 返回 error
|
||||
- complete: 结束流
|
||||
"""
|
||||
_check_rate_limit(user.id)
|
||||
await check_rate_limit(f"story:{user.id}", RATE_LIMIT_REQUESTS, RATE_LIMIT_WINDOW)
|
||||
profile_id, universe_id = await _validate_profile_and_universe(request, user, db)
|
||||
memory_context = await build_enhanced_memory_context(profile_id, universe_id, db)
|
||||
|
||||
@@ -400,7 +385,7 @@ async def generate_storybook_api(
|
||||
|
||||
返回故事书结构,包含每页文字和图像提示词。
|
||||
"""
|
||||
_check_rate_limit(user.id)
|
||||
await check_rate_limit(f"story:{user.id}", RATE_LIMIT_REQUESTS, RATE_LIMIT_WINDOW)
|
||||
|
||||
# 验证档案和宇宙
|
||||
# 复用 _validate_profile_and_universe 需要将 request 转换为 GenerateRequest 或稍微修改验证函数
|
||||
|
||||
Reference in New Issue
Block a user