hello-agents/Co-creation-projects/DeLunnLi-PaperGraph/backend/app/api/routes/search.py

"""智能搜索 API 路由 —— 自然语言论文搜索与 SSE 流式响应."""

from __future__ import annotations

import asyncio
import logging
import os
import time
from typing import Any, Dict, List, Optional

import anyio
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from ...agents.search_agent import SearchAgent, SearchIntent, get_search_agent
from ...api.dependencies import get_searcher
from ...api.search_route_support import (
    ToolCallInfo,
    last_pipeline_tool_error,
    normalize_tool_calls,
    track_tool_call,
    user_facing_error_message,
)
from ...models.schemas import Paper
from ...services.papers.papers_converters import litpapers_to_api_papers
from ...services.retrieval.search_plan import ResolvedSearchPlan
from ...services.retrieval.search_pipeline import run_search_pipeline_async
from ..tool_events import ToolCallTracker, sse_pack

router = APIRouter(prefix="/papers", tags=["智能搜索"])
logger = logging.getLogger(__name__)

_SSE_QUEUE_SIZE = 128
_SEARCH_AGENT_WALL_SEC = max(
    120.0,
    min(900.0, float(os.getenv("PAPERGRAPH_SEARCH_AGENT_WALL_SEC") or 420.0)),
)
_SEARCH_AGENT_INIT_SEC = 25.0
_PREFIX_CONFLICT_MARKER = "为您找到"


class SearchAgentMessage(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000, description="用户搜索需求")
    mode: str = Field(default="accuracy", description="accuracy=准确性优先, novelty=新颖性优先")
    use_tavily: bool = Field(default=False, description="是否使用 Tavily 预搜索")
    history: List[Dict[str, str]] = Field(default_factory=list, description="对话历史")


class SearchAgentResponse(BaseModel):
    success: bool
    response: str
    search_params: Optional[Dict[str, Any]] = None
    tool_calls: List[ToolCallInfo] = Field(default_factory=list)
    papers: List[Paper] = Field(default_factory=list)
    total: int = 0
    message: Optional[str] = None


def _search_params_from_intent(intent: SearchIntent, **extra: Any) -> Dict[str, Any]:
    yf, yt = intent.year_from, intent.year_to
    if isinstance(yf, int) and isinstance(yt, int) and yf > yt:
        yf, yt = yt, yf
    out: Dict[str, Any] = {
        "query": intent.query,
        "keywords": intent.keywords,
        "authors": getattr(intent, "authors", []) or [],
        "arxiv_id_list": getattr(intent, "arxiv_id_list", []) or [],
        "venues": intent.venues,
        "year_from": yf,
        "year_to": yt,
        "sort": intent.sort,
        "use_llm_rank": intent.use_llm_rank,
        "rerank_recall_max": intent.rerank_recall_max,
        "ranking_rationale": intent.ranking_rationale or None,
    }
    out.update(extra)
    return out


def _generate_suggestions(intent: SearchIntent, papers: List[Paper]) -> List[str]:
    if len(papers) < 5:
        return [f"扩大搜索：尝试「{intent.query}」而不限定会议"]
    return []


def _strip_conflicting_search_summary_prefix(text: str) -> str:
    t = (text or "").strip()
    if not t or _PREFIX_CONFLICT_MARKER not in t:
        return t
    return t[: t.find(_PREFIX_CONFLICT_MARKER)].rstrip()


def _explanation_with_suggestions(
    agent: SearchAgent,
    intent: SearchIntent,
    papers: List[Paper],
    profile_mode: str,
    *,
    prefix_plain: str = "",
) -> str:
    base = _strip_conflicting_search_summary_prefix((prefix_plain or "").strip())
    expl = agent.explain_results(intent, papers, profile_mode)
    explanation = base + "\n\n---\n\n" + expl if (papers and base) else (base or expl)
    if papers:
        sug = _generate_suggestions(intent, papers)
        if sug:
            explanation += "\n\n🔍 **您可以这样优化**：\n" + "".join(
                f"{i}. {s}\n" for i, s in enumerate(sug, 1)
            )
    return explanation


def _error_response(msg: str) -> SearchAgentResponse:
    return SearchAgentResponse(
        success=False,
        response=user_facing_error_message(msg),
        message=msg,
    )


async def _prepare_agent_and_query(request: SearchAgentMessage) -> tuple[SearchAgent, str]:
    try:
        with anyio.fail_after(_SEARCH_AGENT_INIT_SEC):
            agent = await anyio.to_thread.run_sync(get_search_agent)
    except TimeoutError as exc:
        logger.warning("search-agent init timeout after %.0fs", _SEARCH_AGENT_INIT_SEC, exc_info=exc)
        raise HTTPException(status_code=504, detail="search_agent_init_timeout") from exc
    return agent, (request.message or "").strip()


