api_server.py

"""
FastAPI Web 服务 - 学术文献调研 Agent
支持 WebSocket 流式响应
"""

import os
import json
import asyncio
from typing import AsyncGenerator, Dict, Any
from datetime import datetime
from contextlib import asynccontextmanager

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 导入 Agent 组件
from research_agent import create_workflow, ResearchState
from llm_config import get_model_config, create_llm_client


# ============== 数据模型 ==============
class ResearchRequest(BaseModel):
    topic: str


class ResearchResponse(BaseModel):
    status: str
    message: str
    data: Dict[str, Any] = {}


# ============== 全局状态 ==============
class ConnectionManager:
    """WebSocket 连接管理器"""
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        # 安全检查：避免重复移除或移除不存在的连接
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def send_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)


manager = ConnectionManager()


# ============== FastAPI 应用 ==============
@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    print("=" * 60)
    print("Research Agent API Server Starting...")
    print("=" * 60)
    try:
        config = get_model_config()
        print(f"[LLM] Model: {config.name} ({config.format})")
        if config.base_url:
            print(f"[LLM] API URL: {config.base_url}")
    except Exception as e:
        print(f"[WARNING] LLM Config Error: {e}")
    print("=" * 60)
    yield
    print("\nServer shutting down...")


app = FastAPI(
    title="Research Agent API",
    description="学术文献调研 Agent Web 服务",
    version="1.0.0",
    lifespan=lifespan
)

# CORS 配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============== API 路由 ==============
@app.get("/")
async def root():
    """返回前端页面"""
    return FileResponse("web/index.html")


@app.get("/api/health")
async def health_check():
    """健康检查"""
    try:
        config = get_model_config()
        return {
            "status": "healthy",
            "model": config.name,
            "format": config.format,
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        return JSONResponse(
            status_code=503,
            content={"status": "unhealthy", "error": str(e)}
        )


@app.post("/api/research")
async def create_research(request: ResearchRequest):
    """同步研究接口 (非流式)"""
    try:
        workflow = create_workflow()
        result = workflow.invoke({
            "topic": request.topic,
            "queries": [],
            "documents": [],
            "final_report": ""
        })

        return {
            "status": "success",
            "data": {
                "topic": request.topic,
                "queries": result.get("queries", []),
                "documents": result.get("documents", []),
                "report": result.get("final_report", "")
            }
        }
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={"status": "error", "message": str(e)}
        )


