Implement chat queue fallback

2026-07-08 16:57:40 +08:00
parent 10b6de0622
commit 6f3bbd3803
10 changed files with 367 additions and 31 deletions
--- a/ai_knowledge_base_v2/apps/backend/app/services/chat_queue_service.py
+++ b/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))