PlantAI-rag/api/advanced_rag_routes.py

"""
高级RAG API路由
集成查询扩展、实体链接、引用溯源、动态知识图谱、跨论文融合、Neo4j图算法
"""

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Optional

router = APIRouter()

# ========== 查询扩展API ==========

class QueryExpansionRequest(BaseModel):
    query: str
    num_expansions: int = 5
    strategy: str = "expansion"  # expansion, intent_focused, keyword_combination


class QueryExpansionResponse(BaseModel):
    original: str
    rewritten: str
    alternatives: List[str]
    intent: str
    confidence: float
    keywords: List[str]
    synonyms: Dict[str, List[str]]


@router.post("/advanced/query/expand")
async def expand_query_endpoint(request: QueryExpansionRequest):
    """查询扩展端点"""
    try:
        from src.rag.advanced import QueryExpander, QueryRewriter
        
        expander = QueryExpander()
        expansion = expander.expand_query(request.query, request.num_expansions)
        
        return QueryExpansionResponse(
            original=request.query,
            rewritten=expansion.expanded_queries[0] if expansion.expanded_queries else request.query,
            alternatives=expansion.expanded_queries[1:] if len(expansion.expanded_queries) > 1 else [],
            intent=expansion.intent.value,
            confidence=expansion.intent_confidence,
            keywords=expansion.keywords,
            synonyms=expansion.synonyms
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询扩展失败: {str(e)}")


@router.post("/advanced/query/rewrite")
async def rewrite_query_endpoint(request: QueryExpansionRequest):
    """查询重写端点"""
    try:
        from src.rag.advanced import QueryRewriter
        
        rewriter = QueryRewriter()
        result = rewriter.rewrite(request.query, request.strategy)
        
        return result
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询重写失败: {str(e)}")


# ========== 实体链接API ==========

class EntityLinkingRequest(BaseModel):
    text: str


class EntityLinkingResponse(BaseModel):
    entities: List[Dict]
    linked_count: int
    total_count: int


@router.post("/advanced/entity/link")
async def entity_linking_endpoint(request: EntityLinkingRequest):
    """实体链接端点"""
    try:
        from src.rag.advanced import EntityLinker
        
        linker = EntityLinker()
        linked_entities = linker.link_entities(request.text)
        
        entities = []
        for le in linked_entities:
            entity_info = {
                "mention": le.mention.text,
                "start_pos": le.mention.start_pos,
                "end_pos": le.mention.end_pos,
                "context": le.mention.context[:100],
                "linked": le.kb_entity is not None,
                "link_confidence": le.link_confidence,
                "disambiguation_score": le.disambiguation_score,
            }
            
            if le.kb_entity:
                entity_info.update({
                    "kb_id": le.kb_entity.entity_id,
                    "kb_name": le.kb_entity.name,
                    "kb_type": le.kb_entity.entity_type.value,
                    "kb_definition": le.kb_entity.definition,
                    "kb_aliases": le.kb_entity.aliases[:5],
                })
            
            entities.append(entity_info)
        
        linked_count = sum(1 for e in entities if e["linked"])
        
        return EntityLinkingResponse(
            entities=entities,
            linked_count=linked_count,
            total_count=len(entities)
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"实体链接失败: {str(e)}")


# ========== 引用溯源API ==========

class CitationRequest(BaseModel):
    answer: str
    sources: List[Dict]
    min_citations: int = 3


class CitationResponse(BaseModel):
    answer: str
    citations: List[Dict]
    confidence: float
    coverage_score: float
    total_citations: int


@router.post("/advanced/citation/add")
async def add_citations_endpoint(request: CitationRequest):
    """添加引用溯源端点"""
    try:
        from src.rag.advanced import CitationIntegrator
        
        integrator = CitationIntegrator()
        result = integrator.integrate_citations(
            request.answer,
            request.sources,
            request.min_citations
        )
        
        return CitationResponse(
            answer=result["answer"],
            citations=result["citations"],
            confidence=result["confidence"],
            coverage_score=result["coverage_score"],
            total_citations=result["total_citations"]
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"添加引用失败: {str(e)}")


# ========== 动态知识图谱API ==========

class DynamicKGRequest(BaseModel):
    sources: List[Dict]


class DynamicKGResponse(BaseModel):
    entities: List[Dict]
    relations: List[Dict]
    total_entities: int
    total_relations: int
    conflicts: int
    resolved: int
    needs_review: int
    graph_data: Dict


@router.post("/advanced/kg/build")
async def build_dynamic_kg_endpoint(request: DynamicKGRequest):
    """动态构建知识图谱端点"""
    try:
        from src.rag.advanced import DynamicKGBuilder
        
        builder = DynamicKGBuilder()
        result = builder.build_from_search_results(request.sources)
        
        # 获取图谱数据
        graph_data = builder.get_graph_data()
        
        return DynamicKGResponse(
            entities=result["entities"],
            relations=result["relations"],
            total_entities=result["total_entities"],
            total_relations=result["total_relations"],
            conflicts=result["conflicts"],
            resolved=result["resolved"],
            needs_review=result["needs_review"],
            graph_data=graph_data
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"构建知识图谱失败: {str(e)}")


# ========== 跨论文融合API ==========

class CrossPaperFusionRequest(BaseModel):
    papers: List[Dict]


class CrossPaperFusionResponse(BaseModel):
    entities: List[Dict]
    relations: List[Dict]
    stats: Dict


@router.post("/advanced/fusion/cross-paper")
async def cross_paper_fusion_endpoint(request: CrossPaperFusionRequest):
    """跨论文知识融合端点"""
    try:
        from src.rag.advanced import CrossPaperFusion
        
        fusion = CrossPaperFusion()
        result = fusion.fuse_knowledge(request.papers)
        
        return CrossPaperFusionResponse(
            entities=result["entities"],
            relations=result["relations"],
            stats=result["stats"]
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"跨论文融合失败: {str(e)}")


@router.post("/advanced/fusion/provenance/{entity_id}")
async def get_provenance_endpoint(entity_id: str, request: CrossPaperFusionRequest):
    """获取实体溯源信息端点"""
    try:
        from src.rag.advanced import CrossPaperFusion
        
        fusion = CrossPaperFusion()
        fusion.fuse_knowledge(request.papers)
        provenance = fusion.get_entity_provenance(entity_id)
        
        return provenance
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取溯源信息失败: {str(e)}")


# ========== Neo4j图算法API ==========

class ShortestPathRequest(BaseModel):
    start_node: str
    end_node: str
    relationship_type: Optional[str] = None


class ShortestPathResponse(BaseModel):
    status: str
    path: Optional[Dict]
    message: Optional[str]


@router.post("/advanced/graph/shortest-path")
async def shortest_path_endpoint(request: ShortestPathRequest):
    """最短路径分析端点"""
    try:
        from src.rag.advanced import GraphAlgorithmAPI
        
        api = GraphAlgorithmAPI()
        result = api.shortest_path_endpoint(
            request.start_node,
            request.end_node,
            request.relationship_type
        )
        
        return ShortestPathResponse(**result)
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"最短路径分析失败: {str(e)}")


