fix: 修复聊天排队和每日额度重置
This commit is contained in:
@@ -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)} 个请求。",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user