Stream chat completions from model providers

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

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 "<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,
+        "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 "<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:
+            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 ""