PlantAI-rag/api/kg_rag_routes.py

"""
KG-RAG融合API路由 - 方案A
"""

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

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from src.rag.kg_rag_fusion import KGRAGFusion, KGRAGContext, KGEntity, KGRelation

router = APIRouter(prefix="/kg-rag", tags=["KG-RAG Fusion"])

# 全局实例
_fusion_system = None

def get_fusion_system():
    """获取融合系统实例"""
    global _fusion_system
    if _fusion_system is None:
        _fusion_system = KGRAGFusion()
    return _fusion_system


# ============ 请求/响应模型 ============

class KGRAGQueryRequest(BaseModel):
    """KG-RAG查询请求"""
    question: str
    top_k: int = 5
    include_answer: bool = True


class EntityInfo(BaseModel):
    """实体信息"""
    name: str
    type: str


class RelationInfo(BaseModel):
    """关系信息"""
    source: str
    target: str
    relation_type: str
    evidence: Optional[str] = ""
    source_paper: Optional[str] = ""


class ChunkInfo(BaseModel):
    """文本块信息"""
    content: str
    score: float
    source: str
    chunk_id: str


class KGRAGQueryResponse(BaseModel):
    """KG-RAG查询响应"""
    question: str
    entities: List[EntityInfo]
    kg_relations: List[RelationInfo]
    vector_chunks: List[ChunkInfo]
    fused_context: str
    sources: List[str]
    answer: Optional[str] = None


class EntityExtractRequest(BaseModel):
    """实体提取请求"""
    text: str


class EntityExtractResponse(BaseModel):
    """实体提取响应"""
    entities: List[EntityInfo]


# ============ API端点 ============

@router.post("/chat", response_model=KGRAGQueryResponse)
async def kg_rag_chat(request: KGRAGQueryRequest):
    """
    KG-RAG融合查询主接口
    
    流程:
    1. 实体识别
    2. KG关系查询
    3. 向量检索
    4. 融合生成
    """
    try:
        system = get_fusion_system()
        
        # 执行查询
        result = system.query(request.question, top_k=request.top_k)
        
        # 转换响应格式
        response = KGRAGQueryResponse(
            question=result.question,
            entities=[EntityInfo(name=e.name, type=e.type) for e in result.entities],
            kg_relations=[RelationInfo(
                source=r.source,
                target=r.target,
                relation_type=r.relation_type,
                evidence=r.evidence,
                source_paper=r.source_paper
            ) for r in result.kg_relations],
            vector_chunks=[ChunkInfo(
                content=c.get("content", ""),
                score=c.get("score", 0.0),
                source=c.get("source", ""),
                chunk_id=c.get("chunk_id", "")
            ) for c in result.vector_chunks],
            fused_context=result.fused_context,
            sources=result.sources
        )
        
        # 可选生成答案
        if request.include_answer:
            from src.rag.kg_rag_fusion import AnswerGenerator
            generator = AnswerGenerator()
            answer_result = generator.generate(request.question, result.fused_context)
            response.answer = answer_result.get("answer", "")
        
        return response
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"KG-RAG查询失败: {str(e)}")


@router.post("/entities", response_model=EntityExtractResponse)
async def extract_entities(request: EntityExtractRequest):
    """
    实体识别接口
    """
    try:
        from src.rag.kg_rag_fusion import EntityRecognizer
        recognizer = EntityRecognizer()
        entities = recognizer.recognize(request.text)
        
        return EntityExtractResponse(
            entities=[EntityInfo(name=e.name, type=e.type) for e in entities]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/kg-relations/{entity_name}")
async def get_entity_relations(entity_name: str, depth: int = 1):
    """
    查询KG中实体的关系
    """
    try:
        from src.rag.kg_rag_fusion import KGRetriever
        retriever = KGRetriever()
        relations = retriever.query_entity_relations(entity_name, depth=depth)
        retriever.close()
        
        return {
            "entity": entity_name,
            "relations": [
                {
                    "source": r.source,
                    "target": r.target,
                    "relation_type": r.relation_type,
                    "evidence": r.evidence,
                    "source_paper": r.source_paper
                }
                for r in relations
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/context/{question}")
async def get_kg_context(question: str):
    """
    获取KG增强上下文（不生成答案）
    """
    try:
        system = get_fusion_system()
        result = system.query(question, top_k=5)
        
        return {
            "question": question,
            "entities": [{"name": e.name, "type": e.type} for e in result.entities],
            "kg_relations_count": len(result.kg_relations),
            "chunks_count": len(result.vector_chunks),
            "fused_context": result.fused_context[:2000]  # 截断
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "neo4j": "bolt://localhost:7687",
        "components": {
            "entity_recognizer": "ready",
            "kg_retriever": "ready",
            "vector_retriever": "ready",
            "context_fusion": "ready"
        }
    }


print("[KG-RAG] 方案A路由加载完成")