From 5aca872129426c0f25cb5333df2ca695d63d7a94 Mon Sep 17 00:00:00 2001 From: Nelson <1475262689@qq.com> Date: Fri, 10 Jul 2026 11:45:37 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E8=81=8A=E5=A4=A9?= =?UTF-8?q?=E6=8E=92=E9=98=9F=E5=92=8C=E6=AF=8F=E6=97=A5=E9=A2=9D=E5=BA=A6?= =?UTF-8?q?=E9=87=8D=E7=BD=AE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../versions/0004_daily_chat_reset_date.py | 17 +++++++++++++++++ .../apps/backend/app/api/chat.py | 2 +- .../apps/backend/app/models/user.py | 5 +++-- .../apps/backend/app/services/chat_service.py | 16 ++++++++++++++++ .../backend/app/services/chat_stream_service.py | 2 ++ 5 files changed, 39 insertions(+), 3 deletions(-) create mode 100644 ai_knowledge_base_v2/apps/backend/alembic/versions/0004_daily_chat_reset_date.py diff --git a/ai_knowledge_base_v2/apps/backend/alembic/versions/0004_daily_chat_reset_date.py b/ai_knowledge_base_v2/apps/backend/alembic/versions/0004_daily_chat_reset_date.py new file mode 100644 index 0000000..5adbd64 --- /dev/null +++ b/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") diff --git a/ai_knowledge_base_v2/apps/backend/app/api/chat.py b/ai_knowledge_base_v2/apps/backend/app/api/chat.py index 0e626ad..4fff921 100644 --- a/ai_knowledge_base_v2/apps/backend/app/api/chat.py +++ b/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)} 个请求。", diff --git a/ai_knowledge_base_v2/apps/backend/app/models/user.py b/ai_knowledge_base_v2/apps/backend/app/models/user.py index 778ba68..eff88b0 100644 --- a/ai_knowledge_base_v2/apps/backend/app/models/user.py +++ b/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) diff --git a/ai_knowledge_base_v2/apps/backend/app/services/chat_service.py b/ai_knowledge_base_v2/apps/backend/app/services/chat_service.py index 8670d00..e890f60 100644 --- a/ai_knowledge_base_v2/apps/backend/app/services/chat_service.py +++ b/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) diff --git a/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py b/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py index be7dbc3..9e59fd9 100644 --- a/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py +++ b/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)