diff --git a/ai_knowledge_base_v2/apps/admin-web/src/App.vue b/ai_knowledge_base_v2/apps/admin-web/src/App.vue index 51b172f..95e17ae 100644 --- a/ai_knowledge_base_v2/apps/admin-web/src/App.vue +++ b/ai_knowledge_base_v2/apps/admin-web/src/App.vue @@ -302,6 +302,15 @@ const systemSettingSections: SystemSettingSection[] = [ max: 3600, description: "请求在排队中超过该时间后自动结束并提示用户稍后再试。", }, + { + key: "chat_active_lease_seconds", + label: "问答执行租约（秒）", + type: "number", + defaultValue: 900, + min: 60, + max: 86400, + description: "Redis 全局队列中执行名额的自动过期时间，防止 worker 异常退出后长期占用并发名额。", + }, ], }, { diff --git a/ai_knowledge_base_v2/apps/backend/.env.example b/ai_knowledge_base_v2/apps/backend/.env.example index 916c55a..a0fa612 100644 --- a/ai_knowledge_base_v2/apps/backend/.env.example +++ b/ai_knowledge_base_v2/apps/backend/.env.example @@ -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 diff --git a/ai_knowledge_base_v2/apps/backend/app/api/chat.py b/ai_knowledge_base_v2/apps/backend/app/api/chat.py index c3aff5b..8da25f5 100644 --- a/ai_knowledge_base_v2/apps/backend/app/api/chat.py +++ b/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() diff --git a/ai_knowledge_base_v2/apps/backend/app/core/config.py b/ai_knowledge_base_v2/apps/backend/app/core/config.py index 305c769..1d47930 100644 --- a/ai_knowledge_base_v2/apps/backend/app/core/config.py +++ b/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 = "系统管理员" diff --git a/ai_knowledge_base_v2/apps/backend/app/services/chat_queue_redis_service.py b/ai_knowledge_base_v2/apps/backend/app/services/chat_queue_redis_service.py new file mode 100644 index 0000000..f2621a5 --- /dev/null +++ b/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 diff --git a/ai_knowledge_base_v2/apps/backend/app/services/chat_queue_runtime.py b/ai_knowledge_base_v2/apps/backend/app/services/chat_queue_runtime.py new file mode 100644 index 0000000..860c4d0 --- /dev/null +++ b/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) diff --git 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 index e1f94f1..b40dc02 100644 --- 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 @@ -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, + ), ) diff --git a/ai_knowledge_base_v2/apps/backend/app/services/redis_client.py b/ai_knowledge_base_v2/apps/backend/app/services/redis_client.py new file mode 100644 index 0000000..b3f1d15 --- /dev/null +++ b/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 diff --git a/ai_knowledge_base_v2/apps/backend/requirements.txt b/ai_knowledge_base_v2/apps/backend/requirements.txt index c5ae3dd..d744d5a 100644 --- a/ai_knowledge_base_v2/apps/backend/requirements.txt +++ b/ai_knowledge_base_v2/apps/backend/requirements.txt @@ -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 diff --git a/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md b/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md index 68014d0..d50a5da 100644 --- a/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md +++ b/ai_knowledge_base_v2/development_records/2026-07-08-concurrency-optimization-plan.md @@ -381,3 +381,38 @@ 1. 不改变知识库召回排序和 Prompt 构建。 2. 不缓存用户问题、模型回答和用户隐私信息。 3. 不把飞书读取失败结果长时间缓存，避免权限/配置修复后仍命中错误缓存。 + +## 阶段 4A 实施记录 + +2026-07-08 完成 Redis 全局问答队列第一版。 + +已完成： + +1. 后端增加 Redis 依赖和 `REDIS_URL` 配置。 +2. dev compose 增加 `redis` 服务，后端默认连接 `redis://redis:6379/0`。 +3. 新增 `RedisChatQueueManager`： + - 使用 Redis sorted set 保存 active 请求 token 和过期时间。 + - 使用 Redis list 保存等待队列 token。 + - 使用 Lua 脚本保证申请名额、入队、队首晋升等操作的原子性。 +4. 新增 `chat_queue_runtime`： + - Redis 可用时优先使用 Redis 全局队列。 + - Redis 不可用或运行时异常时，自动回退到原进程内队列。 +5. active token 增加租约过期时间，默认 `chat_active_lease_seconds=900`，避免 worker 异常退出后永久占用并发名额。 +6. 管理后台系统设置新增“问答执行租约（秒）”配置。 +7. 用户端 SSE 协议未变化。 + +验证记录： + +1. 本机 `compileall` 检查通过。 +2. `git diff --check` 通过。 +3. Docker dev 环境已新增并启动 Redis，Redis 健康检查通过。 +4. 容器内 Redis 队列行为验证通过：第一个请求获得执行名额，第二个请求进入排队，第三个请求被拒绝，释放名额后第二个请求恢复执行。 +5. 后端健康检查通过。 +6. 真实 `/chat/completions` 验证通过：首个事件为 `generating`，随后返回 `content`，最后返回 `[DONE]`。 +7. 真实会话历史验证通过：流式完成后用户消息和助手消息均已落库。 + +阶段限制： + +1. 当前 Redis 队列使用轮询晋升等待请求，下一步可以结合 Pub/Sub 或 blocking pop 降低等待侧轮询。 +2. Redis 飞书共享缓存还未接入，将在阶段 4B 完成。 +3. 生产部署需要确认 Redis 持久化、内存上限、淘汰策略和监控告警。 diff --git a/ai_knowledge_base_v2/docker-compose.dev.yml b/ai_knowledge_base_v2/docker-compose.dev.yml index 8aa8adc..10313cb 100644 --- a/ai_knowledge_base_v2/docker-compose.dev.yml +++ b/ai_knowledge_base_v2/docker-compose.dev.yml @@ -21,6 +21,20 @@ services: timeout: 5s retries: 20 + redis: + image: redis:7-alpine + container_name: ai_kb_v2_redis + ports: + - "6380:6379" + command: ["redis-server", "--appendonly", "yes"] + volumes: + - ai_kb_v2_redis_data:/data + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 20 + backend: build: context: ./apps/backend @@ -29,6 +43,7 @@ services: APP_ENV: docker DEBUG: "true" DATABASE_URL: mysql+pymysql://ai_kb:ai_kb@mysql:3306/ai_knowledge_base_v2?charset=utf8mb4 + REDIS_URL: redis://redis:6379/0 JWT_SECRET_KEY: local-dev-secret-change-before-production MOCK_SMS_ENABLED: "true" MOCK_SMS_CODE: "123456" @@ -44,6 +59,8 @@ services: depends_on: mysql: condition: service_healthy + redis: + condition: service_healthy user-client: build: @@ -69,3 +86,4 @@ services: volumes: ai_kb_v2_mysql_data: + ai_kb_v2_redis_data: