Implement chat queue fallback

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

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

ai_knowledge_base_v2/apps/backend/app/core/config.py

@@ -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 = "系统管理员"

ai_knowledge_base_v2/apps/backend/app/services/chat_queue_service.py

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