Lightstar/backend/api/research.py

"""
研究相关API路由
"""

from fastapi import APIRouter, HTTPException, BackgroundTasks, Query
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional
import logging

from backend.core.research import ResearchAgent
from backend.core.clustering import PaperClusterer
from backend.core.report import ReportGenerator
from backend.config import MAX_SEARCH_RESULTS

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/research", tags=["research"])

# 数据模型
class ResearchRequest(BaseModel):
    research_intent: str = Field(..., description="用户的研究意图")
    max_results: int = Field(MAX_SEARCH_RESULTS, description="最大检索结果数量")

class KeywordsRequest(BaseModel):
    research_intent: str = Field(..., description="用户的研究意图")

class PaperSearchRequest(BaseModel):
    keywords: List[str] = Field(..., description="检索关键词")
    max_results: int = Field(MAX_SEARCH_RESULTS, description="最大检索结果数量")

class ClusterRequest(BaseModel):
    papers: List[Dict[str, Any]] = Field(..., description="要聚类的论文")
    num_clusters: int = Field(0, description="聚类数量，0表示自动确定")

class ReportRequest(BaseModel):
    research_intent: str = Field(..., description="研究意图")
    keywords: List[str] = Field(..., description="关键词")
    papers: List[Dict[str, Any]] = Field(..., description="检索到的论文")
    clusters: Optional[Dict[str, Any]] = Field(None, description="聚类结果")

# 全局实例
research_agent = ResearchAgent()
paper_clusterer = PaperClusterer()
report_generator = ReportGenerator()

# 路由定义
@router.post("/process")
async def process_research(request: ResearchRequest):
    """处理完整的研究流程"""
    try:
        logger.info(f"Processing research intent: {request.research_intent}")
        result = await research_agent.process_research_intent(
            research_intent=request.research_intent,
            max_results=request.max_results
        )
        return result
    except Exception as e:
        logger.error(f"Error in research process: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/extract-keywords")
async def extract_keywords(request: KeywordsRequest):
    """从研究意图中提取关键词"""
    try:
        logger.info(f"Extracting keywords from: {request.research_intent}")
        result = await research_agent.llm_client.extract_keywords(
            research_topic=request.research_intent
        )
        return {"keywords": result}
    except Exception as e:
        logger.error(f"Error extracting keywords: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/search-papers")
async def search_papers(request: PaperSearchRequest):
    """根据关键词检索论文"""
    try:
        logger.info(f"Searching papers with keywords: {request.keywords}")
        papers = []
        for keyword in request.keywords:
            results = await research_agent.arxiv_client.search_papers(
                query=keyword,
                max_results=max(3, request.max_results // len(request.keywords))
            )
            papers.extend(results)
        
        # 去重
        unique_papers = []
        paper_ids = set()
        for paper in papers:
            if paper["id"] not in paper_ids:
                unique_papers.append(paper)
                paper_ids.add(paper["id"])
        
        return {"papers": unique_papers, "count": len(unique_papers)}
    except Exception as e:
        logger.error(f"Error searching papers: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/cluster-papers")
async def cluster_papers(request: ClusterRequest):
    """对论文进行聚类分析"""
    try:
        logger.info(f"Clustering {len(request.papers)} papers")
        result = paper_clusterer.cluster_papers(
            papers=request.papers,
            num_clusters=request.num_clusters
        )
        return result
    except Exception as e:
        logger.error(f"Error clustering papers: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/generate-report")
async def generate_report(request: ReportRequest):
    """生成研究报告"""
    try:
        logger.info(f"Generating report for: {request.research_intent}")
        report = await report_generator.generate_report(
            research_intent=request.research_intent,
            keywords=request.keywords,
            papers=request.papers,
            clusters=request.clusters
        )
        return report
    except Exception as e:
        logger.error(f"Error generating report: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))