Add Redis-backed chat queue

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

@@ -21,7 +21,13 @@ 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
+from app.services.chat_queue_runtime import (
+    cancel_chat_waiter,
+    release_chat_slot,
+    request_chat_slot,
+    wait_for_chat_slot,
+)
+from app.services.chat_queue_service import load_chat_queue_config
 from app.services.chat_stream_service import ChatStreamService
 
 router = APIRouter()
@@ -90,9 +96,10 @@ def stop(
 
 async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> AsyncIterator[str]:
     config = load_chat_queue_config(db)
-    queue_request = chat_queue_manager.request_slot(config)
+    queue_lease = await request_chat_slot(config)
+    queue_request = queue_lease.request
     acquired = queue_request.status == "acquired"
-    queued_token = queue_request.token
+    should_cancel_waiter = queue_request.status == "queued"
 
     if queue_request.status == "rejected":
         yield _sse_event(
@@ -112,12 +119,27 @@ async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user
             activeCount=queue_request.active_count,
             waitingCount=queue_request.waiting_count,
         )
-        queue_request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, queued_token, config)
+        queue_lease = await wait_for_chat_slot(queue_lease, config)
+        queue_request = queue_lease.request
         acquired = queue_request.status == "acquired"
+        should_cancel_waiter = queue_request.status == "queued"
+        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 == "timeout":
             yield _sse_event("error", message="排队等待超时，请稍后再试。")
             yield _sse_done()
             return
+        if queue_request.status == "canceled":
+            yield _sse_event("error", message="排队请求已取消，请稍后再试。")
+            yield _sse_done()
+            return
 
     try:
         yield _sse_event(
@@ -136,9 +158,9 @@ async def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user
         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)
+            await release_chat_slot(queue_lease)
+        elif should_cancel_waiter:
+            await cancel_chat_waiter(queue_lease)
 
     yield _sse_done()
 

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

@@ -25,6 +25,7 @@ class Settings(BaseSettings):
     )
 
     database_url: str = "mysql+pymysql://ai_kb:ai_kb@127.0.0.1:3306/ai_knowledge_base_v2?charset=utf8mb4"
+    redis_url: str = ""
     auto_create_tables: bool = False
 
     jwt_secret_key: str = "change-me-in-production"
@@ -48,6 +49,7 @@ class Settings(BaseSettings):
     chat_max_active_requests: int = 2
     chat_max_queue_size: int = 20
     chat_queue_timeout_seconds: int = 60