async def _run_search_agent_core(
    *,
    agent: SearchAgent,
    request: SearchAgentMessage,
    merged_query: str,
    searcher: Any,
) -> SearchAgentResponse:
    tool_calls: List[ToolCallInfo] = []

    intent = agent.understand_intent(merged_query, request.mode)
    with track_tool_call(tool_calls, "understand_intent", {"query": merged_query}) as tc:
        tc.result_summary = f"sort={intent.sort}, venues={intent.venues}, yf={intent.year_from}, kw={intent.keywords}"

    plan = ResolvedSearchPlan.from_search_intent(intent)
    with track_tool_call(tool_calls, "search_pipeline", {"query": intent.query or merged_query}) as tc:
        tc.result_summary = "intent→SearchPlan→pipeline"
        mr = int(getattr(plan, "max_results", None) or intent.max_results or 10)
        pip = await run_search_pipeline_async(searcher=searcher, plan=plan, max_results=mr)
        tc.result_summary = f"ranked={len(pip.ranked or [])}"

    papers = litpapers_to_api_papers(rp.paper for rp in (pip.ranked or []))
    prefix = f"为您找到 {len(papers)} 篇论文。" if papers else "未找到相关论文。"

    pipeline_err = last_pipeline_tool_error(tool_calls)
    if not papers and pipeline_err:
        body = (
            "主检索未成功返回论文（多源召回或精排阶段出错），与「数据库里确实没有匹配文献」不同。\n\n"
            f"**错误摘要**：{pipeline_err}\n\n"
            "建议稍后重试，或略微改写查询；若频繁出现请查看服务端日志。"
        )
        return SearchAgentResponse(
            success=False,
            response=body,
            search_params=_search_params_from_intent(intent, mode=request.mode),
            tool_calls=normalize_tool_calls(tool_calls),
            papers=[],
            total=0,
            message="search_pipeline_error",
        )

    body = _explanation_with_suggestions(agent, intent, papers, request.mode, prefix_plain=prefix)
    return SearchAgentResponse(
        success=True,
        response=body,
        search_params=_search_params_from_intent(intent, mode=request.mode),
        tool_calls=normalize_tool_calls(tool_calls),
        papers=papers,
        total=len(papers),
    )


async def _search_agent_impl(request: SearchAgentMessage, searcher: Any):
    try:
        agent, merged_query = await _prepare_agent_and_query(request)
        resp = await asyncio.wait_for(
            _run_search_agent_core(
                agent=agent,
                request=request,
                merged_query=merged_query,
                searcher=searcher,
            ),
            timeout=_SEARCH_AGENT_WALL_SEC,
        )
        return resp, None
    except asyncio.TimeoutError as exc:
        logger.warning("search-agent timeout after %.0fs", _SEARCH_AGENT_WALL_SEC, exc_info=exc)
        return _error_response("search_agent_timeout"), HTTPException(status_code=504, detail="search_agent_timeout")
    except HTTPException as e:
        return _error_response(str(e.detail or "search_agent_http_error")), e
    except Exception:
        logger.exception("search-agent unexpected failure")
        return _error_response("search_agent_internal_error"), None


@router.post("/search-agent/stream")
async def search_agent_chat_stream(
    request: SearchAgentMessage,
    searcher=Depends(get_searcher),
):
    async def gen():
        # SSE 流式生成器：通过 anyio 内存通道实现事件驱动的流式推送
        send, recv = anyio.create_memory_object_stream(_SSE_QUEUE_SIZE)
        tracker = ToolCallTracker(sink=lambda ev: send.send_nowait(ev))
        tracker.emit("status", {"message": "search-agent 已接入，开始处理"})

        async def run_once() -> SearchAgentResponse:
            tracker.emit("status", {"message": f"初始化 SearchAgent（mode={request.mode})"})
            t0 = time.time()
            tracker.emit("status", {"message": "正在检索论文…"})
            resp, exc = await _search_agent_impl(request, searcher)
            if exc:
                code = (
                    str(exc.detail or "search_agent_http_error")
                    if isinstance(exc, HTTPException)
                    else "search_agent_internal_error"
                )
                msg = user_facing_error_message(code)
                tracker.emit("error", {"message": msg})
                if not isinstance(exc, HTTPException):
                    logger.exception("search-agent stream run loop failed")
                return _error_response(msg)
            tracker.emit(
                "final",
                {"elapsed_ms": int((time.time() - t0) * 1000), "success": bool(resp.success)},
            )
            return resp

        box: Dict[str, Any] = {"resp": None}
        cancelled_exc = anyio.get_cancelled_exc_class()
        try:
            async with anyio.create_task_group() as tg:

                async def _run() -> None:
                    try:
                        box["resp"] = await run_once()
                    finally:
                        try:
                            await send.aclose()
                        except Exception:
                            pass

                tg.start_soon(_run)

                async for ev in recv:
                    try:
                        yield sse_pack(ev)
                    except (cancelled_exc, asyncio.CancelledError):
                        return
                    except Exception:
                        return
        except (cancelled_exc, asyncio.CancelledError):
            return
        finally:
            try:
                await recv.aclose()
            except Exception:
                pass

        resp: Optional[SearchAgentResponse] = box.get("resp")
        if resp is None:
            resp = _error_response("search_agent_stream_incomplete")
            tracker.emit("error", {"message": resp.message or "search_agent_stream_incomplete"})
        yield sse_pack({"type": "final_result", "result": resp.model_dump(mode="json")})

    return StreamingResponse(
        gen(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )