From 72bd2e1e02231ec63f0dd8017179f1e2e7c3a872 Mon Sep 17 00:00:00 2001 From: Nelson <1475262689@qq.com> Date: Wed, 8 Jul 2026 17:20:29 +0800 Subject: [PATCH] Use async model streaming for chat --- .../apps/backend/app/api/chat.py | 11 +- .../app/services/chat_stream_service.py | 147 ++++++++++- .../app/services/model_stream_service.py | 232 +++++++++++++++++- ...026-07-08-concurrency-optimization-plan.md | 31 +++ 4 files changed, 415 insertions(+), 6 deletions(-) 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 85378dc..c3aff5b 100644 --- a/ai_knowledge_base_v2/apps/backend/app/api/chat.py +++ b/ai_knowledge_base_v2/apps/backend/app/api/chat.py @@ -1,7 +1,8 @@ from __future__ import annotations +import asyncio import json -from collections.abc import Iterator +from collections.abc import AsyncIterator from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse @@ -87,7 +88,7 @@ def stop( return api_success() -def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User) -> Iterator[str]: +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) acquired = queue_request.status == "acquired" @@ -111,7 +112,7 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User activeCount=queue_request.active_count, waitingCount=queue_request.waiting_count, ) - queue_request = chat_queue_manager.wait_for_slot(queued_token, config) + queue_request = await asyncio.to_thread(chat_queue_manager.wait_for_slot, queued_token, config) acquired = queue_request.status == "acquired" if queue_request.status == "timeout": yield _sse_event("error", message="排队等待超时，请稍后再试。") @@ -125,10 +126,12 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User activeCount=queue_request.active_count, waitingCount=queue_request.waiting_count, ) - for chunk in ChatStreamService.stream_answer(db, current_user, payload.sessionId, payload.message): + async for chunk in ChatStreamService.stream_answer_async(db, current_user, payload.sessionId, payload.message): yield _sse_event("content", content=chunk) except HTTPException as exc: yield _sse_event("error", message=str(exc.detail)) + except asyncio.CancelledError: + raise except Exception: yield _sse_event("error", message="AI 回复生成失败，请稍后再试。") finally: diff --git a/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py b/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py index febeff2..a9d7b98 100644 --- a/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py +++ b/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py @@ -1,6 +1,7 @@ from __future__ import annotations -from collections.abc import Iterator +import asyncio +from collections.abc import AsyncIterator, Iterator from datetime import UTC, datetime from inspect import signature from time import perf_counter @@ -138,6 +139,99 @@ class ChatStreamService: ) db.commit() + @staticmethod + async def stream_answer_async(db: Session, user: User, session_id: int, question: str) -> AsyncIterator[str]: + session = ChatService._get_user_session(db, user, session_id) + ChatService._ensure_quota(user) + + now = _now() + normalized_question = question.strip() + user_message = ChatMessage( + session_id=session.id, + user_id=user.id, + role="user", + content=normalized_question, + message_status="FINISHED", + created_at=now, + ) + db.add(user_message) + db.flush() + + history = list( + db.scalars( + select(ChatMessage) + .where( + ChatMessage.session_id == session.id, + ChatMessage.user_id == user.id, + ChatMessage.id < user_message.id, + ) + .order_by(ChatMessage.created_at.asc(), ChatMessage.id.asc()) + ) + ) + _maybe_update_summary(db, session, history) + + started_at = perf_counter() + rag_result = None + model_response = None + answer_parts: list[str] = [] + + try: + rag_result = _build_rag_result(db, user, normalized_question, history, getattr(session, "summary", None)) + model_response = ModelStreamService.stream_async(db, rag_result) + async for chunk in model_response.chunks: + if chunk: + answer_parts.append(chunk) + yield chunk + except asyncio.CancelledError: + db.rollback() + raise + except ExternalServiceError as exc: + cost_ms = int((perf_counter() - started_at) * 1000) + AiRequestLogService.write_failed( + db, + session_id=session.id, + message_id=user_message.id, + user_id=user.id, + model_name=model_response.model_name if model_response is not None else None, + prompt=rag_result.prompt if rag_result is not None else normalized_question, + knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None, + retrieve_count=len(rag_result.chunks) if rag_result is not None else 0, + cost_ms=cost_ms, + error_message=str(exc), + ) + db.commit() + raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc)) from exc + + answer = "".join(answer_parts) + if not answer.strip(): + cost_ms = int((perf_counter() - started_at) * 1000) + AiRequestLogService.write_failed( + db, + session_id=session.id, + message_id=user_message.id, + user_id=user.id, + model_name=model_response.model_name if model_response is not None else None, + prompt=rag_result.prompt if rag_result is not None else normalized_question, + knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None, + retrieve_count=len(rag_result.chunks) if rag_result is not None else 0, + cost_ms=cost_ms, + error_message="模型未返回有效内容", + ) + db.commit() + raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="模型未返回有效内容") + + _write_success( + db, + user=user, + session=session, + question=normalized_question, + answer=answer, + rag_result=rag_result, + model_response=model_response, + started_at=started_at, + now=now, + ) + def _now() -> datetime: return datetime.now(UTC).replace(tzinfo=None) @@ -147,6 +241,57 @@ def _rough_token_count(text: str) -> int: return max(1, len(text.strip()) // 2) +def _write_success( + db: Session, + *, + user: User, + session, + question: str, + answer: str, + rag_result, + model_response, + started_at: float, + now: datetime, +) -> None: + cost_ms = int((perf_counter() - started_at) * 1000) + assistant_message = ChatMessage( + session_id=session.id, + user_id=user.id, + role="assistant", + content=answer, + message_status="FINISHED", + token_input=model_response.input_token if model_response is not None else None, + token_output=_rough_token_count(answer), + response_time_ms=cost_ms, + model_id=model_response.model_id if model_response is not None else None, + created_at=now, + ) + db.add(assistant_message) + db.flush() + + session.message_count += 2 + session.last_message_at = now + if session.title == "新聊天": + session.title = _title_from_question(question) + + user.daily_chat_used += 1 + db.add_all([session, user]) + AiRequestLogService.write_success( + db, + session_id=session.id, + message_id=assistant_message.id, + user_id=user.id, + model_name=model_response.model_name if model_response is not None else "unknown", + prompt=rag_result.prompt if rag_result is not None else question, + knowledge_ids=rag_result.knowledge_ids if rag_result is not None else "", + retrieve_count=len(rag_result.chunks) if rag_result is not None else 0, + input_token=model_response.input_token if model_response is not None else 0, + output_token=_rough_token_count(answer), + cost_ms=cost_ms, + ) + db.commit() + + def _build_rag_result(db: Session, user: User, question: str, history: list[ChatMessage], summary: str | None): parameters = signature(RagService.build_result).parameters if "history" in parameters: diff --git a/ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py b/ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py index 683cdb6..36456dd 100644 --- a/ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py +++ b/ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py @@ -1,7 +1,8 @@ from __future__ import annotations +import asyncio import json -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass from typing import Any @@ -36,6 +37,14 @@ class StreamingModelResponse: chunks: Iterator[str] +@dataclass(frozen=True) +class AsyncStreamingModelResponse: + model_id: int | None + model_name: str + input_token: int + chunks: AsyncIterator[str] + + class ModelStreamService: @staticmethod def stream(db: Session, rag_result: RagResult) -> StreamingModelResponse: @@ -71,6 +80,40 @@ class ModelStreamService: chunks=_stream_configured_model(model, rag_result), ) + @staticmethod + def stream_async(db: Session, rag_result: RagResult) -> AsyncStreamingModelResponse: + model = _get_enabled_model(db) + mock_model_enabled = _system_config_bool(db, "mock_model_enabled", get_settings().mock_model_enabled) + + if mock_model_enabled: + model_name = model.model_name if model is not None else "mock-model" + return AsyncStreamingModelResponse( + model_id=model.id if model is not None else None, + model_name=model_name, + input_token=_rough_token_count(rag_result.prompt), + chunks=_async_chunk_text(_mock_answer(rag_result)), + ) + + if not rag_result.is_hit: + return AsyncStreamingModelResponse( + model_id=model.id if model is not None else None, + model_name=model.model_name if model is not None else "no-hit", + input_token=_rough_token_count(rag_result.prompt), + chunks=_async_chunk_text(NO_HIT_ANSWER), + ) + + if model is None: + raise ExternalServiceError("未启用可用模型，请先在模型管理中启用一个模型。", provider="model") + if not (model.api_url or model.base_url) or not model.api_key: + raise ExternalServiceError("模型 Base URL/API URL 或 API Key 未配置", provider="model") + + return AsyncStreamingModelResponse( + model_id=model.id, + model_name=model.model_name, + input_token=_rough_token_count(rag_result.prompt), + chunks=_stream_configured_model_async(model, rag_result), + ) + def _get_enabled_model(db: Session) -> ModelConfig | None: return db.scalar( @@ -92,6 +135,29 @@ def _stream_configured_model(model: ModelConfig, rag_result: RagResult) -> Itera return _chunk_text(_call_configured_model(model, rag_result)) +async def _stream_configured_model_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]: + api_type = model.api_type or "openai_compatible" + if model.stream_enabled != 1: + answer = await asyncio.to_thread(_call_configured_model, model, rag_result) + async for chunk in _async_chunk_text(answer): + yield chunk + return + + if api_type == "anthropic_messages": + async for chunk in _stream_anthropic_messages_async(model, rag_result): + yield chunk + return + + if api_type == "openai_compatible": + async for chunk in _stream_openai_compatible_model_async(model, rag_result): + yield chunk + return + + answer = await asyncio.to_thread(_call_configured_model, model, rag_result) + async for chunk in _async_chunk_text(answer): + yield chunk + + def _stream_openai_compatible_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]: payload: dict[str, Any] = { "model": model.model_name, @@ -124,6 +190,43 @@ def _stream_openai_compatible_model(model: ModelConfig, rag_result: RagResult) - raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc +async def _stream_openai_compatible_model_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]: + payload = _openai_stream_payload(model, rag_result) + try: + async with httpx.AsyncClient(timeout=model.timeout_second) as client: + async with client.stream( + "POST", + _resolve_openai_endpoint(model), + json=payload, + headers=_auth_headers(model), + ) as response: + response.raise_for_status() + async for chunk in _iter_openai_stream_async(response): + yield chunk + except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc: + raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc + + +def _openai_stream_payload(model: ModelConfig, rag_result: RagResult) -> dict[str, Any]: + payload: dict[str, Any] = { + "model": model.model_name, + "messages": [ + {"role": "system", "content": "你是企业知识库问答助手，只能基于已提供的知识片段回答。"}, + {"role": "user", "content": rag_result.prompt}, + ], + "max_tokens": model.max_token or 1024, + } + _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature)) + _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p)) + _put_if_not_none(payload, "presence_penalty", _decimal_to_float(model.presence_penalty)) + _put_if_not_none(payload, "frequency_penalty", _decimal_to_float(model.frequency_penalty)) + if model.response_format == "json_object": + payload["response_format"] = {"type": "json_object"} + payload.update(_load_extra_params(model.extra_params)) + payload["stream"] = True + return payload + + def _iter_openai_stream(response: httpx.Response) -> Iterator[str]: in_reasoning = False for data_line in _iter_sse_data(response): @@ -167,6 +270,49 @@ def _iter_openai_stream(response: httpx.Response) -> Iterator[str]: yield "" +async def _iter_openai_stream_async(response: httpx.Response) -> AsyncIterator[str]: + in_reasoning = False + async for data_line in _iter_sse_data_async(response): + data = json.loads(data_line) + if isinstance(data, dict) and data.get("error"): + raise ValueError(str(data["error"])) + + choices = data.get("choices") + if not isinstance(choices, list) or not choices: + continue + first_choice = choices[0] + if not isinstance(first_choice, dict): + continue + + delta = first_choice.get("delta") + if isinstance(delta, dict): + reasoning = _string_or_empty(delta.get("reasoning_content") or delta.get("reasoning")) + if reasoning: + if not in_reasoning: + in_reasoning = True + yield "" + yield reasoning + continue + + content = _string_or_empty(delta.get("content")) + if content: + if in_reasoning: + in_reasoning = False + yield "" + yield content + continue + + text = _string_or_empty(first_choice.get("text")) + if text: + if in_reasoning: + in_reasoning = False + yield "" + yield text + + if in_reasoning: + yield "" + + def _stream_anthropic_messages(model: ModelConfig, rag_result: RagResult) -> Iterator[str]: payload: dict[str, Any] = { "model": model.model_name, @@ -194,6 +340,38 @@ def _stream_anthropic_messages(model: ModelConfig, rag_result: RagResult) -> Ite raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc +async def _stream_anthropic_messages_async(model: ModelConfig, rag_result: RagResult) -> AsyncIterator[str]: + payload = _anthropic_stream_payload(model, rag_result) + try: + async with httpx.AsyncClient(timeout=model.timeout_second) as client: + async with client.stream( + "POST", + _resolve_anthropic_endpoint(model), + json=payload, + headers=_anthropic_headers(model), + ) as response: + response.raise_for_status() + async for chunk in _iter_anthropic_stream_async(response): + yield chunk + except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc: + raise ExternalServiceError(f"模型流式调用失败：{exc}", provider="model") from exc + + +def _anthropic_stream_payload(model: ModelConfig, rag_result: RagResult) -> dict[str, Any]: + payload: dict[str, Any] = { + "model": model.model_name, + "max_tokens": model.max_token or 1024, + "system": "你是企业知识库问答助手，只能基于已提供的知识片段回答。", + "messages": [{"role": "user", "content": rag_result.prompt}], + } + _put_if_not_none(payload, "temperature", _decimal_to_float(model.temperature)) + _put_if_not_none(payload, "top_p", _decimal_to_float(model.top_p)) + _put_if_not_none(payload, "top_k", model.top_k) + payload.update(_load_extra_params(model.extra_params)) + payload["stream"] = True + return payload + + def _iter_anthropic_stream(response: httpx.Response) -> Iterator[str]: in_reasoning = False for data_line in _iter_sse_data(response): @@ -227,6 +405,39 @@ def _iter_anthropic_stream(response: httpx.Response) -> Iterator[str]: yield "" +async def _iter_anthropic_stream_async(response: httpx.Response) -> AsyncIterator[str]: + in_reasoning = False + async for data_line in _iter_sse_data_async(response): + data = json.loads(data_line) + if isinstance(data, dict) and data.get("type") == "error": + raise ValueError(str(data.get("error", data))) + + if not isinstance(data, dict) or data.get("type") != "content_block_delta": + continue + delta = data.get("delta") + if not isinstance(delta, dict): + continue + + if delta.get("type") == "thinking_delta": + thinking = _string_or_empty(delta.get("thinking")) + if thinking: + if not in_reasoning: + in_reasoning = True + yield "" + yield thinking + continue + + text = _string_or_empty(delta.get("text")) + if text: + if in_reasoning: + in_reasoning = False + yield "" + yield text + + if in_reasoning: + yield "" + + def _iter_sse_data(response: httpx.Response) -> Iterator[str]: for line in response.iter_lines(): if not line: @@ -243,10 +454,29 @@ def _iter_sse_data(response: httpx.Response) -> Iterator[str]: yield data +async def _iter_sse_data_async(response: httpx.Response) -> AsyncIterator[str]: + async for line in response.aiter_lines(): + if not line: + continue + line = line.strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + break + if data: + yield data + + def _chunk_text(text: str, *, chunk_size: int = 12) -> Iterator[str]: for index in range(0, len(text), chunk_size): yield text[index : index + chunk_size] +async def _async_chunk_text(text: str, *, chunk_size: int = 12) -> AsyncIterator[str]: + for index in range(0, len(text), chunk_size): + yield text[index : index + chunk_size] + + def _string_or_empty(value: Any) -> str: return value if isinstance(value, str) else "" 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 b263e79..dc4c5bc 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 @@ -247,3 +247,34 @@ 2. 不修改数据库 schema。 3. 不把当前工作区已有的 5 个后端未提交文件混入阶段 3A 提交。 4. 如果 async 模型流式验证失败，保留同步流式入口，便于快速回退。 + +## 阶段 3A 实施记录 + +2026-07-08 完成模型长连接异步化第一版。 + +已完成： + +1. `/chat/completions` 的 SSE 生成器已切换为 async 生成器。 +2. 排队等待改为 `asyncio.to_thread(...)`，避免排队中的请求阻塞事件循环。 +3. `ChatStreamService` 新增 `stream_answer_async`，保留原同步 `stream_answer` 作为兼容入口。 +4. `ModelStreamService` 新增 async 模型流式入口： + - OpenAI 兼容协议使用 `httpx.AsyncClient.stream`。 + - Anthropic Messages 协议使用 `httpx.AsyncClient.stream`。 +5. mock、未命中知识库、关闭流式、MiniMax/Gemini 等路径保留安全降级；需要同步调用模型时使用 `asyncio.to_thread(...)` 放到线程中执行。 +6. SSE 事件协议未变化，用户端无需改动。 + +验证记录： + +1. 本机 `compileall` 检查通过。 +2. 容器内 `compileall` 检查通过。 +3. OpenAI 兼容协议 async 模拟流解析通过，输出格式为 `思考答案`。 +4. Anthropic Messages 协议 async 模拟流解析通过，输出格式为 `想答`。 +5. 后端健康检查通过。 +6. 真实 `/chat/completions` 验证通过：首个事件为 `generating`，随后连续返回 `content`，最后返回 `[DONE]`。 +7. 真实会话历史验证通过：流式完成后用户消息和助手消息均已落库。 + +阶段限制： + +1. 飞书检索、RAG 构建、数据库读写仍是同步逻辑；本阶段只优化模型长连接流式调用。 +2. 首字时间仍会受到飞书读取和 RAG 构建耗时影响。 +3. 队列仍是进程内队列，多 worker 全局并发控制需要阶段 4 Redis 队列。