"""Research endpoints — legacy SSE plus plan/run Deep Research workflow. 深度研究接口模块。 提供三阶段研究流程：生成研究计划 (plan) -> 执行研究 (run) -> 单章节重跑 (rerun)， 以及兼容旧版的一键流式研究接口。所有生成结果通过 SSE 实时推送到前端。 """ from __future__ import annotations import json from typing import Annotated, AsyncIterator from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse from app.api.deps import ( UserContext, get_current_user, get_embedding_client, get_es_client, get_llm_client, get_mysql_client, get_neo4j_client, get_redis_client, ) from app.api.schemas.research import ( ResearchChunk, ResearchPlanRequest, ResearchPlanResponse, ResearchRequest, ResearchRunRequest, ResearchSectionRerunRequest, ) from app.core.embedding import EmbeddingService from app.core.graph_query_service import GraphQueryService from app.core.permission import PermissionService from app.core.research_engine import ResearchEngine from app.infrastructure.embedding_client import EmbeddingClient from app.infrastructure.es_client import ESClient from app.infrastructure.llm_client import LLMClient from app.infrastructure.mysql_client import MySQLClient from app.infrastructure.neo4j_client import Neo4jClient from app.infrastructure.redis_client import RedisClient from app.infrastructure.session_store import build_research_session_store from app.utils.logger import get_logger logger = get_logger(__name__) router = APIRouter(prefix="/research", tags=["research"]) # --------------------------------------------------------------------------- # 辅助函数：构建 ResearchEngine 实例，避免各端点重复创建服务对象 # --------------------------------------------------------------------------- def _build_research_engine( es_client: ESClient, neo4j_client: Neo4jClient, redis_client: RedisClient, mysql_client: MySQLClient | None, embedding_client: EmbeddingClient, llm_client: LLMClient, ) -> ResearchEngine: """构建 ResearchEngine 及其依赖服务，各端点共用此辅助函数以减少重复代码。 Build a ResearchEngine with all required services wired up. """ embedding_svc = EmbeddingService(embedding_client) graph_svc = GraphQueryService(neo4j_client) session_store = build_research_session_store( redis_client=redis_client, mysql_client=mysql_client, ) return ResearchEngine( es_client=es_client, embedding_service=embedding_svc, graph_service=graph_svc, llm_client=llm_client, session_store=session_store, ) # --------------------------------------------------------------------------- # SSE helpers # --------------------------------------------------------------------------- async def _sse_stream( gen: AsyncIterator[ResearchChunk], ) -> AsyncIterator[str]: """Wrap an async generator of ResearchChunks into SSE-format strings. 将 ResearchChunk 序列化为 SSE data 帧格式，供 StreamingResponse 使用。 当生成器抛出异常时，向客户端发送 SSE error 事件，避免流被静默截断。 """ try: async for chunk in gen: payload = json.dumps(chunk.model_dump(exclude_none=True), ensure_ascii=False) yield f"data: {payload}\n\n" except Exception as exc: # 生成器异常时向客户端推送 error 事件，防止流被截断后客户端无感知 logger.error("sse_stream_error", error=str(exc), exc_info=exc) error_chunk = ResearchChunk(type="error", content=str(exc)) error_payload = json.dumps(error_chunk.model_dump(exclude_none=True), ensure_ascii=False) yield f"data: {error_payload}\n\n" # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- @router.post( "", summary="研究分析 (SSE 流式)", response_description="Server-Sent Events stream of ResearchChunk objects", ) async def research( body: ResearchRequest, user: Annotated[UserContext, Depends(get_current_user)], es_client: Annotated[ESClient, Depends(get_es_client)], neo4j_client: Annotated[Neo4jClient, Depends(get_neo4j_client)], redis_client: Annotated[RedisClient, Depends(get_redis_client)], mysql_client: Annotated[MySQLClient | None, Depends(get_mysql_client)], embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], llm_client: Annotated[LLMClient, Depends(get_llm_client)], ) -> StreamingResponse: """Stream a deep-research answer via Server-Sent Events. The response is ``text/event-stream`` emitting JSON-encoded :class:`~app.api.schemas.research.ResearchChunk` objects. Chunk types ----------- - ``thinking`` — intermediate status update (not part of final answer) - ``reference`` — a document that was retrieved and used as context - ``text`` — LLM-generated answer token(s) - ``done`` — end-of-stream sentinel - ``error`` — fatal error during processing """ logger.info( "research_request", user_id=user.user_id, question_len=len(body.question), session_id=body.session_id, seed_doc_count=len(body.seed_doc_ids or []), ) # Resolve permission tokens from user context (same as search endpoint) perm_service = PermissionService(redis_client=redis_client) perm = await perm_service.resolve(user) # 通过辅助函数构建 ResearchEngine，避免重复创建服务对象 engine = _build_research_engine( es_client, neo4j_client, redis_client, mysql_client, embedding_client, llm_client, ) gen = engine.research( body.question, perm, session_id=body.session_id, seed_doc_ids=body.seed_doc_ids, ) return StreamingResponse( _sse_stream(gen), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", # Disable Nginx buffering for SSE }, ) @router.post( "/plan", summary="生成研究计划", response_model=ResearchPlanResponse, ) async def research_plan( body: ResearchPlanRequest, user: Annotated[UserContext, Depends(get_current_user)], es_client: Annotated[ESClient, Depends(get_es_client)], neo4j_client: Annotated[Neo4jClient, Depends(get_neo4j_client)], redis_client: Annotated[RedisClient, Depends(get_redis_client)], mysql_client: Annotated[MySQLClient | None, Depends(get_mysql_client)], embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], llm_client: Annotated[LLMClient, Depends(get_llm_client)], ) -> ResearchPlanResponse: """根据研究任务定义生成研究计划，包含目标拆解、子问题、章节大纲等。""" logger.info( "research_plan_request", user_id=user.user_id, topic=body.task.topic, session_id=body.session_id, seed_doc_count=len(body.seed_doc_ids or []), ) # 通过辅助函数构建 ResearchEngine，避免重复创建服务对象 engine = _build_research_engine( es_client, neo4j_client, redis_client, mysql_client, embedding_client, llm_client, ) plan = await engine.build_plan(body.task, seed_doc_ids=body.seed_doc_ids) return ResearchPlanResponse(session_id=body.session_id, task=body.task, plan=plan) @router.post( "/run", summary="执行深度研究 (SSE 流式)", response_description="Server-Sent Events stream of ResearchChunk objects", ) async def research_run( body: ResearchRunRequest, user: Annotated[UserContext, Depends(get_current_user)], es_client: Annotated[ESClient, Depends(get_es_client)], neo4j_client: Annotated[Neo4jClient, Depends(get_neo4j_client)], redis_client: Annotated[RedisClient, Depends(get_redis_client)], mysql_client: Annotated[MySQLClient | None, Depends(get_mysql_client)], embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], llm_client: Annotated[LLMClient, Depends(get_llm_client)], ) -> StreamingResponse: """基于已确认的研究计划执行深度研究，以 SSE 流式返回研究报告。""" logger.info( "research_run_request", user_id=user.user_id, topic=body.task.topic, session_id=body.session_id, explicit_doc_count=len(body.plan.included_doc_ids or body.seed_doc_ids or []), ) perm_service = PermissionService(redis_client=redis_client) perm = await perm_service.resolve(user) # 通过辅助函数构建 ResearchEngine，避免重复创建服务对象 engine = _build_research_engine( es_client, neo4j_client, redis_client, mysql_client, embedding_client, llm_client, ) gen = engine.run_deep_research( body.task, body.plan, perm, session_id=body.session_id, ) return StreamingResponse( _sse_stream(gen), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) @router.post( "/sections/rerun", summary="重跑研究章节 (SSE 流式)", response_description="Server-Sent Events stream of ResearchChunk objects", ) async def research_rerun_section( body: ResearchSectionRerunRequest, user: Annotated[UserContext, Depends(get_current_user)], es_client: Annotated[ESClient, Depends(get_es_client)], neo4j_client: Annotated[Neo4jClient, Depends(get_neo4j_client)], redis_client: Annotated[RedisClient, Depends(get_redis_client)], mysql_client: Annotated[MySQLClient | None, Depends(get_mysql_client)], embedding_client: Annotated[EmbeddingClient, Depends(get_embedding_client)], llm_client: Annotated[LLMClient, Depends(get_llm_client)], ) -> StreamingResponse: """对研究报告中的单个章节进行重新生成，支持指定来源文档。""" logger.info( "research_section_rerun_request", user_id=user.user_id, topic=body.task.topic, section_title=body.section_title, session_id=body.session_id, source_doc_count=len(body.source_doc_ids or []), explicit_doc_count=len(body.plan.included_doc_ids or body.seed_doc_ids or []), ) perm_service = PermissionService(redis_client=redis_client) perm = await perm_service.resolve(user) # 通过辅助函数构建 ResearchEngine，避免重复创建服务对象 engine = _build_research_engine( es_client, neo4j_client, redis_client, mysql_client, embedding_client, llm_client, ) gen = engine.rerun_section( body.task, body.plan, body.section_title, perm, section_summary=body.section_summary, source_doc_ids=body.source_doc_ids, session_id=body.session_id, ) return StreamingResponse( _sse_stream(gen), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, )