@router.get("/advanced/graph/communities")
async def community_detection_endpoint():
    """社区发现端点"""
    try:
        from src.rag.advanced import GraphAlgorithmAPI
        
        api = GraphAlgorithmAPI()
        result = api.community_detection_endpoint()
        
        return result
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"社区发现失败: {str(e)}")


class CentralityRequest(BaseModel):
    algorithm: str = "degree"  # degree, betweenness, pagerank
    top_k: int = 10


@router.post("/advanced/graph/centrality")
async def centrality_analysis_endpoint(request: CentralityRequest):
    """中心性分析端点"""
    try:
        from src.rag.advanced import GraphAlgorithmAPI
        
        api = GraphAlgorithmAPI()
        result = api.centrality_analysis_endpoint(request.algorithm, request.top_k)
        
        return result
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"中心性分析失败: {str(e)}")


@router.get("/advanced/graph/stats")
async def graph_stats_endpoint():
    """图统计端点"""
    try:
        from src.rag.advanced import GraphAlgorithmAPI
        
        api = GraphAlgorithmAPI()
        result = api.graph_stats_endpoint()
        
        return result
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取图统计失败: {str(e)}")


# ========== 综合RAG API ==========

class AdvancedRAGRequest(BaseModel):
    query: str
    enable_expansion: bool = True
    enable_entity_linking: bool = True
    enable_citation: bool = True
    enable_dynamic_kg: bool = True
    top_k: int = 5


class AdvancedRAGResponse(BaseModel):
    query: str
    expanded_query: Optional[str]
    answer: str
    entities: List[Dict]
    citations: List[Dict]
    knowledge_graph: Optional[Dict]
    confidence: float


@router.post("/advanced/rag")
async def advanced_rag_endpoint(request: AdvancedRAGRequest):
    """高级RAG综合端点"""
    try:
        from src.rag.advanced import (
            QueryExpander,
            EntityLinker,
            CitationIntegrator,
            DynamicKGBuilder
        )
        
        # 1. 查询扩展
        expanded_query = request.query
        if request.enable_expansion:
            expander = QueryExpander()
            expansion = expander.expand_query(request.query, num_expansions=3)
            expanded_query = expansion.expanded_queries[0] if expansion.expanded_queries else request.query
        
        # 2. 实体链接
        entities = []
        if request.enable_entity_linking:
            linker = EntityLinker()
            linked = linker.link_entities(request.query)
            entities = [
                {
                    "mention": le.mention.text,
                    "linked": le.kb_entity is not None,
                    "kb_name": le.kb_entity.name if le.kb_entity else None,
                    "kb_type": le.kb_entity.entity_type.value if le.kb_entity else None,
                    "confidence": le.link_confidence
                }
                for le in linked
            ]
        
        # 3. 模拟检索结果（实际应该调用检索器）
        mock_sources = [
            {
                "doc_id": f"PMID{10000000 + i}",
                "content": f"Sample content for {request.query}",
                "score": 0.9 - i * 0.1,
                "chunk_id": f"chunk_{i}",
                "metadata": {"title": f"Paper {i}", "year": "2023"}
            }
            for i in range(request.top_k)
        ]
        
        # 4. 生成答案
        answer = f"基于检索到的{len(mock_sources)}条文献，关于'{request.query}'的分析如下..."
        
        # 5. 添加引用
        citations = []
        if request.enable_citation:
            integrator = CitationIntegrator()
            citation_result = integrator.integrate_citations(answer, mock_sources, min_citations=3)
            answer = citation_result["answer"]
            citations = citation_result["citations"]
        
        # 6. 构建动态知识图谱
        kg_data = None
        if request.enable_dynamic_kg:
            builder = DynamicKGBuilder()
            kg_result = builder.build_from_search_results(mock_sources)
            kg_data = builder.get_graph_data()
        
        return AdvancedRAGResponse(
            query=request.query,
            expanded_query=expanded_query if request.enable_expansion else None,
            answer=answer,
            entities=entities,
            citations=citations,
            knowledge_graph=kg_data,
            confidence=0.85
        )
    
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"高级RAG处理失败: {str(e)}")