Use async model streaming for chat

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:

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 "</think>"
 
 
+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 "<think>"
+                yield reasoning
+                continue
+
+            content = _string_or_empty(delta.get("content"))
+            if content:
+                if in_reasoning:
+                    in_reasoning = False
+                    yield "</think>"
+                yield content
+                continue
+
+        text = _string_or_empty(first_choice.get("text"))
+        if text:
+            if in_reasoning:
+                in_reasoning = False
+                yield "</think>"
+            yield text
+
+    if in_reasoning:
+        yield "</think>"
+
+
 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 "</think>"
 
 
+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 "<think>"
+                yield thinking
+            continue
+
+        text = _string_or_empty(delta.get("text"))
+        if text:
+            if in_reasoning:
+                in_reasoning = False
+                yield "</think>"
+            yield text
+
+    if in_reasoning:
+        yield "</think>"
+
+
 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 ""