PlantAI-rag/api/temporal_routes.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
时序知识图谱API路由
"""

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

router = APIRouter()

# 全局时序知识图谱实例
_temporal_kg = None


def get_temporal_kg():
    """获取时序知识图谱实例"""
    global _temporal_kg
    if _temporal_kg is None:
        from src.rag.temporal import TemporalKnowledgeGraph
        _temporal_kg = TemporalKnowledgeGraph(name="gene_family_temporal_kg")
    return _temporal_kg


class CreateEntityRequest(BaseModel):
    entity_id: str
    name: str
    entity_type: str
    start_year: Optional[int] = None
    end_year: Optional[int] = None
    properties: Dict = {}


class CreateRelationRequest(BaseModel):
    relation_id: str
    source_id: str
    target_id: str
    relation_type: str
    year: Optional[int] = None
    properties: Dict = {}


class TimeRangeQueryRequest(BaseModel):
    start_year: Optional[int] = None
    end_year: Optional[int] = None
    entity_types: List[str] = []
    relation_types: List[str] = []


class TimelineQueryRequest(BaseModel):
    entity_id: Optional[str] = None
    topic: Optional[str] = None
    start_year: Optional[int] = None
    end_year: Optional[int] = None


@router.post("/temporal/entity")
async def create_entity(request: CreateEntityRequest):
    """创建时序实体"""
    try:
        from src.rag.temporal import TemporalEntity, TimePeriod, TimePoint
        
        kg = get_temporal_kg()
        
        # 构建时间范围
        valid_time = None
        if request.start_year:
            valid_time = TimePeriod(
                start=TimePoint.from_year(request.start_year),
                end=TimePoint.from_year(request.end_year) if request.end_year else TimePoint.from_year(request.start_year)
            )
        
        entity = TemporalEntity(
            entity_id=request.entity_id,
            name=request.name,
            entity_type=request.entity_type,
            valid_time=valid_time,
            properties=request.properties
        )
        
        kg.add_entity(entity)
        
        return {
            "status": "success",
            "entity_id": request.entity_id,
            "message": f"实体 {request.name} 创建成功"
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建实体失败: {str(e)}")


@router.post("/temporal/relation")
async def create_relation(request: CreateRelationRequest):
    """创建时序关系"""
    try:
        from src.rag.temporal import TemporalRelation, TimePoint
        
        kg = get_temporal_kg()
        
        # 构建时间戳
        timestamp = None
        if request.year:
            timestamp = TimePoint.from_year(request.year)
        
        relation = TemporalRelation(
            relation_id=request.relation_id,
            source_id=request.source_id,
            target_id=request.target_id,
            relation_type=request.relation_type,
            timestamp=timestamp,
            properties=request.properties
        )
        
        kg.add_relation(relation)
        
        return {
            "status": "success",
            "relation_id": request.relation_id,
            "message": f"关系 {request.relation_type} 创建成功"
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建关系失败: {str(e)}")


@router.post("/temporal/query/range")
async def query_by_time_range(request: TimeRangeQueryRequest):
    """按时间范围查询"""
    try:
        from src.rag.temporal import TimeRangeQuery, TimePoint
        
        kg = get_temporal_kg()
        
        query = TimeRangeQuery(
            start=TimePoint.from_year(request.start_year) if request.start_year else None,
            end=TimePoint.from_year(request.end_year) if request.end_year else None,
            entity_types=request.entity_types,
            relation_types=request.relation_types
        )
        
        result = kg.query(query)
        
        return {
            "query": query.to_dict(),
            "entities": [e.to_dict() for e in result["entities"]],
            "relations": [r.to_dict() for r in result["relations"]],
            "total_entities": len(result["entities"]),
            "total_relations": len(result["relations"])
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@router.post("/temporal/timeline")
async def get_timeline(request: TimelineQueryRequest):
    """获取时间线"""
    try:
        kg = get_temporal_kg()
        
        timeline = kg.get_timeline(
            entity_id=request.entity_id,
            topic=request.topic,
            start_year=request.start_year,
            end_year=request.end_year
        )
        
        return timeline.to_dict()
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取时间线失败: {str(e)}")


@router.get("/temporal/timeline/visualization")
async def get_timeline_visualization(
    topic: Optional[str] = None,
    start_year: Optional[int] = None,
    end_year: Optional[int] = None
):
    """获取时间线可视化数据"""
    try:
        kg = get_temporal_kg()
        
        visual_data = kg.visualize_timeline(topic=topic)
        
        return {
            "topic": topic,
            "start_year": start_year,
            "end_year": end_year,
            "visualization_data": visual_data
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取可视化数据失败: {str(e)}")


@router.get("/temporal/stats")
async def get_stats():
    """获取统计信息"""
    try:
        kg = get_temporal_kg()
        stats = kg.get_stats()
        
        return stats.to_dict()
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计失败: {str(e)}")


@router.get("/temporal/entity/{entity_id}")
async def get_entity(entity_id: str):
    """获取实体详情"""
    try:
        kg = get_temporal_kg()
        entity = kg.get_entity(entity_id)
        
        if not entity:
            raise HTTPException(status_code=404, detail=f"实体不存在: {entity_id}")
        
        return entity.to_dict()
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取实体失败: {str(e)}")


@router.get("/temporal/entity/{entity_id}/history")
async def get_entity_history(entity_id: str):
    """获取实体历史"""
    try:
        kg = get_temporal_kg()
        timeline = kg.get_timeline(entity_id=entity_id)
        
        return timeline.to_dict()
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取历史失败: {str(e)}")


@router.get("/temporal/entities/by-type/{entity_type}")
async def get_entities_by_type(entity_type: str):
    """按类型获取实体"""
    try:
        kg = get_temporal_kg()
        entities = kg.get_entities_by_type(entity_type)
        
        return {
            "entity_type": entity_type,
            "entities": [e.to_dict() for e in entities],
            "total": len(entities)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取实体失败: {str(e)}")


@router.get("/temporal/entities/by-year/{year}")
async def get_entities_by_year(year: int):
    """按年份获取实体"""
    try:
        kg = get_temporal_kg()
        entities = kg.get_entities_by_year(year)
        
        return {
            "year": year,
            "entities": [e.to_dict() for e in entities],
            "total": len(entities)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取实体失败: {str(e)}")


@router.post("/temporal/consistency/check")
async def check_consistency():
    """检查图谱一致性"""
    try:
        kg = get_temporal_kg()
        conflicts = kg.check_consistency()
        
        return {
            "consistent": len(conflicts) == 0,
            "conflicts": conflicts,
            "total_conflicts": len(conflicts)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"检查一致性失败: {str(e)}")


@router.post("/temporal/infer")
async def infer_relations():
    """推断新关系"""
    try:
        kg = get_temporal_kg()
        inferred = kg.infer_relations()
        
        return {
            "inferred_relations": [r.to_dict() for r in inferred],
            "total_inferred": len(inferred)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"推断关系失败: {str(e)}")


@router.post("/temporal/save")
async def save_graph(filepath: str = "/tmp/temporal_kg.json"):
    """保存图谱"""
    try:
        kg = get_temporal_kg()
        kg.save(filepath)
        
        return {
            "status": "success",
            "filepath": filepath
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"保存失败: {str(e)}")


@router.post("/temporal/load")
async def load_graph(filepath: str = "/tmp/temporal_kg.json"):
    """加载图谱"""
    try:
        global _temporal_kg
        from src.rag.temporal import TemporalKnowledgeGraph
        
        _temporal_kg = TemporalKnowledgeGraph.load(filepath)
        
        return {
            "status": "success",
            "filepath": filepath,
            "name": _temporal_kg.name
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"加载失败: {str(e)}")


@router.post("/temporal/demo/populate")
async def populate_demo_data():
    """填充演示数据"""
    try:
        from src.rag.temporal import (
            TemporalEntity, TemporalRelation,
            TimePoint, TimePeriod
        )
        
        kg = get_temporal_kg()
        
        # 添加演示实体
        demo_entities = [
            ("spl_arabidopsis", "SPL基因家族(拟南芥)", "gene_family", 2018, 2023),
            ("spl_rice", "SPL基因家族(水稻)", "gene_family", 2019, 2023),
            ("spl_poplar", "SPL基因家族(杨树)", "gene_family", 2020, 2023),
            ("spl3", "SPL3基因", "gene", 2018, None),
            ("spl9", "SPL9基因", "gene", 2019, None),
            ("flowering", "开花调控", "process", 2018, 2023),
            ("stress_response", "胁迫响应", "process", 2020, 2023),
        ]
        
        for entity_id, name, entity_type, start_year, end_year in demo_entities:
            valid_time = TimePeriod(
                start=TimePoint.from_year(start_year),
                end=TimePoint.from_year(end_year) if end_year else TimePoint.from_year(2023)
            )
            
            entity = TemporalEntity(
                entity_id=entity_id,
                name=name,
                entity_type=entity_type,
                valid_time=valid_time,
                properties={"topic": "SPL基因家族"}
            )
            kg.add_entity(entity)
        
        # 添加演示关系
        demo_relations = [
            ("rel1", "spl3", "flowering", "regulates", 2019),
            ("rel2", "spl9", "flowering", "regulates", 2020),
            ("rel3", "spl_arabidopsis", "stress_response", "involved_in", 2021),
            ("rel4", "spl3", "spl9", "interacts_with", 2020),
        ]
        
        for rel_id, source, target, rel_type, year in demo_relations:
            relation = TemporalRelation(
                relation_id=rel_id,
                source_id=source,
                target_id=target,
                relation_type=rel_type,
                timestamp=TimePoint.from_year(year),
                properties={}
            )
            kg.add_relation(relation)
        
        return {
            "status": "success",
            "entities_added": len(demo_entities),
            "relations_added": len(demo_relations),
            "message": "演示数据已填充"
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"填充演示数据失败: {str(e)}")