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

ai_knowledge_base_v2/apps/backend/alembic/versions/0004_daily_chat_reset_date.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from alembic import op
+import sqlalchemy as sa
+
+revision = "0004_daily_chat_reset_date"
+down_revision = "0003_ai_log_retrieved_chunks"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column("sys_user", sa.Column("daily_chat_reset_date", sa.Date(), nullable=True))
+
+
+def downgrade() -> None:
+    op.drop_column("sys_user", "daily_chat_reset_date")

ai_knowledge_base_v2/apps/backend/app/api/chat.py

@@ -113,7 +113,7 @@ async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user
         yield _sse_done()
         return
 
-    if queue_request.status == "queued" and queued_token is not None:
+    if queue_request.status == "queued" and queue_request.token is not None:
         yield _sse_event(
             "queued",
             message=f"当前请求较多，正在排队中，前方约 {max(queue_request.position - 1, 0)} 个请求。",

ai_knowledge_base_v2/apps/backend/app/models/user.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import date, datetime
 
-from sqlalchemy import BigInteger, DateTime, Integer, String
+from sqlalchemy import BigInteger, Date, DateTime, Integer, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from app.models.base import Base, SoftDeleteMixin, TimestampMixin
@@ -18,6 +18,7 @@ class User(Base, TimestampMixin, SoftDeleteMixin):
     status: Mapped[int] = mapped_column(default=1, nullable=False)
     daily_chat_limit: Mapped[int] = mapped_column(Integer, default=100, nullable=False)
     daily_chat_used: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    daily_chat_reset_date: Mapped[date | None] = mapped_column(Date, nullable=True)
     effective_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
     expired_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
     last_login_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)

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)