From 84e6f1f51a0a4543d5bba65e3070624824ba3866 Mon Sep 17 00:00:00 2001 From: gqt <3217233537@qq.com> Date: Mon, 13 Jul 2026 11:43:23 +0800 Subject: [PATCH] feat: stream dual-agent answers over sse --- backend/src/kbqa/api/router.py | 3 +- backend/src/kbqa/chat/routes.py | 29 +++- backend/src/kbqa/chat/schemas.py | 5 + backend/src/kbqa/chat/service.py | 220 +++++++++++++++++++++++++++++++ backend/src/kbqa/chat/sse.py | 11 ++ backend/src/kbqa/lifecycle.py | 3 + 6 files changed, 269 insertions(+), 2 deletions(-) create mode 100644 backend/src/kbqa/chat/service.py create mode 100644 backend/src/kbqa/chat/sse.py diff --git a/backend/src/kbqa/api/router.py b/backend/src/kbqa/api/router.py index 035f77c..ca9ac80 100644 --- a/backend/src/kbqa/api/router.py +++ b/backend/src/kbqa/api/router.py @@ -1,8 +1,9 @@ from fastapi import APIRouter -from kbqa.chat.routes import router as sessions_router +from kbqa.chat.routes import chat_router, router as sessions_router from kbqa.documents.routes import router as documents_router api_router = APIRouter(prefix="/api/v1") api_router.include_router(documents_router) api_router.include_router(sessions_router) +api_router.include_router(chat_router) diff --git a/backend/src/kbqa/chat/routes.py b/backend/src/kbqa/chat/routes.py index 7fc528f..8885640 100644 --- a/backend/src/kbqa/chat/routes.py +++ b/backend/src/kbqa/chat/routes.py @@ -1,6 +1,7 @@ from typing import Annotated -from fastapi import APIRouter, Depends, Query, Response, status +from fastapi import APIRouter, Depends, Query, Request, Response, status +from fastapi.responses import StreamingResponse from sqlalchemy.ext.asyncio import AsyncSession from kbqa.api.errors import AppError @@ -9,6 +10,7 @@ from kbqa.chat.schemas import ( MessagePage, MessageResponse, MessageSourceResponse, + ChatStreamRequest, SessionCreate, SessionPage, SessionResponse, @@ -16,6 +18,7 @@ from kbqa.chat.schemas import ( from kbqa.database import get_session router = APIRouter(prefix="/sessions", tags=["sessions"]) +chat_router = APIRouter(prefix="/chat", tags=["chat"]) SessionDependency = Annotated[AsyncSession, Depends(get_session)] @@ -91,3 +94,27 @@ async def delete_session(session_id: str, session: SessionDependency) -> Respons if not await ChatRepository(session).delete_session(session_id): raise AppError("SESSION_NOT_FOUND", "会话不存在", 404) return Response(status_code=status.HTTP_204_NO_CONTENT) + + +@chat_router.post("/stream") +async def stream_chat( + payload: ChatStreamRequest, + request: Request, + session: SessionDependency, +) -> StreamingResponse: + if await ChatRepository(session).get_session(payload.session_id) is None: + raise AppError("SESSION_NOT_FOUND", "会话不存在", 404) + query = payload.query.strip() + if not query: + raise AppError("VALIDATION_ERROR", "问题不能为空", 422) + service = request.app.state.chat_service + service.reserve(payload.session_id) + return StreamingResponse( + service.stream(payload.session_id, query, request.state.request_id), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + "X-Request-ID": request.state.request_id, + }, + ) diff --git a/backend/src/kbqa/chat/schemas.py b/backend/src/kbqa/chat/schemas.py index 7fc7698..8f5c62f 100644 --- a/backend/src/kbqa/chat/schemas.py +++ b/backend/src/kbqa/chat/schemas.py @@ -9,6 +9,11 @@ class SessionCreate(BaseModel): title: str | None = Field(default=None, max_length=200) +class ChatStreamRequest(BaseModel): + session_id: str + query: str = Field(min_length=1, max_length=10_000) + + class SessionResponse(BaseModel): model_config = ConfigDict(from_attributes=True) diff --git a/backend/src/kbqa/chat/service.py b/backend/src/kbqa/chat/service.py new file mode 100644 index 0000000..ff4e7a0 --- /dev/null +++ b/backend/src/kbqa/chat/service.py @@ -0,0 +1,220 @@ +import asyncio +import logging +import re +from collections.abc import AsyncIterator, Sequence +from contextlib import suppress +from typing import Any + +from langchain.messages import AIMessage, AIMessageChunk, HumanMessage +from sqlalchemy.ext.asyncio import async_sessionmaker + +from kbqa.agents.context import ChatRunContext +from kbqa.agents.factory import AgentBundle, create_agent_bundle +from kbqa.api.errors import AppError +from kbqa.chat.repository import ChatRepository, MessageSourceInput +from kbqa.chat.sse import encode_event, encode_ping +from kbqa.config import Settings +from kbqa.rag.retriever import KnowledgeRetriever + +INSUFFICIENT_ANSWER = "知识库中没有足够资料回答这个问题。" +logger = logging.getLogger(__name__) + + +class ChatService: + def __init__( + self, + settings: Settings, + session_factory: async_sessionmaker, + retriever: KnowledgeRetriever, + ) -> None: + self.settings = settings + self.session_factory = session_factory + self.retriever = retriever + self._active_sessions: set[str] = set() + self._agents: AgentBundle | None = None + + def reserve(self, session_id: str) -> None: + if session_id in self._active_sessions: + raise AppError("SESSION_BUSY", "当前会话正在生成回答", 409) + self._active_sessions.add(session_id) + + def _get_agents(self) -> AgentBundle: + if self._agents is None: + self._agents = create_agent_bundle(self.settings) + return self._agents + + async def stream( + self, + session_id: str, + query: str, + request_id: str, + ) -> AsyncIterator[bytes]: + assistant_id: str | None = None + graph_stream = None + try: + async with self.session_factory() as session: + repository = ChatRepository(session) + if await repository.get_session(session_id) is None: + yield encode_event( + "error", + { + "code": "SESSION_NOT_FOUND", + "message": "会话不存在", + "retryable": False, + }, + ) + return + _, assistant = await repository.begin_turn(session_id, query) + assistant_id = assistant.id + history = await repository.load_recent_messages( + session_id, self.settings.history_message_limit + ) + + messages = [ + HumanMessage(content=message.content) + if message.role == "user" + else AIMessage(content=message.content) + for message in history + if message.id != assistant_id + ] + sources: list[dict[str, Any]] = [] + answer_parts: list[str] = [] + answering = False + no_evidence = False + yield encode_event("status", {"phase": "researching"}) + + graph_stream = self._get_agents().main_agent.astream( + {"messages": messages}, + context=ChatRunContext( + retriever=self.retriever, + request_id=request_id, + session_id=session_id, + assistant_message_id=assistant_id, + ), + config={ + "recursion_limit": 12, + "tags": ["main-agent"], + "metadata": { + "request_id": request_id, + "session_id": session_id, + "assistant_message_id": assistant_id, + }, + }, + stream_mode=["messages", "custom", "updates"], + version="v2", + ) + iterator = graph_stream.__aiter__() + pending = asyncio.create_task(anext(iterator)) + while True: + done, _ = await asyncio.wait({pending}, timeout=15.0) + if not done: + yield encode_ping() + continue + try: + chunk = pending.result() + except StopAsyncIteration: + break + pending = asyncio.create_task(anext(iterator)) + chunk_type = chunk.get("type") + data = chunk.get("data") + if chunk_type == "custom" and isinstance(data, dict): + if data.get("kind") == "no_evidence": + no_evidence = True + pending.cancel() + with suppress(asyncio.CancelledError): + await pending + await graph_stream.aclose() + graph_stream = None + break + if data.get("kind") == "sources": + sources = list(data.get("sources", [])) + yield encode_event("sources", {"sources": sources}) + yield encode_event("status", {"phase": "answering"}) + answering = True + continue + if chunk_type != "messages" or not answering: + continue + token, _metadata = data + if not isinstance(token, AIMessageChunk): + continue + text = self._message_text(token.content) + if text: + answer_parts.append(text) + yield encode_event("token", {"content": text}) + + if no_evidence: + answer = INSUFFICIENT_ANSWER + yield encode_event("status", {"phase": "answering"}) + yield encode_event("token", {"content": answer}) + else: + answer = "".join(answer_parts).strip() + self._validate_citations(answer, sources) + + source_inputs = [ + MessageSourceInput( + document_id=str(source["document_id"]), + chunk_id=str(source["chunk_id"]), + citation_label=str(source["label"]), + score=float(source["score"]), + ) + for source in sources + ] + async with self.session_factory() as session: + await ChatRepository(session).complete_answer(assistant_id, answer, source_inputs) + yield encode_event("done", {"message_id": assistant_id}) + except asyncio.CancelledError: + if assistant_id is not None: + await self._mark_failed(assistant_id, "CLIENT_DISCONNECTED") + raise + except AppError as exc: + if assistant_id is not None: + await self._mark_failed(assistant_id, exc.code) + yield encode_event( + "error", + {"code": exc.code, "message": exc.message, "retryable": exc.retryable}, + ) + except Exception: + logger.exception("Chat stream failed request_id=%s", request_id) + if assistant_id is not None: + await self._mark_failed(assistant_id, "MODEL_UPSTREAM_ERROR") + yield encode_event( + "error", + { + "code": "MODEL_UPSTREAM_ERROR", + "message": "模型调用失败，请稍后重试", + "retryable": True, + }, + ) + finally: + if graph_stream is not None: + await graph_stream.aclose() + self._active_sessions.discard(session_id) + + async def _mark_failed(self, message_id: str, error_code: str) -> None: + async with self.session_factory() as session: + await ChatRepository(session).fail_answer(message_id, error_code) + + @staticmethod + def _message_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, Sequence): + parts: list[str] = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + return "".join(parts) + return "" + + @staticmethod + def _validate_citations(answer: str, sources: list[dict[str, Any]]) -> None: + if not answer: + raise AppError("EMPTY_MODEL_ANSWER", "模型没有返回有效回答", 502, retryable=True) + allowed = {str(source["label"]) for source in sources} + used = set(re.findall(r"\[(S\d+)]", answer)) + if sources and not used: + raise AppError("MISSING_CITATION", "模型回答缺少知识库引用", 502, retryable=True) + if not used.issubset(allowed): + raise AppError("INVALID_CITATION", "模型回答包含无效引用", 502, retryable=True) diff --git a/backend/src/kbqa/chat/sse.py b/backend/src/kbqa/chat/sse.py new file mode 100644 index 0000000..7709f6b --- /dev/null +++ b/backend/src/kbqa/chat/sse.py @@ -0,0 +1,11 @@ +import json +from collections.abc import Mapping + + +def encode_event(event: str, data: Mapping[str, object]) -> bytes: + payload = json.dumps(data, ensure_ascii=False, separators=(",", ":")) + return f"event: {event}\ndata: {payload}\n\n".encode() + + +def encode_ping() -> bytes: + return b": ping\n\n" diff --git a/backend/src/kbqa/lifecycle.py b/backend/src/kbqa/lifecycle.py index 2247f2c..975b732 100644 --- a/backend/src/kbqa/lifecycle.py +++ b/backend/src/kbqa/lifecycle.py @@ -7,6 +7,7 @@ from fastapi import FastAPI from pymilvus import MilvusClient from kbqa.config import get_settings +from kbqa.chat.service import ChatService from kbqa.database import Database from kbqa.documents.service import DocumentService from kbqa.indexing.worker import IndexWorker @@ -53,6 +54,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: vector_store, vector_mutation_lock, ) + chat_service = ChatService(settings, database.session_factory, retriever) app.state.settings = settings app.state.database = database app.state.milvus_client = milvus_client @@ -61,6 +63,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.retriever = retriever app.state.index_worker = worker app.state.document_service = document_service + app.state.chat_service = chat_service await worker.start() await document_service.recover_deletions() try: