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

Add Redis-backed chat queue

Browse Source
This commit is contained in:
Nelson
2026-07-08 17:36:20 +08:00
parent db1d83631b
commit 23d5f42972
11 changed files with 349 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,7 @@ DEBUG=true
API_PREFIX=/api
DATABASE_URL=mysql+pymysql://ai_kb:ai_kb@127.0.0.1:3306/ai_knowledge_base_v2?charset=utf8mb4
REDIS_URL=
AUTO_CREATE_TABLES=false
JWT_SECRET_KEY=change-me-in-production
@@ -27,6 +28,7 @@ DEFAULT_USER_NAME_PREFIX=用户
CHAT_MAX_ACTIVE_REQUESTS=2
CHAT_MAX_QUEUE_SIZE=20
CHAT_QUEUE_TIMEOUT_SECONDS=60
CHAT_ACTIVE_LEASE_SECONDS=900
BOOTSTRAP_ADMIN_USERNAME=admin
BOOTSTRAP_ADMIN_PASSWORD=admin123456

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

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

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

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

155
ai_knowledge_base_v2/apps/backend/app/services/chat_queue_redis_service.py Normal file
View File

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

65
ai_knowledge_base_v2/apps/backend/app/services/chat_queue_runtime.py Normal file
View File

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

8
ai_knowledge_base_v2/apps/backend/app/services/chat_queue_service.py
View File

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

25
ai_knowledge_base_v2/apps/backend/app/services/redis_client.py Normal file
View File

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

1
ai_knowledge_base_v2/apps/backend/requirements.txt
View File

@@ -9,5 +9,6 @@ PyJWT>=2.10.0
PyMySQL>=1.1.1
cryptography>=45.0.0
httpx>=0.28.0
redis>=5.0.0
openpyxl>=3.1.5
python-multipart>=0.0.20