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

This commit is contained in:
2026-07-10 11:45:37 +08:00
parent 5819830e33
commit 5aca872129
5 changed files with 39 additions and 3 deletions

View File

@@ -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)} 个请求。",

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)