GuangYu/QuestionProject
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement chat queue fallback

Browse Source
This commit is contained in:
Nelson
2026-07-08 16:57:40 +08:00
parent 10b6de0622
commit 6f3bbd3803
10 changed files with 367 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
ai_knowledge_base_v2/apps/backend/.env.example
View File

@@ -24,6 +24,9 @@ FEISHU_RETRY_COUNT=2
DEFAULT_DAILY_CHAT_LIMIT=100
DEFAULT_USER_NAME_PREFIX=用户
CHAT_MAX_ACTIVE_REQUESTS=2
CHAT_MAX_QUEUE_SIZE=20
CHAT_QUEUE_TIMEOUT_SECONDS=60
BOOTSTRAP_ADMIN_USERNAME=admin
BOOTSTRAP_ADMIN_PASSWORD=admin123456

70
ai_knowledge_base_v2/apps/backend/app/api/chat.py
View File

@@ -3,7 +3,7 @@ from __future__ import annotations
import json
from collections.abc import Iterator
from fastapi import APIRouter, Depends, Query
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
@@ -20,6 +20,7 @@ from app.schemas.chat import (
UpdateSessionTitleRequest,
)
from app.services.chat_service import ChatService
from app.services.chat_queue_service import chat_queue_manager, load_chat_queue_config
router = APIRouter()
@@ -72,8 +73,7 @@ def completions(
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> StreamingResponse:
answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message)
return StreamingResponse(_sse_chunks(answer), media_type="text/event-stream")
return StreamingResponse(_chat_stream(payload, db, current_user), media_type="text/event-stream")
@router.post("/stop")
@@ -86,9 +86,69 @@ def stop(
return api_success()
def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> Iterator[str]:
config = load_chat_queue_config(db)
queue_request = chat_queue_manager.request_slot(config)
acquired = queue_request.status == "acquired"
queued_token = queue_request.token
if queue_request.status == "rejected":
yield _sse_event(
"error",
message="当前请求过多，请稍后再试。",
activeCount=queue_request.active_count,
waitingCount=queue_request.waiting_count,
)
yield _sse_done()
return
if queue_request.status == "queued" and queued_token is not None:
yield _sse_event(
"queued",
message=f"当前请求较多，正在排队中，前方约 {max(queue_request.position - 1, 0)} 个请求。",
position=queue_request.position,
activeCount=queue_request.active_count,
waitingCount=queue_request.waiting_count,
)
queue_request = chat_queue_manager.wait_for_slot(queued_token, config)
acquired = queue_request.status == "acquired"
if queue_request.status == "timeout":
yield _sse_event("error", message="排队等待超时，请稍后再试。")
yield _sse_done()
return
try:
yield _sse_event(
"generating",
message="已进入生成队列，正在生成回答。",
activeCount=queue_request.active_count,
waitingCount=queue_request.waiting_count,
)
answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message)
yield from _sse_chunks(answer)
except HTTPException as exc:
yield _sse_event("error", message=str(exc.detail))
except Exception:
yield _sse_event("error", message="AI 回复生成失败，请稍后再试。")
finally:
if acquired:
chat_queue_manager.release_slot()
elif queued_token is not None:
chat_queue_manager.cancel_waiter(queued_token)
yield _sse_done()
def _sse_chunks(answer: str) -> Iterator[str]:
chunk_size = 12
for index in range(0, len(answer), chunk_size):
chunk = answer[index : index + chunk_size]
yield f"data: {json.dumps({'content': chunk}, ensure_ascii=False)}\n\n"
yield "data: [DONE]\n\n"
yield _sse_event("content", content=chunk)
def _sse_event(event_type: str, **payload: object) -> str:
return f"data: {json.dumps({'type': event_type, **payload}, ensure_ascii=False)}\n\n"
def _sse_done() -> str:
return "data: [DONE]\n\n"

3
ai_knowledge_base_v2/apps/backend/app/core/config.py
View File