+    chat_active_lease_seconds: int = 900
     bootstrap_admin_username: str = "admin"
     bootstrap_admin_password: str = "admin123456"
     bootstrap_admin_name: str = "系统管理员"

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

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from uuid import uuid4
+
+from redis.asyncio import Redis
+
+from app.services.chat_queue_service import ChatQueueConfig, QueueRequest
+from app.services.redis_client import get_redis_client
+
+
+class RedisChatQueueManager:
+    _active_key = "ai_kb:chat_queue:active"
+    _waiting_key = "ai_kb:chat_queue:waiting"
+    _request_script = """
+local active_key = KEYS[1]
+local waiting_key = KEYS[2]
+local now = tonumber(ARGV[1])
+local token = ARGV[2]
+local max_active = tonumber(ARGV[3])
+local max_waiting = tonumber(ARGV[4])
+local active_expire_at = tonumber(ARGV[5])
+local key_ttl = tonumber(ARGV[6])
+
+redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
+local active_count = redis.call('ZCARD', active_key)
+local waiting_count = redis.call('LLEN', waiting_key)
+
+if active_count < max_active and waiting_count == 0 then
+  redis.call('ZADD', active_key, active_expire_at, token)
+  redis.call('EXPIRE', active_key, key_ttl)
+  return {'acquired', token, 0, active_count + 1, 0}
+end
+
+if waiting_count >= max_waiting then
+  redis.call('EXPIRE', active_key, key_ttl)
+  redis.call('EXPIRE', waiting_key, key_ttl)
+  return {'rejected', token, 0, active_count, waiting_count}
+end
+
+redis.call('RPUSH', waiting_key, token)
+redis.call('EXPIRE', waiting_key, key_ttl)
+redis.call('EXPIRE', active_key, key_ttl)
+return {'queued', token, waiting_count + 1, active_count, waiting_count + 1}
+"""
+    _promote_script = """
+local active_key = KEYS[1]
+local waiting_key = KEYS[2]
+local now = tonumber(ARGV[1])
+local token = ARGV[2]
+local max_active = tonumber(ARGV[3])
+local active_expire_at = tonumber(ARGV[4])
+local key_ttl = tonumber(ARGV[5])
+
+redis.call('ZREMRANGEBYSCORE', active_key, '-inf', now)
+local active_count = redis.call('ZCARD', active_key)
+local waiting_count = redis.call('LLEN', waiting_key)
+local head = redis.call('LINDEX', waiting_key, 0)
+
+if head == token and active_count < max_active then
+  redis.call('LPOP', waiting_key)
+  redis.call('ZADD', active_key, active_expire_at, token)
+  redis.call('EXPIRE', active_key, key_ttl)
+  redis.call('EXPIRE', waiting_key, key_ttl)
+  return {'acquired', token, 0, active_count + 1, waiting_count - 1}
+end
+
+local position = redis.call('LPOS', waiting_key, token)
+if not position then
+  return {'canceled', token, 0, active_count, waiting_count}
+end
+
+redis.call('EXPIRE', active_key, key_ttl)
+redis.call('EXPIRE', waiting_key, key_ttl)
+return {'queued', token, position + 1, active_count, waiting_count}
+"""
+
+    async def request_slot(self, config: ChatQueueConfig) -> QueueRequest:
+        client = _require_client()
+        now = time.time()
+        token = uuid4().hex
+        result = await client.eval(
+            self._request_script,
+            2,
+            self._active_key,
+            self._waiting_key,
+            now,
+            token,
+            config.max_active_requests,
+            config.max_queue_size,
+            now + config.active_lease_seconds,
+            _key_ttl(config),
+        )
+        return _queue_request(result)
+
+    async def wait_for_slot(self, token: str, config: ChatQueueConfig) -> QueueRequest:
+        client = _require_client()
+        deadline = time.monotonic() + config.queue_timeout_seconds
+        while True:
+            now = time.time()
+            result = await client.eval(
+                self._promote_script,
+                2,
+                self._active_key,
+                self._waiting_key,
+                now,
+                token,
+                config.max_active_requests,
+                now + config.active_lease_seconds,
+                _key_ttl(config),
+            )
+            request = _queue_request(result)
+            if request.status in {"acquired", "canceled"}:
+                return request
+
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                await self.cancel_waiter(token)
+                return QueueRequest(status="timeout", token=token)
+            await asyncio.sleep(min(0.25, remaining))
+
+    async def cancel_waiter(self, token: str) -> None:
+        client = _require_client()
+        await client.lrem(self._waiting_key, 0, token)
+
+    async def release_slot(self, token: str) -> None:
+        client = _require_client()
+        await client.zrem(self._active_key, token)
+
+
+redis_chat_queue_manager = RedisChatQueueManager()
+
+
+def _require_client() -> Redis:
+    client = get_redis_client()
+    if client is None:
+        raise RuntimeError("Redis is not configured")
+    return client
+
+
+def _queue_request(result: list) -> QueueRequest:
+    status = str(result[0])
+    token = str(result[1]) if len(result) > 1 and result[1] else None
+    return QueueRequest(
+        status=status,
+        token=token,
+        position=int(result[2]) if len(result) > 2 else 0,
+        active_count=int(result[3]) if len(result) > 3 else 0,
+        waiting_count=int(result[4]) if len(result) > 4 else 0,
+    )
+
+
+def _key_ttl(config: ChatQueueConfig) -> int:
+    return max(config.active_lease_seconds, config.queue_timeout_seconds) + 300

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

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+
+from app.services.chat_queue_redis_service import redis_chat_queue_manager
+from app.services.chat_queue_service import ChatQueueConfig, QueueRequest, chat_queue_manager
+from app.services.redis_client import ping_redis
+
+
+@dataclass
+class ChatQueueLease:
+    backend: str
+    request: QueueRequest
+
+
+async def request_chat_slot(config: ChatQueueConfig) -> ChatQueueLease:
+    if await ping_redis():
+        try:
+            return ChatQueueLease("redis", await redis_chat_queue_manager.request_slot(config))
+        except Exception:
+            pass
+    request = await asyncio.to_thread(chat_queue_manager.request_slot, config)
+    return ChatQueueLease("memory", request)
+
+
+async def wait_for_chat_slot(lease: ChatQueueLease, config: ChatQueueConfig) -> ChatQueueLease:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            lease.request = await redis_chat_queue_manager.wait_for_slot(token, config)
+            return lease
+        except Exception:
+            await cancel_chat_waiter(lease)
+            lease.backend = "memory"
+            lease.request = await asyncio.to_thread(chat_queue_manager.request_slot, config)
+            if lease.request.status == "queued":
+                lease.request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, lease.request.token, config)
+            return lease
+
+    lease.request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, token, config)
+    return lease
+
+
+async def cancel_chat_waiter(lease: ChatQueueLease) -> None:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            await redis_chat_queue_manager.cancel_waiter(token)
+        except Exception:
+            pass
+        return
+    if token is not None:
+        await asyncio.to_thread(chat_queue_manager.cancel_waiter, token)
+
+
+async def release_chat_slot(lease: ChatQueueLease) -> None:
+    token = lease.request.token
+    if lease.backend == "redis" and isinstance(token, str):
+        try:
+            await redis_chat_queue_manager.release_slot(token)
+        except Exception:
+            pass
+        return
+    await asyncio.to_thread(chat_queue_manager.release_slot)

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

@@ -17,6 +17,7 @@ class ChatQueueConfig:
     max_active_requests: int
     max_queue_size: int
     queue_timeout_seconds: int
+    active_lease_seconds: int
 
 
 @dataclass(frozen=True)
@@ -132,6 +133,13 @@ def load_chat_queue_config(db: Session) -> ChatQueueConfig:
             minimum=1,
             maximum=3600,
         ),
+        active_lease_seconds=_config_int(
+            db,
+            "chat_active_lease_seconds",
+            settings.chat_active_lease_seconds,
+            minimum=60,
+            maximum=86400,
+        ),
     )
 
 

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

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+from redis.asyncio import Redis
+
+from app.core.config import get_settings
+
+
+@lru_cache(maxsize=1)
+def get_redis_client() -> Redis | None:
+    redis_url = get_settings().redis_url.strip()
+    if not redis_url:
+        return None
+    return Redis.from_url(redis_url, decode_responses=True)
+
+
+async def ping_redis() -> bool:
+    client = get_redis_client()
+    if client is None:
+        return False
+    try:
+        return bool(await client.ping())
+    except Exception:
+        return False