fix: 修复聊天排队和每日额度重置

ai_knowledge_base_v2/apps/backend/app/services/chat_service.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from datetime import UTC, datetime
 from time import perf_counter
+from zoneinfo import ZoneInfo
 
 from fastapi import HTTPException, status
 from sqlalchemy import select
@@ -70,6 +71,7 @@ class ChatService:
 
     @staticmethod
     def create_answer(db: Session, user: User, session_id: int, question: str) -> str:
+        user = ChatService.prepare_daily_quota(db, user)
         session = ChatService._get_user_session(db, user, session_id)
         ChatService._ensure_quota(user)
 
@@ -241,6 +243,20 @@ class ChatService:
         if user.daily_chat_used >= user.daily_chat_limit:
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="今日提问次数已用完")
 
+    @staticmethod
+    def prepare_daily_quota(db: Session, user: User) -> User:
+        locked_user = db.scalar(select(User).where(User.id == user.id).with_for_update())
+        if locked_user is None:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="用户不存在")
+
+        today = datetime.now(ZoneInfo("Asia/Shanghai")).date()
+        if locked_user.daily_chat_reset_date != today:
+            locked_user.daily_chat_used = 0
+            locked_user.daily_chat_reset_date = today
+            db.add(locked_user)
+            db.flush()
+        return locked_user
+
 
 def _now() -> datetime:
     return datetime.now(UTC).replace(tzinfo=None)

ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py

@@ -23,6 +23,7 @@ from app.services.rag_service import RagService
 class ChatStreamService:
     @staticmethod
     def stream_answer(db: Session, user: User, session_id: int, question: str) -> Iterator[str]:
+        user = ChatService.prepare_daily_quota(db, user)
         session = ChatService._get_user_session(db, user, session_id)
         ChatService._ensure_quota(user)
 
@@ -145,6 +146,7 @@ class ChatStreamService:
 
     @staticmethod
     async def stream_answer_async(db: Session, user: User, session_id: int, question: str) -> AsyncIterator[str]:
+        user = ChatService.prepare_daily_quota(db, user)
         session = ChatService._get_user_session(db, user, session_id)
         ChatService._ensure_quota(user)