@@ -45,6 +45,9 @@ class Settings(BaseSettings):
default_daily_chat_limit: int = 100
default_user_name_prefix: str = "用户"
chat_max_active_requests: int = 2
chat_max_queue_size: int = 20
chat_queue_timeout_seconds: int = 60
bootstrap_admin_username: str = "admin"
bootstrap_admin_password: str = "admin123456"
bootstrap_admin_name: str = "系统管理员"

146
ai_knowledge_base_v2/apps/backend/app/services/chat_queue_service.py Normal file
View File

@@ -0,0 +1,146 @@
from __future__ import annotations
from collections import deque
from dataclasses import dataclass
from threading import Condition
import time
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.core.config import get_settings
from app.models.ai_config import SystemConfig
@dataclass(frozen=True)
class ChatQueueConfig:
max_active_requests: int
max_queue_size: int
queue_timeout_seconds: int
@dataclass(frozen=True)
class QueueRequest:
status: str
token: object | None = None
position: int = 0
active_count: int = 0
waiting_count: int = 0
class ChatQueueManager:
def __init__(self) -> None:
self._condition = Condition()
self._active_count = 0
self._waiters: deque[object] = deque()
def request_slot(self, config: ChatQueueConfig) -> QueueRequest:
with self._condition:
if self._active_count < config.max_active_requests and not self._waiters:
self._active_count += 1
return QueueRequest(
status="acquired",
active_count=self._active_count,
waiting_count=len(self._waiters),
)
if len(self._waiters) >= config.max_queue_size:
return QueueRequest(
status="rejected",
active_count=self._active_count,
waiting_count=len(self._waiters),
)
token = object()
self._waiters.append(token)
return QueueRequest(
status="queued",
token=token,
position=len(self._waiters),
active_count=self._active_count,
waiting_count=len(self._waiters),
)
def wait_for_slot(self, token: object, config: ChatQueueConfig) -> QueueRequest:
deadline = time.monotonic() + config.queue_timeout_seconds
with self._condition:
while True:
if self._waiters and self._waiters[0] is token and self._active_count < config.max_active_requests:
self._waiters.popleft()
self._active_count += 1
self._condition.notify_all()
return QueueRequest(
status="acquired",
active_count=self._active_count,
waiting_count=len(self._waiters),
)
remaining = deadline - time.monotonic()
if remaining <= 0:
try:
self._waiters.remove(token)
except ValueError:
pass
self._condition.notify_all()
return QueueRequest(
status="timeout",
active_count=self._active_count,
waiting_count=len(self._waiters),
)
self._condition.wait(timeout=remaining)
def cancel_waiter(self, token: object) -> None:
with self._condition:
try:
self._waiters.remove(token)
except ValueError:
return
self._condition.notify_all()
def release_slot(self) -> None:
with self._condition:
if self._active_count > 0:
self._active_count -= 1
self._condition.notify_all()
chat_queue_manager = ChatQueueManager()
def load_chat_queue_config(db: Session) -> ChatQueueConfig:
settings = get_settings()
return ChatQueueConfig(
max_active_requests=_config_int(
db,
"chat_max_active_requests",
settings.chat_max_active_requests,
minimum=1,
maximum=1000,
),
max_queue_size=_config_int(
db,
"chat_max_queue_size",
settings.chat_max_queue_size,
minimum=0,
maximum=10000,
),
queue_timeout_seconds=_config_int(
db,
"chat_queue_timeout_seconds",
settings.chat_queue_timeout_seconds,
minimum=1,
maximum=3600,
),
)
def _config_int(db: Session, key: str, default: int, *, minimum: int, maximum: int) -> int:
config = db.scalar(select(SystemConfig).where(SystemConfig.config_key == key))
if config is None or not config.config_value.strip():
return default
try:
value = int(float(config.config_value.strip()))
except ValueError:
return default
return max(minimum, min(maximum, value))