# ============== WebSocket 流式接口 ==============
async def run_research_stream(topic: str) -> AsyncGenerator[Dict[str, Any], None]:
    """
    流式运行研究任务，产生进度事件
    """
    from langgraph.graph import StateGraph, END
    from langchain_core.prompts import ChatPromptTemplate

    # 初始化
    yield {"type": "init", "message": "初始化研究任务...", "topic": topic}
    await asyncio.sleep(0.1)

    # 手动执行每个节点以便流式输出
    llm = create_llm_client()

    # Step 1: Planner
    yield {"type": "progress", "step": "planner", "status": "running", "message": "正在生成检索词..."}

    prompt = ChatPromptTemplate.from_template("""You are an academic search assistant.

Given a research topic:
{topic}

Generate 3–5 high-quality English search queries suitable for:
- Arxiv
- PubMed
- Google Scholar

Requirements:
- Use precise academic terminology
- Cover different aspects if possible

Return ONLY a list of queries, one per line, without numbering or bullet points.
""")

    try:
        chain = prompt | llm
        result = chain.invoke({"topic": topic})
        queries = [q.strip() for q in result.content.strip().split('\n') if q.strip()]
        queries = queries[:5]

        yield {
            "type": "progress",
            "step": "planner",
            "status": "completed",
            "message": f"生成 {len(queries)} 个检索词",
            "queries": queries
        }
    except Exception as e:
        yield {"type": "error", "step": "planner", "message": str(e)}
        return

    await asyncio.sleep(0.2)

    # Step 2: Searcher
    yield {"type": "progress", "step": "searcher", "status": "running", "message": "正在检索文献..."}

    try:
        from research_agent import search_arxiv, search_pubmed

        all_documents = []
        for query in queries:
            yield {"type": "progress", "step": "searcher", "status": "running", "message": f"搜索: {query}"}

            arxiv_papers = search_arxiv(query, max_results=2)
            all_documents.extend(arxiv_papers)

            pubmed_papers = search_pubmed(query, max_results=2)
            all_documents.extend(pubmed_papers)

            await asyncio.sleep(0.1)

        # 去重
        seen_titles = set()
        unique_docs = []
        for doc in all_documents:
            title_lower = doc["title"].lower()
            if title_lower not in seen_titles:
                seen_titles.add(title_lower)
                unique_docs.append(doc)

        yield {
            "type": "progress",
            "step": "searcher",
            "status": "completed",
            "message": f"检索到 {len(unique_docs)} 篇文献",
            "documents": unique_docs
        }
    except Exception as e:
        yield {"type": "error", "step": "searcher", "message": str(e)}
        return

    await asyncio.sleep(0.2)

    # Step 3: Writer
    yield {"type": "progress", "step": "writer", "status": "running", "message": "正在生成综述报告..."}

    try:
        if not unique_docs:
            report = "未检索到相关文献，请尝试其他关键词。"
        else:
            docs_text = "\n\n".join([
                f"[{i+1}] {doc['title']}\nSource: {doc['source']} ({doc.get('published', 'N/A')})\n{doc['summary']}\nURL: {doc['url']}"
                for i, doc in enumerate(unique_docs)
            ])

            writer_prompt = ChatPromptTemplate.from_template("""You are writing an academic research summary.

Based ONLY on the provided papers:

{documents}

Rules:
1. Do NOT fabricate any papers
2. Only use the given documents
3. Cite using [number] format
4. If information is insufficient, say "limited evidence"

Output structure in Chinese:

## 研究背景
简要介绍该领域的背景和重要性。

## 最新进展
总结文献中报告的最新研究进展。

## 关键方法
归纳文献中使用的主要研究方法和技术。

## 存在挑战
分析该领域面临的主要挑战和局限。

## 文献列表
- [1] 标题 - 来源 - URL
（列出所有参考文献）
""")

            # 在异步上下文中执行可能耗时的 LLM 调用，避免事件循环阻塞导致 WebSocket 被上游超时断开
            async def invoke_writer():
                return await asyncio.get_running_loop().run_in_executor(
                    None,
                    lambda: (writer_prompt | llm).invoke({"documents": docs_text})
                )

            writer_task = asyncio.create_task(invoke_writer())

            # 心跳保活：每隔10秒发送一次状态，避免中间无消息导致连接被代理/浏览器关闭
            while not writer_task.done():
                await asyncio.sleep(10)
                yield {"type": "heartbeat", "message": "正在生成报告，请稍候（连接保活中）..."}

            result = await writer_task
            report = result.content

        yield {
            "type": "progress",
            "step": "writer",
            "status": "completed",
            "message": "报告生成完成",
            "report": report
        }

        # 最终完成
        yield {
            "type": "complete",
            "message": "研究任务完成",
            "data": {
                "topic": topic,
                "queries": queries,
                "documents": unique_docs,
                "report": report
            }
        }

    except Exception as e:
        yield {"type": "error", "step": "writer", "message": str(e)}


@app.websocket("/ws/research")
async def websocket_research(websocket: WebSocket):
    """WebSocket 流式研究接口"""
    await manager.connect(websocket)
    # 配置心跳超时检测（30秒）
    websocket.ping_interval = 20
    websocket.ping_timeout = 30

    try:
        while True:
            # 接收消息
            try:
                data = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=60.0  # 接收消息超时60秒
                )
            except asyncio.TimeoutError:
                # 发送 ping 保活
                try:
                    await websocket.send_json({"type": "ping"})
                    continue
                except:
                    break

            request = json.loads(data)

            # 处理 ping 响应
            if request.get("type") == "pong":
                continue

            topic = request.get("topic", "").strip()

            if not topic:
                await websocket.send_json({
                    "type": "error",
                    "message": "请提供研究主题"
                })
                continue

            # 流式执行研究
            async for event in run_research_stream(topic):
                await websocket.send_json(event)

    except WebSocketDisconnect:
        pass
    except Exception as e:
        try:
            await websocket.send_json({
                "type": "error",
                "message": f"服务器错误: {str(e)}"
            })
        except:
            pass
    finally:
        manager.disconnect(websocket)


# ============== 静态文件服务 ==============
if os.path.exists("web"):
    app.mount("/static", StaticFiles(directory="web/static"), name="static")


# ============== 启动入口 ==============
if __name__ == "__main__":
    import uvicorn

    # 确保 web 目录存在
    os.makedirs("web/static", exist_ok=True)

    uvicorn.run(
        "api_server:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )