From 465a481e7aaf9e0feab1f426fb4886d851ca4b19 Mon Sep 17 00:00:00 2001 From: Nelson <1475262689@qq.com> Date: Wed, 8 Jul 2026 17:13:53 +0800 Subject: [PATCH] Stream chat completions from model providers --- .../apps/backend/app/api/chat.py | 13 +- .../app/services/chat_stream_service.py | 170 ++++++++++++ .../app/services/model_stream_service.py | 252 ++++++++++++++++++ ...026-07-08-concurrency-optimization-plan.md | 26 ++ 4 files changed, 451 insertions(+), 10 deletions(-) create mode 100644 ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py create mode 100644 ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py 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 d497f33..85378dc 100644 --- a/ai_knowledge_base_v2/apps/backend/app/api/chat.py +++ b/ai_knowledge_base_v2/apps/backend/app/api/chat.py @@ -21,6 +21,7 @@ from app.schemas.chat import ( ) 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_stream_service import ChatStreamService router = APIRouter() @@ -124,8 +125,8 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User activeCount=queue_request.active_count, waitingCount=queue_request.waiting_count, ) - answer = ChatService.create_answer(db, current_user, payload.sessionId, payload.message) - yield from _sse_chunks(answer) + for chunk in ChatStreamService.stream_answer(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 Exception: @@ -138,14 +139,6 @@ def _chat_stream(payload: ChatCompletionRequest, db: Session, current_user: User yield _sse_done() - -def _sse_chunks(answer: str) -> Iterator[str]: - chunk_size = 12 - for index in range(0, len(answer), chunk_size): - chunk = answer[index : index + chunk_size] - yield _sse_event("content", content=chunk) - - def _sse_event(event_type: str, **payload: object) -> str: return f"data: {json.dumps({'type': event_type, **payload}, ensure_ascii=False)}\n\n" 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 new file mode 100644 index 0000000..febeff2 --- /dev/null +++ b/ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from collections.abc import Iterator +from datetime import UTC, datetime +from inspect import signature +from time import perf_counter + +from fastapi import HTTPException, status +from sqlalchemy import select +from sqlalchemy.orm import Session + +from app.models.chat import ChatMessage +from app.models.user import User +from app.services.ai_request_log_service import AiRequestLogService +from app.services.chat_service import ChatService, _title_from_question +from app.services.external_errors import ExternalServiceError +from app.services.model_stream_service import ModelStreamService +from app.services.rag_service import RagService + + +class ChatStreamService: + @staticmethod + def stream_answer(db: Session, user: User, session_id: int, question: str) -> Iterator[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(db, rag_result) + for chunk in model_response.chunks: + if chunk: + answer_parts.append(chunk) + yield chunk + except GeneratorExit: + 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="模型未返回有效内容") + + 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(normalized_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 normalized_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 _now() -> datetime: + return datetime.now(UTC).replace(tzinfo=None) + + +def _rough_token_count(text: str) -> int: + return max(1, len(text.strip()) // 2) + + +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: + return RagService.build_result(db, user, question, history=history, session_summary=summary) + return RagService.build_result(db, user, question) + + +def _maybe_update_summary(db: Session, session, history: list[ChatMessage]) -> None: + summary_handler = getattr(ChatService, "_maybe_update_summary", None) + if not callable(summary_handler) or not hasattr(session, "summary_up_to_message_id"): + return + try: + from app.services import rag_service + + trigger_rounds = getattr(rag_service, "SUMMARY_TRIGGER_ROUNDS", None) + if trigger_rounds is None: + return + if len(history) // 2 >= trigger_rounds: + summary_handler(db, session, history) + except Exception: + return 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 new file mode 100644 index 0000000..683cdb6 --- /dev/null +++ b/ai_knowledge_base_v2/apps/backend/app/services/model_stream_service.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import json +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any + +import httpx +from sqlalchemy import select +from sqlalchemy.orm import Session + +from app.core.config import get_settings +from app.models.ai_config import ModelConfig +from app.services.external_errors import ExternalServiceError +from app.services.model_service import ( + _anthropic_headers, + _auth_headers, + _call_configured_model, + _decimal_to_float, + _load_extra_params, + _mock_answer, + _put_if_not_none, + _resolve_anthropic_endpoint, + _resolve_openai_endpoint, + _rough_token_count, + _system_config_bool, +) +from app.services.rag_service import NO_HIT_ANSWER, RagResult + + +@dataclass(frozen=True) +class StreamingModelResponse: + model_id: int | None + model_name: str + input_token: int + chunks: Iterator[str] + + +class ModelStreamService: + @staticmethod + def stream(db: Session, rag_result: RagResult) -> StreamingModelResponse: + 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 StreamingModelResponse( + model_id=model.id if model is not None else None, + model_name=model_name, + input_token=_rough_token_count(rag_result.prompt), + chunks=_chunk_text(_mock_answer(rag_result)), + ) + + if not rag_result.is_hit: + return StreamingModelResponse( + 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=_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 StreamingModelResponse( + model_id=model.id, + model_name=model.model_name, + input_token=_rough_token_count(rag_result.prompt), + chunks=_stream_configured_model(model, rag_result), + ) + + +def _get_enabled_model(db: Session) -> ModelConfig | None: + return db.scalar( + select(ModelConfig) + .where(ModelConfig.enabled == 1) + .order_by(ModelConfig.id.desc()) + .limit(1) + ) + + +def _stream_configured_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]: + api_type = model.api_type or "openai_compatible" + if model.stream_enabled != 1: + return _chunk_text(_call_configured_model(model, rag_result)) + if api_type == "anthropic_messages": + return _stream_anthropic_messages(model, rag_result) + if api_type == "openai_compatible": + return _stream_openai_compatible_model(model, rag_result) + return _chunk_text(_call_configured_model(model, rag_result)) + + +def _stream_openai_compatible_model(model: ModelConfig, rag_result: RagResult) -> Iterator[str]: + 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 + + try: + with httpx.stream( + "POST", + _resolve_openai_endpoint(model), + json=payload, + headers=_auth_headers(model), + timeout=model.timeout_second, + ) as response: + response.raise_for_status() + yield from _iter_openai_stream(response) + except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc: + raise ExternalServiceError(f"模型流式调用失败:{exc}", provider="model") from exc + + +def _iter_openai_stream(response: httpx.Response) -> Iterator[str]: + in_reasoning = False + for data_line in _iter_sse_data(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, + "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 + + try: + with httpx.stream( + "POST", + _resolve_anthropic_endpoint(model), + json=payload, + headers=_anthropic_headers(model), + timeout=model.timeout_second, + ) as response: + response.raise_for_status() + yield from _iter_anthropic_stream(response) + except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc: + raise ExternalServiceError(f"模型流式调用失败:{exc}", provider="model") from exc + + +def _iter_anthropic_stream(response: httpx.Response) -> Iterator[str]: + in_reasoning = False + for data_line in _iter_sse_data(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: + continue + if isinstance(line, bytes): + line = line.decode("utf-8") + 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] + + +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 377d7c4..df2366b 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 @@ -171,3 +171,29 @@ 3. 非流式协议保留降级路径,保证模型配置多样性不被破坏。 4. 本阶段仍使用同步 `httpx` 的流式能力;真正 async 化放到阶段 3。 5. 本阶段不引入 Redis,全局多 worker 协调继续放到阶段 4。 + +已完成: + +1. 新增 `ModelStreamService`,负责模型流式协议适配。 +2. 新增 `ChatStreamService`,负责流式问答过程中的用户消息保存、模型 chunk 输出、助手消息落库和 AI 请求日志记录。 +3. `/chat/completions` 已从“完整回答后切片输出”改为“边接收模型 chunk 边发送 SSE `content` 事件”。 +4. OpenAI 兼容协议支持 `stream=true` 并解析 `choices[].delta.content`。 +5. Anthropic Messages 协议支持 `stream=true` 并解析 `content_block_delta` 文本事件。 +6. DeepSeek 等兼容协议返回的 `reasoning_content` 会被转换成 `...`,保持当前 H5 思考态显示逻辑可用。 +7. MiniMax、Gemini 或关闭流式的模型配置会自动降级为非流式调用,再通过统一 SSE 事件输出。 +8. 如果当前后端已具备历史对话/摘要能力,流式服务会自动传入历史上下文;旧版本签名不支持时自动回退。 + +验证记录: + +1. 容器内 `compileall` 检查通过。 +2. OpenAI 兼容协议模拟流式事件解析通过,输出格式为 `思考答案`。 +3. Anthropic Messages 协议模拟流式事件解析通过,输出格式为 `答`。 +4. 后端健康检查通过。 +5. 真实 `/chat/completions` 验证通过:接口先返回 `generating`,随后连续返回上游模型 chunk,最后返回 `[DONE]`。 +6. 真实会话历史验证通过:流式完成后用户消息和助手消息均已落库。 + +阶段限制: + +1. 本阶段仍是同步 `httpx.stream`,一个请求生成期间仍会占用当前 worker 执行资源。 +2. 用户端能更快看到模型输出,但飞书检索阶段仍在首字之前完成;首字时间还会受飞书读取耗时影响。 +3. 多 worker 全局排队、共享飞书缓存、异步飞书读取仍需在阶段 3 和阶段 4 继续处理。