ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/structured-index/main.py

@@ -0,0 +1,202 @@
+"""
+结构化索引工具的主入口：构建 / 查询 RAPTOR 与 GraphRAG 索引，或运行离线对比演示。
+
+说明：RAPTOR、GraphRAG 的**索引构建**需要调用 LLM（实体抽取、递归摘要），因此
+build / query 依赖 OPENAI_API_KEY 及相应重型依赖（umap、sentence-transformers 等）。
+若只想直观理解「结构化索引解决了扁平检索的什么问题」，可运行无需 API 的 `demo` 子命令。
+"""
+
+import argparse
+import asyncio
+from pathlib import Path
+import json
+import sys
+
+from loguru import logger
+
+
+async def build_indexes(file_path: Path, index_type: str = "both",
+                        output: str = None):
+    """Build RAPTOR and/or GraphRAG indexes from a document."""
+    # 重型依赖延迟导入：保证 --help / demo 在缺少 umap 等依赖时仍可用
+    from config import get_raptor_config, get_graphrag_config
+    from raptor_indexer import RaptorIndexer
+    from graphrag_indexer import GraphRAGIndexer
+    from document_processor import DocumentProcessor
+
+    logger.info(f"Building {index_type} index(es) from {file_path}")
+
+    # Process document
+    processor = DocumentProcessor()
+    text = await processor.process_file(file_path)
+    logger.info(f"Processed document: {len(text)} characters")
+
+    all_stats = {}
+
+    # Build RAPTOR index
+    if index_type in ["raptor", "both"]:
+        logger.info("Building RAPTOR tree index...")
+        raptor_config = get_raptor_config()
+        raptor = RaptorIndexer(raptor_config)
+        raptor.build_index(text)
+        raptor.save_index()
+        stats = raptor.get_tree_statistics()
+        all_stats["raptor"] = stats
+        logger.info(f"RAPTOR index built: {stats}")
+
+    # Build GraphRAG index
+    if index_type in ["graphrag", "both"]:
+        logger.info("Building GraphRAG knowledge graph...")
+        graphrag_config = get_graphrag_config()
+        graphrag = GraphRAGIndexer(graphrag_config)
+        graphrag.build_knowledge_graph(text)
+        graphrag.detect_communities()
+        graphrag.hierarchical_summarization()
+        graphrag.save_index()
+        stats = graphrag.get_graph_statistics()
+        all_stats["graphrag"] = stats
+        logger.info(f"GraphRAG index built: {stats}")
+
+    if output:
+        with open(output, "w", encoding="utf-8") as f:
+            json.dump(all_stats, f, ensure_ascii=False, indent=2)
+        logger.info(f"索引统计已写入：{output}")
+
+    logger.info("Indexing complete!")
+
+
+async def query_indexes(query: str, index_type: str = "both", top_k: int = 5,
+                        multi_hop: int = 0):
+    """Query RAPTOR and/or GraphRAG indexes."""
+    from config import get_raptor_config, get_graphrag_config
+    from raptor_indexer import RaptorIndexer
+    from graphrag_indexer import GraphRAGIndexer
+
+    results = {}
+
+    # Query RAPTOR
+    if index_type in ["raptor", "both"]:
+        try:
+            raptor_config = get_raptor_config()
+            raptor = RaptorIndexer(raptor_config)
+            raptor.load_index()
+            raptor_results = raptor.search(query, top_k)
+            results["raptor"] = raptor_results
+            logger.info(f"RAPTOR returned {len(raptor_results)} results")
+        except Exception as e:
+            logger.error(f"Error querying RAPTOR: {e}")
+
+    # Query GraphRAG
+    if index_type in ["graphrag", "both"]:
+        try:
+            graphrag_config = get_graphrag_config()
+            graphrag = GraphRAGIndexer(graphrag_config)
+            graphrag.load_index()
+            graphrag_results = graphrag.search(query, top_k)
+            results["graphrag"] = graphrag_results
+            logger.info(f"GraphRAG returned {len(graphrag_results)} results")
+
+            # 多跳关系检索：以召回的最佳实体为起点，沿关系边遍历
+            if multi_hop > 0 and graphrag_results:
+                start = next((r.get("name") for r in graphrag_results
+                              if r.get("type") == "entity"), None)
+                if start:
+                    paths = graphrag.multi_hop_search(start, max_hops=multi_hop)
+                    results["graphrag_multi_hop"] = paths
+                    logger.info(f"GraphRAG multi-hop from '{start}' "
+                                f"returned {len(paths)} paths")
+        except Exception as e:
+            logger.error(f"Error querying GraphRAG: {e}")
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="结构化索引工具：在统一框架下构建并查询 RAPTOR（树状层次）与 "
+                    "GraphRAG（实体关系图）索引，对应本书实验 3-7。",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    subparsers = parser.add_subparsers(dest="command", help="要执行的子命令")
+
+    # Build command
+    build_parser = subparsers.add_parser(
+        "build", help="从文档构建结构化索引（需要 OPENAI_API_KEY）")
+    build_parser.add_argument("file", type=str,
+                              help="待索引的文档路径（支持 .pdf/.txt/.md/.html）")
+    build_parser.add_argument("--type", choices=["raptor", "graphrag", "both"],
+                              default="both", help="要构建的索引类型（默认 both）")
+    build_parser.add_argument("--output", type=str, default=None,
+                              help="将索引统计信息写入指定 JSON 文件")
+
+    # Query command
+    query_parser = subparsers.add_parser(
+        "query", help="查询已构建的索引（需要 OPENAI_API_KEY 及已有索引）")
+    query_parser.add_argument("query", type=str, help="检索查询语句")
+    query_parser.add_argument("--type", choices=["raptor", "graphrag", "both"],
+                              default="both", help="要查询的索引类型（默认 both）")
+    query_parser.add_argument("--top-k", type=int, default=5,
+                              help="返回结果条数（默认 5）")
+    query_parser.add_argument("--multi-hop", type=int, default=0, metavar="N",
+                              help="对 GraphRAG 额外执行 N 跳关系遍历（0 表示关闭）")
+    query_parser.add_argument("--output", type=str, default=None,
+                              help="将查询结果写入指定 JSON 文件")
+
+    # Demo command（离线，无需 API）
+    demo_parser = subparsers.add_parser(
+        "demo", help="离线对比演示：结构化索引 vs 扁平检索（无需 API Key）")
+    demo_parser.add_argument("--query", type=str, default=None,
+                             help="自定义查询；缺省时运行内置的三组对比查询")
+    demo_parser.add_argument("--top-k", type=int, default=3,
+                             help="扁平检索展示的结果条数（默认 3）")
+    demo_parser.add_argument("--output", type=str, default=None,
+                             help="将演示结果写入指定 JSON 文件")
+
+    # Server command
+    subparsers.add_parser("serve", help="启动 HTTP API 服务")
+
+    args = parser.parse_args()
+
+    if args.command == "build":
+        asyncio.run(build_indexes(Path(args.file), args.type, args.output))
+    elif args.command == "query":
+        results = asyncio.run(query_indexes(args.query, args.type, args.top_k,
+                                            args.multi_hop))
+
+        # Display results
+        for index_type, index_results in results.items():
+            print(f"\n{index_type.upper()} Results:")
+            print("-" * 50)
+            if index_type == "graphrag_multi_hop":
+                for i, r in enumerate(index_results, 1):
+                    chain = r["path"][0]["source"]
+                    for step in r["path"]:
+                        chain += f" --{step['relation']}--> {step['target']}"
+                    print(f"\n{i}. [{r['hops']} 跳] {chain}")
+                continue
+            for i, result in enumerate(index_results, 1):
+                print(f"\n{i}. Score: {result.get('score', 'N/A'):.3f}")
+                if 'summary' in result:
+                    print(f"   Summary: {result['summary'][:200]}...")
+                elif 'description' in result:
+                    print(f"   Description: {result['description'][:200]}...")
+                if 'level' in result:
+                    print(f"   Level: {result['level']}")
+
+        if args.output:
+            with open(args.output, "w", encoding="utf-8") as f:
+                json.dump(results, f, ensure_ascii=False, indent=2, default=str)
+            print(f"\n查询结果已写入：{args.output}")
+    elif args.command == "demo":
+        from structured_vs_flat_demo import run_demo
+        run_demo(top_k=args.top_k, custom_query=args.query, output=args.output)
+    elif args.command == "serve":
+        from api_service import run_server
+        run_server()
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()