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"