#!/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)}")