PlantAI-rag/api/raptor_routes.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RAPTOR层次索引API路由
"""

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Optional
from pathlib import Path
from src.rag.project_config import PROJECT_ROOT

router = APIRouter()

# 全局RAPTOR树实例缓存
_raptor_trees: Dict[str, any] = {}


class BuildTreeRequest(BaseModel):
    doc_id: str
    chunks: List[Dict]
    max_levels: int = 3
    summarization_ratio: float = 0.3


class BuildTreeResponse(BaseModel):
    status: str
    root_id: str
    total_nodes: int
    level_stats: Dict
    message: str


class RaptorSearchRequest(BaseModel):
    query: str
    doc_id: Optional[str] = None  # None表示搜索所有树
    top_k: int = 5
    strategy: str = "top_down"  # top_down, cross_level, leaf_only
    min_level: Optional[int] = None
    max_level: Optional[int] = None


class RaptorSearchResponse(BaseModel):
    query: str
    results: List[Dict]
    total_results: int
    strategy: str
    levels_covered: List[str]


class TreeStatsResponse(BaseModel):
    doc_id: str
    total_nodes: int
    root_id: str
    level_stats: Dict
    tree_depth: int


class MultiLevelSearchRequest(BaseModel):
    query: str
    doc_id: Optional[str] = None
    levels: List[int] = [0, 1, 2]


class MultiLevelSearchResponse(BaseModel):
    query: str
    results_by_level: Dict[str, List[Dict]]
    total_results: int


@router.post("/raptor/build", response_model=BuildTreeResponse)
async def build_raptor_tree(request: BuildTreeRequest):
    """构建RAPTOR树"""
    try:
        from src.rag.advanced.raptor import RaptorTreeBuilder
        
        builder = RaptorTreeBuilder(
            max_levels=request.max_levels,
            summarization_ratio=request.summarization_ratio
        )
        
        root_id = builder.build_tree(
            chunks=request.chunks,
            doc_id=request.doc_id
        )
        
        # 缓存树
        _raptor_trees[request.doc_id] = builder
        
        # 获取统计
        stats = builder.get_tree_stats()
        
        return BuildTreeResponse(
            status="success",
            root_id=root_id,
            total_nodes=stats['total_nodes'],
            level_stats=stats['levels'],
            message=f"RAPTOR树构建成功: {request.doc_id}"
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"构建树失败: {str(e)}")


@router.post("/raptor/search", response_model=RaptorSearchResponse)
async def search_raptor(request: RaptorSearchRequest):
    """RAPTOR层次检索"""
    try:
        from src.rag.advanced.raptor import RaptorRetriever
        
        if not _raptor_trees:
            raise HTTPException(status_code=404, detail="没有可用的RAPTOR树，请先构建")
        
        all_results = []
        levels_covered = set()
        
        # 搜索指定的树或所有树
        trees_to_search = [request.doc_id] if request.doc_id else list(_raptor_trees.keys())
        
        for doc_id in trees_to_search:
            if doc_id not in _raptor_trees:
                continue
            
            builder = _raptor_trees[doc_id]
            retriever = RaptorRetriever(builder)
            
            results = retriever.search(
                query=request.query,
                top_k=request.top_k,
                strategy=request.strategy,
                min_level=request.min_level,
                max_level=request.max_level
            )
            
            for result in results:
                all_results.append({
                    'node_id': result.node.node_id,
                    'level': result.level,
                    'score': result.score,
                    'content': result.node.content[:500],
                    'summary': result.node.summary,
                    'is_summary': result.is_summary,
                    'path': result.path.path,
                    'doc_id': doc_id
                })
                levels_covered.add(result.level)
        
        # 排序
        all_results.sort(key=lambda x: x['score'], reverse=True)
        
        return RaptorSearchResponse(
            query=request.query,
            results=all_results[:request.top_k],
            total_results=len(all_results),
            strategy=request.strategy,
            levels_covered=list(levels_covered)
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"检索失败: {str(e)}")


@router.post("/raptor/search/multi-level", response_model=MultiLevelSearchResponse)
async def search_multi_level(request: MultiLevelSearchRequest):
    """多层级检索"""
    try:
        from src.rag.advanced.raptor import RaptorRetriever
        
        if not _raptor_trees:
            raise HTTPException(status_code=404, detail="没有可用的RAPTOR树")
        
        results_by_level = {}
        total = 0
        
        trees_to_search = [request.doc_id] if request.doc_id else list(_raptor_trees.keys())
        
        for doc_id in trees_to_search:
            if doc_id not in _raptor_trees:
                continue
            
            builder = _raptor_trees[doc_id]
            retriever = RaptorRetriever(builder)
            
            multi_level_results = retriever.get_multi_level_results(
                query=request.query,
                levels=request.levels
            )
            
            for level_num, results in multi_level_results.items():
                level_key = f"level_{level_num}"
                if level_key not in results_by_level:
                    results_by_level[level_key] = []
                
                for result in results:
                    results_by_level[level_key].append({
                        'node_id': result.node.node_id,
                        'level': result.level,
                        'score': result.score,
                        'content': result.node.content[:300],
                        'summary': result.node.summary,
                        'doc_id': doc_id
                    })
                    total += 1
        
        return MultiLevelSearchResponse(
            query=request.query,
            results_by_level=results_by_level,
            total_results=total
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"多层级检索失败: {str(e)}")


@router.get("/raptor/tree/{doc_id}/stats", response_model=TreeStatsResponse)
async def get_tree_stats(doc_id: str):
    """获取树统计信息"""
    try:
        if doc_id not in _raptor_trees:
            raise HTTPException(status_code=404, detail=f"树不存在: {doc_id}")
        
        builder = _raptor_trees[doc_id]
        stats = builder.get_tree_stats()
        
        # 计算树深度
        max_height = 0
        for node in builder.nodes.values():
            max_height = max(max_height, node.get_height())
        
        return TreeStatsResponse(
            doc_id=doc_id,
            total_nodes=stats['total_nodes'],
            root_id=stats['root_id'],
            level_stats=stats['levels'],
            tree_depth=max_height
        )
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计失败: {str(e)}")


@router.get("/raptor/tree/{doc_id}/structure")
async def get_tree_structure(doc_id: str):
    """获取树结构"""
    try:
        if doc_id not in _raptor_trees:
            raise HTTPException(status_code=404, detail=f"树不存在: {doc_id}")
        
        builder = _raptor_trees[doc_id]
        
        # 构建树结构
        def build_node_tree(node_id: str) -> Dict:
            node = builder.get_node(node_id)
            if not node:
                return {}
            
            children = []
            for child_id in node.children_ids:
                child_tree = build_node_tree(child_id)
                if child_tree:
                    children.append(child_tree)
            
            return {
                'id': node.node_id,
                'level': node.level.value,
                'summary': node.summary[:200] if node.summary else node.content[:200],
                'children': children
            }
        
        structure = build_node_tree(builder.root_id)
        
        return {
            'doc_id': doc_id,
            'root_id': builder.root_id,
            'structure': structure
        }
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取结构失败: {str(e)}")


@router.post("/raptor/tree/{doc_id}/save")
async def save_tree(doc_id: str, filepath: Optional[str] = None):
    """保存树到文件"""
    try:
        if doc_id not in _raptor_trees:
            raise HTTPException(status_code=404, detail=f"树不存在: {doc_id}")
        
        builder = _raptor_trees[doc_id]
        
        if not filepath:
            filepath = str(PROJECT_ROOT / "data" / "raptor" / f"{doc_id}_tree.json")
        
        builder.save(filepath)
        
        return {
            'status': 'success',
            'doc_id': doc_id,
            'filepath': filepath
        }
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"保存失败: {str(e)}")


@router.post("/raptor/tree/load")
async def load_tree(filepath: str):
    """从文件加载树"""
    try:
        from src.rag.advanced.raptor import RaptorTreeBuilder
        
        builder = RaptorTreeBuilder.load(filepath)
        
        # 从文件路径提取doc_id
        doc_id = Path(filepath).stem.replace('_tree', '')
        _raptor_trees[doc_id] = builder
        
        return {
            'status': 'success',
            'doc_id': doc_id,
            'root_id': builder.root_id,
            'total_nodes': len(builder.nodes)
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"加载失败: {str(e)}")


@router.get("/raptor/trees")
async def list_trees():
    """列出所有可用的树"""
    return {
        'trees': [
            {
                'doc_id': doc_id,
                'root_id': builder.root_id,
                'total_nodes': len(builder.nodes)
            }
            for doc_id, builder in _raptor_trees.items()
        ]
    }


@router.delete("/raptor/tree/{doc_id}")
async def delete_tree(doc_id: str):
    """删除树"""
    try:
        if doc_id not in _raptor_trees:
            raise HTTPException(status_code=404, detail=f"树不存在: {doc_id}")
        
        del _raptor_trees[doc_id]
        
        return {
            'status': 'success',
            'doc_id': doc_id,
            'message': '树已删除'
        }
    
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除失败: {str(e)}")