structured-index¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/structured-index
项目说明¶
Structured Document Indexing with RAPTOR and GraphRAG¶
This educational project demonstrates two advanced approaches for indexing and querying large technical documents:
- RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) - Creates a hierarchical tree structure with recursive summarization
- GraphRAG (Graph-based Retrieval Augmented Generation) - Builds a knowledge graph with entities, relationships, and community detection
Both approaches are optimized for handling large technical documentation like the Intel® 64 and IA-32 Architectures Software Developer's Manual.
Features¶
RAPTOR Tree-Based Indexing¶
- Hierarchical tree structure with multiple levels of abstraction
- Recursive summarization for information compression
- Multi-level search capability (leaf nodes to root summaries)
- Clustering-based node grouping using Gaussian Mixture Models
- UMAP dimensionality reduction for efficient clustering
GraphRAG Knowledge Graph Indexing¶
- Entity and relationship extraction using LLMs
- Community detection for identifying related concepts
- Hierarchical community summarization
- Graph-based search with multiple strategies
- Multi-hop relation traversal (
GraphRAGIndexer.multi_hop_search)：沿关系边做多跳遍历， 回答扁平向量检索无法表达的「A 通过什么与 B 相连」这类关系性问题（对应书中「多跳关系推理」） - Support for different entity types (instructions, registers, features, etc.)
HTTP API Service¶
- RESTful API for building and querying indexes
- Support for file uploads
- Asynchronous processing for large documents
- Hybrid search across both index types
- Real-time index statistics and status monitoring
Installation¶
-
Clone the repository and navigate to the project directory:
-
Install dependencies:
-
Copy and configure the environment file:
Quick Start¶
命令行接口（CLI）¶
所有子命令都提供中文 --help：python main.py --help、python main.py demo --help 等。
usage: main.py [-h] {build,query,demo,serve} ...
build 从文档构建结构化索引（需要 OPENAI_API_KEY）
query 查询已构建的索引（需要 OPENAI_API_KEY 及已有索引）
demo 离线对比演示：结构化索引 vs 扁平检索（无需 API Key）
serve 启动 HTTP API 服务
0. 离线对比演示（无需 API Key，推荐先跑这个）¶
这是理解实验 3-8 的最快入口：它用一个手工整理的 Intel x86 SIMD 小知识库， 直观对比「扁平检索」与「结构化索引」在三类查询上的差异，全程无需 OpenAI API：
# 运行内置的三组对比查询（多跳关系推理 / 跨节点综合对比 / 多层次导航）
python main.py demo
# 自定义查询，同时给出扁平检索与图多跳遍历两种视角
python main.py demo --query "VADDPS 用到哪个寄存器"
# 把结果写入 JSON
python main.py demo --output demo_result.json
演示输出示例（多跳关系推理，扁平检索答不了、图检索沿关系边可达）：
【查询 1｜多跳关系推理】运行 ADDPS 指令前，操作系统必须把哪个控制寄存器位置 1？
-- 扁平检索（按词面相似度返回独立片段）--
1. [control-bit] CR4.OSFXSR (score=0.459)
...
✗ 只能召回词面相近的孤立片段，无法把 ADDPS 与某个控制位「连」起来。
-- 结构化图检索（沿关系边多跳遍历）--
ADDPS --属于--> SSE --需要启用--> CR4.OSFXSR
✓ 答案：CR4.OSFXSR（从 ADDPS 经 2 跳可达）
说明：
build与query需要真实索引，而索引构建依赖 LLM（实体抽取、递归摘要）， 因此需要设置OPENAI_API_KEY（嵌入用本地sentence-transformers）。demo则把索引结果预先手工写好，让读者无需 API Key 也能看到结构化索引解决的问题。
1. 构建索引（需要 OPENAI_API_KEY）¶
# 同时构建 RAPTOR 与 GraphRAG 索引
python main.py build path/to/document.pdf
# 只构建 RAPTOR，或只构建 GraphRAG
python main.py build path/to/document.pdf --type raptor
python main.py build path/to/document.pdf --type graphrag
# 将索引统计写入 JSON
python main.py build path/to/document.pdf --output stats.json
2. 查询索引¶
# 查询两种索引
python main.py query "What are the MOV instruction variants?"
# 指定索引类型与返回条数
python main.py query "explain SSE instructions" --type raptor --top-k 10
# GraphRAG 多跳关系遍历：以召回的最佳实体为起点，沿关系边走 N 跳
python main.py query "SSE registers" --type graphrag --multi-hop 2
# 将查询结果写入 JSON
python main.py query "control registers" --output result.json
3. 启动 API 服务¶
Using the HTTP API¶
-
Start the server:
-
Build an index via API:
# Upload a file and build index curl -X POST "http://localhost:4242/upload" \ -F "file=@path/to/intel_manual.pdf" \ -F "index_type=both" # Build from text curl -X POST "http://localhost:4242/build" \ -H "Content-Type: application/json" \ -d '{ "file_path": "/path/to/document.pdf", "index_type": "both", "force_rebuild": false }' -
Query the index:
-
Check index status:
API Endpoints¶
| Endpoint | Method | Description |
|---|---|---|
/ |
GET | API information and available endpoints |
/build |
POST | Build index from text or file |
/upload |
POST | Upload file and build index |
/query |
POST | Query the indexes |
/status |
GET | Get index status |
/statistics |
GET | Get detailed index statistics |
/indexes |
DELETE | Clear indexes |
Project Structure¶
structured-index/
├── config.py # Configuration management
├── raptor_indexer.py # RAPTOR tree-based indexing
├── graphrag_indexer.py # GraphRAG graph-based indexing
├── document_processor.py # Document parsing and preprocessing
├── api_service.py # HTTP API service
├── structured_vs_flat_demo.py # 离线对比演示：结构化索引 vs 扁平检索（无需 API）
├── main.py # CLI interface
├── requirements.txt # Python dependencies
├── env.example # Environment variables template
├── indexes/ # Saved index files
│ ├── raptor/ # RAPTOR index storage
│ └── graphrag/ # GraphRAG index storage
└── cache/ # Temporary cache directory
How It Works¶
RAPTOR Indexing Process¶
- Text Chunking: Document is split into manageable chunks with overlap
- Embedding Generation: Each chunk is converted to vector embeddings
- Leaf Node Creation: Chunks become leaf nodes with summaries
- Hierarchical Clustering: Nodes are clustered using GMM
- Parent Node Generation: Clusters are summarized to create parent nodes
- Tree Building: Process repeats for multiple levels
- Multi-level Search: Queries search across all tree levels
GraphRAG Indexing Process¶
- Entity Extraction: LLM identifies entities (instructions, registers, etc.)
- Relationship Discovery: Connections between entities are extracted
- Graph Construction: NetworkX graph built from entities and relationships
- Community Detection: Related entities grouped using Leiden/Louvain
- Community Summarization: Each community gets a descriptive summary
- Hierarchical Aggregation: Similar communities are merged and summarized
- Graph Search: Queries match against entities and community summaries
Example: Processing Intel Architecture Manual¶
import asyncio
from pathlib import Path
from config import get_raptor_config, get_graphrag_config
from raptor_indexer import RaptorIndexer
from graphrag_indexer import GraphRAGIndexer
from document_processor import DocumentProcessor
async def process_intel_manual():
# Process the Intel manual PDF
processor = DocumentProcessor()
intel_manual_path = Path("intel_x86_64_manual.pdf")
text = await processor.process_file(intel_manual_path)
# Build RAPTOR index
raptor_config = get_raptor_config()
raptor = RaptorIndexer(raptor_config)
raptor.build_index(text)
raptor.save_index()
# Build GraphRAG index
graphrag_config = get_graphrag_config()
graphrag = GraphRAGIndexer(graphrag_config)
graphrag.build_knowledge_graph(text)
graphrag.detect_communities()
graphrag.hierarchical_summarization()
graphrag.save_index()
# Example queries
queries = [
"What are the different addressing modes?",
"Explain SIMD instructions",
"How does the MOV instruction work?",
"What are control registers?"
]
for query in queries:
print(f"\nQuery: {query}")
print("-" * 50)
# RAPTOR search
raptor_results = raptor.search(query, top_k=3)
print("RAPTOR Results:")
for r in raptor_results:
print(f" Level {r['level']}: {r['summary'][:100]}...")
# GraphRAG search
graphrag_results = graphrag.search(query, top_k=3)
print("\nGraphRAG Results:")
for r in graphrag_results:
if r['type'] == 'entity':
print(f" Entity: {r['name']} - {r['description'][:100]}...")
else:
print(f" Community: {r['summary'][:100]}...")
# Run the example
asyncio.run(process_intel_manual())
Advanced Configuration¶
RAPTOR Parameters¶
chunk_size: Size of text chunks (default: 1000 words)chunk_overlap: Overlap between chunks (default: 200 words)tree_depth: Maximum tree depth (default: 3)summarization_length: Target summary length (default: 200 words)
GraphRAG Parameters¶
chunk_size: Size of text chunks (default: 1200 words)max_knowledge_triples: Max triples per chunk (default: 10)community_detection_algorithm: "leiden" or "louvain"summarization_model: Model for generating summaries
Performance Considerations¶
- Large Documents: Processing 5000+ page documents may take significant time
- API Rate Limits: Consider OpenAI API rate limits when processing
- Memory Usage: Large graphs require substantial memory
- Caching: Results are cached to improve subsequent query performance
- Parallel Processing: Use the API service for concurrent operations
Integration with Agentic RAG¶
This project provides backend services for the agentic-rag project. See the agentic-rag README for integration details.
Troubleshooting¶
- Out of Memory: Reduce chunk_size or process document in sections
- API Errors: Check API keys and rate limits
- Slow Indexing: Consider using faster/smaller models for initial testing
- Import Errors: Ensure all dependencies are installed correctly
References¶
OpenRouter 通用回退 / Universal OpenRouter fallback¶
This experiment now supports a universal OpenRouter fallback for its chat LLM.
- If the primary provider key (e.g.
MOONSHOT_API_KEY/KIMI_API_KEY/OPENAI_API_KEY/DOUBAO_API_KEY…) is present, behavior is unchanged. - Else if
OPENROUTER_API_KEYis set, the chat LLM is automatically routed through OpenRouter (https://openrouter.ai/api/v1). Model names are mapped automatically:gpt-*/o1-*→openai/…,claude-*→anthropic/claude-opus-4.8,kimi-*→moonshotai/kimi-k2.6, ids already containing/are kept as-is, and other provider-native ids (e.g.doubao-*) fall back toopenai/gpt-5.6-luna. SetOPENROUTER_MODELto force a specific OpenRouter model id. - Else a clear error lists the accepted keys.
Add OPENROUTER_API_KEY=... to your .env (see env.example) to enable it.
Note: embeddings here are local SentenceTransformers (all-MiniLM-L6-v2), so they are unaffected — only the chat LLM used for RAPTOR summarization and GraphRAG entity extraction is routed through OpenRouter.
源代码¶
api_service.py¶
"""
HTTP API service for querying RAPTOR and GraphRAG indexes.
"""
from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Literal
from pathlib import Path
import asyncio
from concurrent.futures import ThreadPoolExecutor
import uvicorn
from loguru import logger
import json
import aiofiles
import tempfile
import shutil
from config import get_raptor_config, get_graphrag_config, get_api_config
from raptor_indexer import RaptorIndexer
from graphrag_indexer import GraphRAGIndexer
from document_processor import DocumentProcessor
# Initialize FastAPI app
app = FastAPI(
title="Structured Index API",
description="API for querying RAPTOR tree-based and GraphRAG graph-based document indexes",
version="1.0.0"
)
# Thread pool for CPU-intensive operations
executor = ThreadPoolExecutor(max_workers=4)
# Global indexers
raptor_indexer: Optional[RaptorIndexer] = None
graphrag_indexer: Optional[GraphRAGIndexer] = None
document_processor: Optional[DocumentProcessor] = None
class BuildIndexRequest(BaseModel):
"""Request model for building an index."""
text: Optional[str] = Field(None, description="Text content to index")
file_path: Optional[str] = Field(None, description="Path to document file")
index_type: Literal["raptor", "graphrag", "both"] = Field("both", description="Type of index to build")
force_rebuild: bool = Field(False, description="Force rebuild even if index exists")
class QueryRequest(BaseModel):
"""Request model for querying an index."""
query: str = Field(..., description="Search query")
index_type: Literal["raptor", "graphrag", "hybrid"] = Field("hybrid", description="Index to query")
top_k: int = Field(5, description="Number of results to return")
search_type: Optional[str] = Field("hybrid", description="Search type for GraphRAG")
class IndexResponse(BaseModel):
"""Response model for index operations."""
status: str
message: str
statistics: Optional[Dict[str, Any]] = None
class QueryResponse(BaseModel):
"""Response model for query operations."""
query: str
results: List[Dict[str, Any]]
index_type: str
total_results: int
@app.on_event("startup")
async def startup_event():
"""Initialize indexers on startup."""
global raptor_indexer, graphrag_indexer, document_processor
logger.info("Initializing indexers...")
# Initialize configurations
raptor_config = get_raptor_config()
graphrag_config = get_graphrag_config()
# Initialize indexers
raptor_indexer = RaptorIndexer(raptor_config)
graphrag_indexer = GraphRAGIndexer(graphrag_config)
document_processor = DocumentProcessor()
# Try to load existing indexes
try:
if (raptor_config.index_dir / "raptor_index.pkl").exists():
raptor_indexer.load_index()
logger.info("Loaded existing RAPTOR index")
except Exception as e:
logger.warning(f"Could not load RAPTOR index: {e}")
try:
if (graphrag_config.index_dir / "graphrag_index.pkl").exists():
graphrag_indexer.load_index()
logger.info("Loaded existing GraphRAG index")
except Exception as e:
logger.warning(f"Could not load GraphRAG index: {e}")
logger.info("API service started successfully")
@app.get("/")
async def root():
"""Root endpoint."""
return {
"service": "Structured Index API",
"version": "1.0.0",
"endpoints": {
"build": "/build",
"query": "/query",
"status": "/status",
"statistics": "/statistics"
}
}
@app.post("/build", response_model=IndexResponse)
async def build_index(
request: BuildIndexRequest,
background_tasks: BackgroundTasks
):
"""Build RAPTOR and/or GraphRAG index from text or file."""
try:
# Get text content
if request.text:
text_content = request.text
elif request.file_path:
file_path = Path(request.file_path)
if not file_path.exists():
raise HTTPException(status_code=404, detail=f"File not found: {request.file_path}")
text_content = await document_processor.process_file(file_path)
else:
raise HTTPException(status_code=400, detail="Either text or file_path must be provided")
# Check if we should rebuild
if not request.force_rebuild:
existing_indexes = []
if request.index_type in ["raptor", "both"]:
if (raptor_indexer.config.index_dir / "raptor_index.pkl").exists():
existing_indexes.append("RAPTOR")
if request.index_type in ["graphrag", "both"]:
if (graphrag_indexer.config.index_dir / "graphrag_index.pkl").exists():
existing_indexes.append("GraphRAG")
if existing_indexes:
return IndexResponse(
status="exists",
message=f"Indexes already exist: {', '.join(existing_indexes)}. Use force_rebuild=true to rebuild."
)
# Build indexes in background
async def build_indexes():
results = {}
if request.index_type in ["raptor", "both"]:
logger.info("Building RAPTOR index...")
loop = asyncio.get_event_loop()
await loop.run_in_executor(executor, raptor_indexer.build_index, text_content)
await loop.run_in_executor(executor, raptor_indexer.save_index)
results["raptor"] = raptor_indexer.get_tree_statistics()
if request.index_type in ["graphrag", "both"]:
logger.info("Building GraphRAG index...")
loop = asyncio.get_event_loop()
await loop.run_in_executor(executor, graphrag_indexer.build_knowledge_graph, text_content)
await loop.run_in_executor(executor, graphrag_indexer.detect_communities)
await loop.run_in_executor(executor, graphrag_indexer.hierarchical_summarization)
await loop.run_in_executor(executor, graphrag_indexer.save_index)
results["graphrag"] = graphrag_indexer.get_graph_statistics()
return results
# Start building in background
background_tasks.add_task(build_indexes)
return IndexResponse(
status="building",
message=f"Started building {request.index_type} index(es) in background"
)
except Exception as e:
logger.error(f"Error building index: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/upload", response_model=IndexResponse)
async def upload_and_build(
file: UploadFile = File(...),
index_type: Literal["raptor", "graphrag", "both"] = "both",
background_tasks: BackgroundTasks = None
):
"""Upload a document and build index."""
try:
# Save uploaded file temporarily
with tempfile.NamedTemporaryFile(delete=False, suffix=Path(file.filename).suffix) as tmp_file:
content = await file.read()
tmp_file.write(content)
tmp_path = tmp_file.name
# Process the file
text_content = await document_processor.process_file(Path(tmp_path))
# Clean up temp file
Path(tmp_path).unlink()
# Build index
request = BuildIndexRequest(
text=text_content,
index_type=index_type,
force_rebuild=True
)
return await build_index(request, background_tasks)
except Exception as e:
logger.error(f"Error processing uploaded file: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query", response_model=QueryResponse)
async def query_index(request: QueryRequest):
"""Query the RAPTOR or GraphRAG index."""
try:
results = []
if request.index_type == "raptor":
# Query RAPTOR index
if not raptor_indexer.nodes:
raise HTTPException(status_code=404, detail="RAPTOR index not built")
loop = asyncio.get_event_loop()
raptor_results = await loop.run_in_executor(
executor,
raptor_indexer.search,
request.query,
request.top_k
)
results = raptor_results
elif request.index_type == "graphrag":
# Query GraphRAG index
if not graphrag_indexer.entities:
raise HTTPException(status_code=404, detail="GraphRAG index not built")
loop = asyncio.get_event_loop()
graphrag_results = await loop.run_in_executor(
executor,
graphrag_indexer.search,
request.query,
request.top_k,
request.search_type
)
results = graphrag_results
elif request.index_type == "hybrid":
# Query both indexes and combine results
all_results = []
# Query RAPTOR
if raptor_indexer.nodes:
loop = asyncio.get_event_loop()
raptor_results = await loop.run_in_executor(
executor,
raptor_indexer.search,
request.query,
request.top_k
)
for r in raptor_results:
r["source"] = "raptor"
all_results.extend(raptor_results)
# Query GraphRAG
if graphrag_indexer.entities:
loop = asyncio.get_event_loop()
graphrag_results = await loop.run_in_executor(
executor,
graphrag_indexer.search,
request.query,
request.top_k,
request.search_type
)
for r in graphrag_results:
r["source"] = "graphrag"
all_results.extend(graphrag_results)
# Sort by score and return top-k
all_results.sort(key=lambda x: x.get("score", 0), reverse=True)
results = all_results[:request.top_k]
return QueryResponse(
query=request.query,
results=results,
index_type=request.index_type,
total_results=len(results)
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error querying index: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/status")
async def get_status():
"""Get the status of the indexes."""
status = {
"raptor": {
"built": len(raptor_indexer.nodes) > 0 if raptor_indexer else False,
"node_count": len(raptor_indexer.nodes) if raptor_indexer else 0
},
"graphrag": {
"built": len(graphrag_indexer.entities) > 0 if graphrag_indexer else False,
"entity_count": len(graphrag_indexer.entities) if graphrag_indexer else 0,
"relationship_count": len(graphrag_indexer.relationships) if graphrag_indexer else 0
}
}
return status
@app.get("/statistics")
async def get_statistics():
"""Get detailed statistics about the indexes."""
stats = {}
if raptor_indexer and raptor_indexer.nodes:
stats["raptor"] = raptor_indexer.get_tree_statistics()
if graphrag_indexer and graphrag_indexer.entities:
stats["graphrag"] = graphrag_indexer.get_graph_statistics()
if not stats:
raise HTTPException(status_code=404, detail="No indexes built")
return stats
@app.delete("/indexes")
async def clear_indexes(index_type: Literal["raptor", "graphrag", "both"] = "both"):
"""Clear the specified indexes."""
try:
cleared = []
if index_type in ["raptor", "both"]:
raptor_indexer.nodes = {}
raptor_indexer.root_nodes = []
# Delete saved index
index_file = raptor_indexer.config.index_dir / "raptor_index.pkl"
if index_file.exists():
index_file.unlink()
cleared.append("RAPTOR")
if index_type in ["graphrag", "both"]:
graphrag_indexer.entities = {}
graphrag_indexer.relationships = []
graphrag_indexer.communities = {}
graphrag_indexer.graph.clear()
# Delete saved index
index_file = graphrag_indexer.config.index_dir / "graphrag_index.pkl"
if index_file.exists():
index_file.unlink()
cleared.append("GraphRAG")
return {
"status": "success",
"message": f"Cleared indexes: {', '.join(cleared)}"
}
except Exception as e:
logger.error(f"Error clearing indexes: {e}")
raise HTTPException(status_code=500, detail=str(e))
def run_server():
"""Run the API server."""
config = get_api_config()
logger.info(f"Starting API server on {config.host}:{config.port}")
uvicorn.run(
"api_service:app",
host=config.host,
port=config.port,
reload=config.reload
)
if __name__ == "__main__":
run_server()
config.py¶
"""
Configuration for structured index project.
"""
import os
from pathlib import Path
from typing import Optional
from dataclasses import dataclass
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def _openrouter_model_id(model) -> str:
"""Map a provider-native model name to an OpenRouter model id, used by the
universal OpenRouter fallback. An explicit OPENROUTER_MODEL env var wins."""
override = os.getenv("OPENROUTER_MODEL")
if override:
return override
m = (model or "").strip()
if not m:
return "openai/gpt-5.6-luna"
if "/" in m:
return m
ml = m.lower()
if ml.startswith(("gpt-", "o1", "o3", "o4", "chatgpt")):
return "openai/" + m
if ml.startswith("claude-"):
return "anthropic/claude-opus-4.8"
if ml.startswith("kimi"):
# kimi-k3 is not on OpenRouter; moonshotai/kimi-k2.6 is the closest hosted id.
return "moonshotai/kimi-k2.6"
return "openai/gpt-5.6-luna"
def _resolve_llm(api_key: str, *models):
"""Return (api_key, base_url, *mapped_models). When the OpenAI key is
absent but OPENROUTER_API_KEY is present, route the chat LLM (used for
RAPTOR summarization / GraphRAG entity extraction) through OpenRouter.
Embeddings here are local SentenceTransformers, so they are unaffected."""
openrouter_key = os.getenv("OPENROUTER_API_KEY")
# gpt-5.x (incl. gpt-5.6*) needs OpenAI org-verification on the direct API;
# when an OpenRouter key is present, prefer routing these ids through it.
prefer_openrouter = bool(openrouter_key) and any(
str(m or "").lower().startswith("gpt-5") for m in models)
if (not api_key or prefer_openrouter) and openrouter_key:
base_url = "https://openrouter.ai/api/v1"
return (openrouter_key, base_url,
*[_openrouter_model_id(m) for m in models])
return (api_key, None, *models)
@dataclass
class RaptorConfig:
"""Configuration for RAPTOR tree-based indexing."""
openai_api_key: str
model_name: str = "gpt-5.6-luna"
embedding_model: str = "text-embedding-3-small"
max_tokens: int = 2048
temperature: float = 0.1
chunk_size: int = 1000
chunk_overlap: int = 200
tree_depth: int = 3
summarization_length: int = 200
index_dir: Path = Path("indexes/raptor")
base_url: Optional[str] = None
@dataclass
class GraphRAGConfig:
"""Configuration for GraphRAG graph-based indexing."""
llm_api_key: str
llm_model: str = "gpt-5.6-luna"
embedding_model: str = "text-embedding-3-small"
chunk_size: int = 1200
chunk_overlap: int = 100
max_knowledge_triples: int = 10
community_detection_algorithm: str = "leiden"
summarization_model: str = "gpt-5.6-luna"
index_dir: Path = Path("indexes/graphrag")
cache_dir: Path = Path("cache/graphrag")
base_url: Optional[str] = None
@dataclass
class APIConfig:
"""Configuration for HTTP API service."""
host: str = "127.0.0.1"
port: int = 4242
reload: bool = True
max_results: int = 10
timeout_seconds: int = 30
def get_raptor_config() -> RaptorConfig:
"""Get RAPTOR configuration from environment."""
api_key, base_url, model_name = _resolve_llm(
os.getenv("OPENAI_API_KEY", ""),
os.getenv("RAPTOR_MODEL", "gpt-5.6-luna"),
)
return RaptorConfig(
openai_api_key=api_key,
model_name=model_name,
embedding_model=os.getenv("RAPTOR_EMBEDDING_MODEL", "text-embedding-3-small"),
max_tokens=int(os.getenv("RAPTOR_MAX_TOKENS", "2048")),
temperature=float(os.getenv("RAPTOR_TEMPERATURE", "0.1")),
chunk_size=int(os.getenv("RAPTOR_CHUNK_SIZE", "1000")),
chunk_overlap=int(os.getenv("RAPTOR_CHUNK_OVERLAP", "200")),
tree_depth=int(os.getenv("RAPTOR_TREE_DEPTH", "3")),
summarization_length=int(os.getenv("RAPTOR_SUMMARY_LENGTH", "200")),
base_url=base_url,
)
def get_graphrag_config() -> GraphRAGConfig:
"""Get GraphRAG configuration from environment."""
api_key, base_url, llm_model, summ_model = _resolve_llm(
os.getenv("OPENAI_API_KEY", ""),
os.getenv("GRAPHRAG_MODEL", "gpt-5.6-luna"),
os.getenv("GRAPHRAG_SUMMARY_MODEL", "gpt-5.6-luna"),
)
return GraphRAGConfig(
llm_api_key=api_key,
llm_model=llm_model,
embedding_model=os.getenv("GRAPHRAG_EMBEDDING_MODEL", "text-embedding-3-small"),
chunk_size=int(os.getenv("GRAPHRAG_CHUNK_SIZE", "1200")),
chunk_overlap=int(os.getenv("GRAPHRAG_CHUNK_OVERLAP", "100")),
max_knowledge_triples=int(os.getenv("GRAPHRAG_MAX_TRIPLES", "10")),
community_detection_algorithm=os.getenv("GRAPHRAG_COMMUNITY_ALG", "leiden"),
summarization_model=summ_model,
base_url=base_url,
)
def get_api_config() -> APIConfig:
"""Get API configuration from environment."""
return APIConfig(
host=os.getenv("API_HOST", "127.0.0.1"),
port=int(os.getenv("API_PORT", "4242")),
reload=os.getenv("API_RELOAD", "true").lower() == "true",
max_results=int(os.getenv("API_MAX_RESULTS", "10")),
timeout_seconds=int(os.getenv("API_TIMEOUT", "30"))
)
document_processor.py¶
"""
Document processor for handling various file formats.
Specializes in processing technical documentation like Intel manuals.
"""
import re
from pathlib import Path
from typing import List, Optional, Dict, Any
import pypdf
import pdfplumber
from bs4 import BeautifulSoup
import markdown
from loguru import logger
import asyncio
import aiofiles
class DocumentProcessor:
"""Process various document formats into text for indexing."""
def __init__(self):
self.supported_formats = {
'.pdf': self.process_pdf,
'.txt': self.process_text,
'.md': self.process_markdown,
'.html': self.process_html
}
logger.info("Initialized document processor")
async def process_file(self, file_path: Path) -> str:
"""Process a file based on its extension."""
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
ext = file_path.suffix.lower()
if ext not in self.supported_formats:
raise ValueError(f"Unsupported file format: {ext}")
processor = self.supported_formats[ext]
# Run processor (some are async, some are sync)
if asyncio.iscoroutinefunction(processor):
return await processor(file_path)
else:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, processor, file_path)
def process_pdf(self, file_path: Path) -> str:
"""
Process PDF files with special handling for technical documentation.
Optimized for Intel manuals with complex formatting.
"""
logger.info(f"Processing PDF: {file_path}")
try:
# Try pdfplumber first for better table extraction
return self._process_pdf_with_pdfplumber(file_path)
except Exception as e:
logger.warning(f"pdfplumber failed, falling back to pypdf: {e}")
return self._process_pdf_with_pypdf(file_path)
def _process_pdf_with_pdfplumber(self, file_path: Path) -> str:
"""Process PDF using pdfplumber for better structure preservation."""
text_content = []
with pdfplumber.open(file_path) as pdf:
total_pages = len(pdf.pages)
logger.info(f"Processing {total_pages} pages...")
for i, page in enumerate(pdf.pages):
if i % 100 == 0:
logger.info(f"Processing page {i}/{total_pages}")
# Extract text
page_text = page.extract_text()
if page_text:
# Clean up the text
page_text = self._clean_pdf_text(page_text)
text_content.append(page_text)
# Extract tables if present
tables = page.extract_tables()
for table in tables:
if table:
# Convert table to structured text
table_text = self._format_table(table)
if table_text:
text_content.append(table_text)
return "\n\n".join(text_content)
def _process_pdf_with_pypdf(self, file_path: Path) -> str:
"""Fallback PDF processing using pypdf."""
text_content = []
with open(file_path, 'rb') as file:
reader = pypdf.PdfReader(file)
total_pages = len(reader.pages)
logger.info(f"Processing {total_pages} pages with pypdf...")
for i, page in enumerate(reader.pages):
if i % 100 == 0:
logger.info(f"Processing page {i}/{total_pages}")
text = page.extract_text()
if text:
text = self._clean_pdf_text(text)
text_content.append(text)
return "\n\n".join(text_content)
def _clean_pdf_text(self, text: str) -> str:
"""Clean extracted PDF text."""
# Remove excessive whitespace
text = re.sub(r'\s+', ' ', text)
# Fix common PDF extraction issues
text = re.sub(r'(\w)-\s+(\w)', r'\1\2', text) # Fix hyphenated words
text = re.sub(r'\s*\n\s*', '\n', text) # Clean up newlines
# Remove page numbers and headers (common in Intel manuals)
text = re.sub(r'^[\d\s]*Intel.*?Manual.*?\n', '', text, flags=re.MULTILINE)
text = re.sub(r'^\d+-\d+\s*$', '', text, flags=re.MULTILINE)
# Extract instruction definitions (Intel manual specific)
text = self._extract_intel_instructions(text)
return text.strip()
def _extract_intel_instructions(self, text: str) -> str:
"""Extract and format Intel x86/x64 instructions."""
# Pattern for Intel instruction format
instruction_pattern = r'([A-Z]{2,}[A-Z0-9]*)\s*[-—]\s*([^\n]+)'
# Find all instruction definitions
matches = re.finditer(instruction_pattern, text)
formatted_parts = []
last_end = 0
for match in matches:
# Add text before the match
formatted_parts.append(text[last_end:match.start()])
# Format the instruction
instruction = match.group(1)
description = match.group(2)
formatted_parts.append(f"\n**{instruction}**: {description}")
last_end = match.end()
# Add remaining text
formatted_parts.append(text[last_end:])
return ''.join(formatted_parts)
def _format_table(self, table: List[List]) -> str:
"""Format a table into structured text."""
if not table or not table[0]:
return ""
formatted = []
# Assume first row is header
headers = table[0]
formatted.append("Table: " + " | ".join(str(h) for h in headers if h))
# Format data rows
for row in table[1:]:
if row and any(cell for cell in row):
formatted.append(" " + " | ".join(str(cell) if cell else "-" for cell in row))
return "\n".join(formatted)
async def process_text(self, file_path: Path) -> str:
"""Process plain text files."""
logger.info(f"Processing text file: {file_path}")
async with aiofiles.open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
content = await f.read()
return content
def process_markdown(self, file_path: Path) -> str:
"""Process Markdown files."""
logger.info(f"Processing Markdown file: {file_path}")
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
# Convert Markdown to plain text
html = markdown.markdown(content)
soup = BeautifulSoup(html, 'html.parser')
text = soup.get_text()
return text
def process_html(self, file_path: Path) -> str:
"""Process HTML files."""
logger.info(f"Processing HTML file: {file_path}")
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
content = f.read()
soup = BeautifulSoup(content, 'html.parser')
# Remove script and style elements
for element in soup(['script', 'style']):
element.decompose()
# Get text
text = soup.get_text()
# Clean up whitespace
lines = (line.strip() for line in text.splitlines())
chunks = (phrase.strip() for line in lines for phrase in line.split(" "))
text = '\n'.join(chunk for chunk in chunks if chunk)
return text
def extract_sections(self, text: str, section_pattern: Optional[str] = None) -> Dict[str, str]:
"""
Extract sections from text based on patterns.
Useful for structured documents like Intel manuals.
"""
if section_pattern is None:
# Default pattern for sections like "Chapter 1", "Section 2.3", etc.
section_pattern = r'^(Chapter|Section|Part|\d+\.)\s+[\d\w\.]+.*$'
sections = {}
current_section = "Introduction"
current_content = []
for line in text.split('\n'):
if re.match(section_pattern, line, re.IGNORECASE):
# Save previous section
if current_content:
sections[current_section] = '\n'.join(current_content)
# Start new section
current_section = line.strip()
current_content = []
else:
current_content.append(line)
# Save last section
if current_content:
sections[current_section] = '\n'.join(current_content)
return sections
def extract_code_blocks(self, text: str) -> List[str]:
"""Extract code blocks or instruction examples from text."""
code_blocks = []
# Pattern for code blocks (various formats)
patterns = [
r'```[\s\S]*?```', # Markdown code blocks
r'<code>[\s\S]*?</code>', # HTML code blocks
r'^\s{4,}.*$', # Indented code blocks
r'^\t+.*$', # Tab-indented blocks
]
for pattern in patterns:
matches = re.finditer(pattern, text, re.MULTILINE)
for match in matches:
code_blocks.append(match.group(0))
return code_blocks
def extract_intel_opcodes(self, text: str) -> List[Dict[str, str]]:
"""
Extract Intel instruction opcodes and their descriptions.
Specific to Intel architecture manuals.
"""
opcodes = []
# Pattern for Intel opcode format
opcode_pattern = r'([0-9A-F]{2}(?:\s+[0-9A-F]{2})*)\s+(/[0-7]|/r)?\s+([A-Z]+[A-Z0-9]*)\s+([^\n]+)'
matches = re.finditer(opcode_pattern, text)
for match in matches:
opcodes.append({
'opcode': match.group(1),
'mod': match.group(2) or '',
'instruction': match.group(3),
'description': match.group(4).strip()
})
return opcodes
download_sample.py¶
"""
Script to download sample technical documentation for testing.
Since the full Intel manual is very large, this creates a sample document.
"""
import requests
from pathlib import Path
import json
def create_sample_intel_doc():
"""Create a sample Intel architecture documentation for testing."""
sample_doc = """
Intel® 64 and IA-32 Architectures Software Developer's Manual
Volume 1: Basic Architecture
CHAPTER 3: BASIC EXECUTION ENVIRONMENT
3.1 MODES OF OPERATION
The IA-32 architecture supports three basic operating modes: protected mode, real-address mode, and system management mode. The operating mode determines which instructions and architectural features are accessible.
Protected mode — This mode is the native state of the processor. Among the capabilities of protected mode is the ability to directly execute "real-address mode" 8086 software in a protected, multi-tasking environment. This feature is called virtual-8086 mode.
Real-address mode — This mode implements the programming environment of the Intel 8086 processor with extensions. The processor is placed in real-address mode following power-up or a reset.
System management mode (SMM) — This mode provides an operating system or executive with a transparent mechanism for implementing platform-specific functions such as power management and system security.
3.2 OVERVIEW OF THE BASIC EXECUTION ENVIRONMENT
3.2.1 64-Bit Mode Execution Environment
When in 64-bit mode, the following architectural features become available:
• 64-bit linear addressing
• Physical address extensions to 52 bits
• 16 general-purpose registers (GPRs) in 64-bit mode
• 64-bit-wide GPRs
• 64-bit instruction pointer (RIP)
• New operating mode (64-bit mode)
• Uniform byte-register addressing
• Additional SSE registers
• Fast interrupt-prioritization mechanism
3.3 MEMORY ORGANIZATION
3.3.1 IA-32 Memory Models
When employing the processor's memory management facilities, programs do not directly address physical memory. Instead, they access memory using one of three memory models: flat, segmented, or real-address mode.
Flat memory model — Memory appears to a program as a single, continuous address space. This space is called a linear address space. Code, data, and stacks are all contained in this address space. Linear address space is byte addressable.
Segmented memory model — Memory appears to a program as a group of independent address spaces called segments. Code, data, and stacks are typically contained in separate segments.
Real-address mode memory model — This is the memory model for the Intel 8086 processor. It supports a nominally 64-KByte register-based memory model.
3.4 GENERAL-PURPOSE REGISTERS
The 64-bit extensions expand the general-purpose registers to 64 bits and add 8 new registers (R8-R15). All 16 general-purpose registers can be accessed at the byte, word, dword, and qword level.
3.4.1 General-Purpose Registers in 64-Bit Mode
In 64-bit mode, there are 16 general-purpose registers and the default operand size is 32 bits. However, general-purpose registers can be accessed as 64-bit, 32-bit, 16-bit, or 8-bit values.
Register set includes:
• RAX, RBX, RCX, RDX - Extended versions of EAX, EBX, ECX, EDX
• RBP, RSI, RDI, RSP - Extended versions of EBP, ESI, EDI, ESP
• R8-R15 - New registers introduced with 64-bit extensions
3.4.2 Register Operand-Size Encoding
In 64-bit mode, the default operand size for most instructions is 32 bits. A REX prefix specifies a 64-bit operand size. Operand sizes of 8 bits and 16 bits are also available.
CHAPTER 4: INSTRUCTION SET REFERENCE
4.1 INSTRUCTION FORMAT
All Intel 64 and IA-32 instruction encodings are subsets of the general instruction format shown below. Instructions consist of optional instruction prefixes, primary opcode bytes, an addressing-form specifier (if required), a displacement (if required), and an immediate data field (if required).
4.2 DATA MOVEMENT INSTRUCTIONS
MOV—Move
Copies the second operand (source operand) to the first operand (destination operand). The source operand can be an immediate value, general-purpose register, segment register, or memory location.
Operation:
DEST ← SRC;
Flags Affected:
None
Protected Mode Exceptions:
#GP(0) If the destination operand is in a non-writable segment
#GP(0) If a memory operand effective address is outside the CS, DS, ES, FS, or GS segment limit
#SS(0) If a memory operand effective address is outside the SS segment limit
#PF(fault-code) If a page fault occurs
#AC(0) If alignment checking is enabled
MOVSX/MOVSXD—Move with Sign-Extension
Copies the contents of the source operand to the destination operand and sign extends the value. The size of the converted value depends on the operand-size attribute.
MOVZX—Move with Zero-Extend
Copies the contents of the source operand to the destination operand and zero extends the value.
XCHG—Exchange Register/Memory with Register
Exchanges the contents of the destination (first) and source (second) operands. The operands can be two general-purpose registers or a register and a memory location.
4.3 ARITHMETIC INSTRUCTIONS
ADD—Add
Adds the destination operand (first operand) and the source operand (second operand) and then stores the result in the destination operand.
Operation:
DEST ← DEST + SRC;
Flags Affected:
The OF, SF, ZF, AF, CF, and PF flags are set according to the result.
SUB—Subtract
Subtracts the second operand (source operand) from the first operand (destination operand) and stores the result in the destination operand.
MUL—Unsigned Multiply
Performs an unsigned multiplication of the first operand (destination operand) and the second operand (source operand) and stores the result in the destination operand.
IMUL—Signed Multiply
Performs a signed multiplication and stores the result in the destination.
DIV—Unsigned Divide
Divides unsigned the value in the AX, DX:AX, EDX:EAX, or RDX:RAX registers by the source operand and stores the result in the AX (AH:AL), DX:AX, EDX:EAX, or RDX:RAX registers.
4.4 LOGICAL INSTRUCTIONS
AND—Logical AND
Performs a bitwise AND operation on the destination operand and the source operand and stores the result in the destination operand location.
OR—Logical Inclusive OR
Performs a bitwise inclusive OR operation between the destination operand and the source operand and stores the result in the destination operand location.
XOR—Logical Exclusive OR
Performs a bitwise exclusive OR operation on the destination operand and the source operand and stores the result in the destination operand.
NOT—One's Complement Negation
Performs a bitwise NOT operation on the destination operand and stores the result in the destination operand location.
4.5 CONTROL TRANSFER INSTRUCTIONS
JMP—Jump
Transfers program control to a different point in the code segment. The destination operand specifies the address of the target instruction.
Jcc—Jump if Condition Is Met
Checks the state of one or more status flags in the EFLAGS register and, if the flags are in the specified state (condition), performs a jump to the target instruction specified by the destination operand.
CALL—Call Procedure
Saves procedure linking information on the stack and branches to the called procedure specified using the target operand.
RET—Return from Procedure
Transfers program control to a return address located on the top of the stack.
4.6 STRING INSTRUCTIONS
MOVS/MOVSB/MOVSW/MOVSD/MOVSQ—Move String
Moves the byte, word, doubleword, or quadword specified with the second operand to the location specified with the first operand.
CMPS/CMPSB/CMPSW/CMPSD/CMPSQ—Compare String Operands
Compares the byte, word, doubleword, or quadword specified with the first source operand with the second source operand and sets the status flags in the EFLAGS register according to the results.
CHAPTER 5: SIMD INSTRUCTIONS
5.1 SSE INSTRUCTIONS
SSE instructions operate on packed single-precision floating-point values contained in XMM registers or memory.
MOVAPS—Move Aligned Packed Single-Precision Floating-Point Values
Moves 128 bits of packed single-precision floating-point values from the source operand to the destination operand.
MOVUPS—Move Unaligned Packed Single-Precision Floating-Point Values
Moves 128 bits of packed single-precision floating-point values from the source operand to the destination operand.
ADDPS—Add Packed Single-Precision Floating-Point Values
Performs addition of the packed single-precision floating-point values from the source operand and the destination operand, and stores the packed single-precision floating-point results in the destination operand.
SUBPS—Subtract Packed Single-Precision Floating-Point Values
Performs subtraction of the packed single-precision floating-point values in the source operand from the packed single-precision floating-point values in the destination operand.
MULPS—Multiply Packed Single-Precision Floating-Point Values
Performs multiplication of the packed single-precision floating-point values from the source operand and the destination operand.
5.2 AVX INSTRUCTIONS
AVX instructions extend SSE functionality with 256-bit YMM registers.
VMOVAPS—Move Aligned Packed Single-Precision Floating-Point Values
Moves 256 bits of packed single-precision floating-point values from the source operand to the destination operand.
VADDPS—Add Packed Single-Precision Floating-Point Values
Performs SIMD addition of the packed single-precision floating-point values from the first source operand and second source operand, and stores the packed single-precision floating-point results in the destination operand.
CHAPTER 6: SYSTEM PROGRAMMING
6.1 SYSTEM REGISTERS
Control Registers
Control registers (CR0, CR2, CR3, and CR4) control the operation of the processor and the characteristics of the currently executing task.
CR0—Contains system control flags that control operating mode and states of the processor
CR1—Reserved
CR2—Contains the page-fault linear address
CR3—Contains the physical address of the base of the paging-structure hierarchy
CR4—Contains a group of flags that enable several architectural extensions
6.2 SYSTEM INSTRUCTIONS
CPUID—CPU Identification
Returns processor identification and feature information in the EAX, EBX, ECX, and EDX registers. The instruction's output depends on the contents of the EAX register upon execution.
RDTSC—Read Time-Stamp Counter
Reads the current value of the processor's time-stamp counter (a 64-bit MSR) into the EDX:EAX registers.
RDMSR—Read from Model Specific Register
Reads the contents of a 64-bit model specific register (MSR) specified in the ECX register into registers EDX:EAX.
WRMSR—Write to Model Specific Register
Writes the contents of registers EDX:EAX into the 64-bit model specific register (MSR) specified in the ECX register.
"""
# Save as a text file
output_path = Path("sample_intel_manual.txt")
with open(output_path, 'w', encoding='utf-8') as f:
f.write(sample_doc)
print(f"Created sample Intel documentation: {output_path}")
print(f"File size: {len(sample_doc)} characters")
return str(output_path)
def create_sample_queries():
"""Create sample queries for testing."""
queries = [
# Basic instruction queries
"What is the MOV instruction and how does it work?",
"Explain the difference between MOVSX and MOVZX",
"What are the arithmetic instructions in x86?",
# Register queries
"What are the general-purpose registers in 64-bit mode?",
"How many general-purpose registers are available in x86-64?",
"What is the purpose of control registers CR0-CR4?",
# Memory model queries
"What are the different memory models in IA-32 architecture?",
"Explain the flat memory model",
"What is the difference between segmented and flat memory models?",
# SIMD queries
"What are SSE instructions?",
"What is the difference between SSE and AVX?",
"How do MOVAPS and MOVUPS differ?",
# System programming queries
"What does the CPUID instruction do?",
"How do I read the time-stamp counter?",
"What are model specific registers (MSRs)?",
# Complex queries
"How do string instructions work in x86?",
"What are the different operating modes in IA-32?",
"Explain the instruction format in Intel 64 architecture",
# Relationship queries (good for GraphRAG)
"What is the relationship between MOV and XCHG instructions?",
"How are ADD and SUB instructions related?",
"Which instructions affect the FLAGS register?",
]
# Save queries
queries_path = Path("sample_queries.json")
with open(queries_path, 'w', encoding='utf-8') as f:
json.dump(queries, f, indent=2)
print(f"Created {len(queries)} sample queries: {queries_path}")
return queries
if __name__ == "__main__":
print("Creating sample Intel architecture documentation...")
doc_path = create_sample_intel_doc()
print("\nCreating sample queries...")
queries = create_sample_queries()
print("\n" + "="*60)
print("Sample data created successfully!")
print("="*60)
print("\nTo test the system:")
print("1. Build indexes:")
print(f" python main.py build {doc_path} --type both")
print("\n2. Start API server:")
print(" python main.py serve")
print("\n3. Test with queries:")
print(" python main.py query \"What is the MOV instruction?\"")
print("\n4. Run comprehensive test:")
print(" python test_indexing.py")
graphrag_indexer.py¶
"""
GraphRAG (Graph-based Retrieval Augmented Generation) implementation.
This creates a knowledge graph with entities, relationships, and community detection.
"""
import os
import json
import pickle
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple, Set
from dataclasses import dataclass, asdict
import numpy as np
from tqdm import tqdm
import networkx as nx
from openai import OpenAI
from sentence_transformers import SentenceTransformer
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity
from loguru import logger
import re
from collections import defaultdict
from config import GraphRAGConfig
@dataclass
class Entity:
"""Represents an entity in the knowledge graph."""
id: str
name: str
type: str
description: str
embedding: Optional[np.ndarray]
attributes: Dict[str, Any]
@dataclass
class Relationship:
"""Represents a relationship between entities."""
id: str
source: str # Entity ID
target: str # Entity ID
type: str
description: str
weight: float = 1.0
@dataclass
class Community:
"""Represents a community of related entities."""
id: str
entity_ids: List[str]
summary: str
embedding: Optional[np.ndarray]
level: int
class GraphRAGIndexer:
"""GraphRAG knowledge graph indexer with entity extraction and community detection."""
def __init__(self, config: GraphRAGConfig):
self.config = config
self.client = OpenAI(api_key=config.llm_api_key, base_url=config.base_url)
self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
# Knowledge graph components
self.entities: Dict[str, Entity] = {}
self.relationships: List[Relationship] = []
self.communities: Dict[str, Community] = {}
self.graph = nx.Graph()
# Ensure directories exist
self.config.index_dir.mkdir(parents=True, exist_ok=True)
self.config.cache_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Initialized GraphRAG indexer with model: {config.llm_model}")
def chunk_text(self, text: str) -> List[str]:
"""Split text into chunks with overlap."""
# Split by sentences first for better context preservation
sentences = re.split(r'(?<=[.!?])\s+', text)
chunks = []
current_chunk = []
current_size = 0
for sentence in sentences:
words = sentence.split()
if current_size + len(words) > self.config.chunk_size:
if current_chunk:
chunks.append(" ".join(current_chunk))
# Start new chunk with overlap
overlap_size = min(self.config.chunk_overlap, len(current_chunk))
current_chunk = current_chunk[-overlap_size:] if overlap_size > 0 else []
current_size = sum(len(s.split()) for s in current_chunk)
current_chunk.append(sentence)
current_size += len(words)
if current_chunk:
chunks.append(" ".join(current_chunk))
logger.info(f"Created {len(chunks)} text chunks")
return chunks
def extract_entities_relationships(self, text: str) -> Tuple[List[Dict], List[Dict]]:
"""Extract entities and relationships from text using LLM."""
prompt = f"""
Extract entities and relationships from the following technical text about Intel x86/x64 architecture.
Focus on instructions, registers, CPU features, and architectural concepts.
For entities, identify:
- Intel instructions (type: "instruction")
- Registers (type: "register")
- CPU features (type: "feature")
- Architectural components (type: "component")
- Data types (type: "datatype")
For relationships, identify how entities are connected (e.g., "uses", "modifies", "depends_on", "part_of").
Text: {text[:2000]} # Limit text length for API
Return the result as JSON with the following structure:
{{
"entities": [
{{"name": "entity_name", "type": "entity_type", "description": "brief description"}}
],
"relationships": [
{{"source": "entity1", "target": "entity2", "type": "relationship_type", "description": "brief description"}}
]
}}
Return only valid JSON, no additional text.
"""
try:
response = self.client.chat.completions.create(
model=self.config.llm_model,
messages=[
{"role": "system", "content": "You are an expert at analyzing technical documentation and extracting structured knowledge."},
{"role": "user", "content": prompt}
],
max_tokens=1000,
temperature=0.1
)
result = response.choices[0].message.content.strip()
# Extract JSON from response
json_match = re.search(r'\{[\s\S]*\}', result)
if json_match:
data = json.loads(json_match.group())
return data.get("entities", []), data.get("relationships", [])
else:
logger.warning("Could not parse JSON from LLM response")
return [], []
except Exception as e:
logger.error(f"Error extracting entities: {e}")
return [], []
def build_knowledge_graph(self, text: str):
"""Build knowledge graph from text."""
logger.info("Building knowledge graph...")
# Chunk the text
chunks = self.chunk_text(text)
# Extract entities and relationships from each chunk
all_entities = {}
all_relationships = []
for i, chunk in enumerate(tqdm(chunks, desc="Extracting entities")):
entities, relationships = self.extract_entities_relationships(chunk)
# Process entities
for entity_data in entities:
entity_name = entity_data.get("name", "").lower()
if entity_name and entity_name not in all_entities:
# Create embedding for entity description
desc = entity_data.get("description", entity_name)
embedding = self.embedding_model.encode([desc])[0]
entity = Entity(
id=f"entity_{len(all_entities)}",
name=entity_name,
type=entity_data.get("type", "unknown"),
description=desc,
embedding=embedding,
attributes={"chunk_id": i}
)
all_entities[entity_name] = entity
self.entities[entity.id] = entity
# Process relationships
for rel_data in relationships:
source_name = rel_data.get("source", "").lower()
target_name = rel_data.get("target", "").lower()
if source_name in all_entities and target_name in all_entities:
relationship = Relationship(
id=f"rel_{len(all_relationships)}",
source=all_entities[source_name].id,
target=all_entities[target_name].id,
type=rel_data.get("type", "related"),
description=rel_data.get("description", ""),
weight=1.0
)
all_relationships.append(relationship)
self.relationships.append(relationship)
# Build NetworkX graph
logger.info("Building NetworkX graph...")
for entity_id, entity in self.entities.items():
self.graph.add_node(entity_id, **asdict(entity))
for rel in self.relationships:
self.graph.add_edge(rel.source, rel.target,
type=rel.type,
description=rel.description,
weight=rel.weight)
logger.info(f"Built graph with {len(self.entities)} entities and {len(self.relationships)} relationships")
def detect_communities(self):
"""Detect communities in the knowledge graph."""
logger.info("Detecting communities...")
if len(self.graph.nodes) == 0:
logger.warning("Graph is empty, cannot detect communities")
return
# Use different community detection algorithms
if self.config.community_detection_algorithm == "leiden":
try:
import leidenalg
import igraph as ig
# Convert NetworkX to igraph
ig_graph = ig.Graph.from_networkx(self.graph)
partitions = leidenalg.find_partition(ig_graph, leidenalg.ModularityVertexPartition)
communities = {}
for i, community in enumerate(partitions):
communities[i] = [list(self.graph.nodes())[idx] for idx in community]
except ImportError:
logger.warning("Leiden algorithm not available, falling back to Louvain")
communities = nx.community.louvain_communities(self.graph, seed=42)
communities = {i: list(comm) for i, comm in enumerate(communities)}
else:
# Use Louvain algorithm
communities = nx.community.louvain_communities(self.graph, seed=42)
communities = {i: list(comm) for i, comm in enumerate(communities)}
# Create community summaries
for comm_id, entity_ids in communities.items():
if not entity_ids:
continue
# Get entities in community
community_entities = [self.entities[eid] for eid in entity_ids if eid in self.entities]
# Create community summary
entity_descriptions = [e.description for e in community_entities[:10]] # Limit for API
summary_prompt = f"""
Summarize the following group of related entities from Intel x86/x64 documentation:
Entities:
{chr(10).join(entity_descriptions)}
Provide a concise summary (max 150 words) describing what these entities have in common and their role in the architecture.
"""
try:
response = self.client.chat.completions.create(
model=self.config.summarization_model,
messages=[
{"role": "system", "content": "You are an expert at summarizing technical documentation."},
{"role": "user", "content": summary_prompt}
],
max_tokens=200,
temperature=0.1
)
summary = response.choices[0].message.content.strip()
except Exception as e:
logger.error(f"Error creating community summary: {e}")
summary = f"Community containing {len(entity_ids)} related entities"
# Create embedding for community
embedding = self.embedding_model.encode([summary])[0]
community = Community(
id=f"community_{comm_id}",
entity_ids=entity_ids,
summary=summary,
embedding=embedding,
level=0
)
self.communities[community.id] = community
logger.info(f"Detected {len(self.communities)} communities")
def hierarchical_summarization(self):
"""Create hierarchical summaries of communities."""
if len(self.communities) <= 1:
return
logger.info("Creating hierarchical community summaries...")
# Group communities by similarity
community_embeddings = np.array([c.embedding for c in self.communities.values()])
similarity_matrix = cosine_similarity(community_embeddings)
# Simple hierarchical clustering
threshold = 0.7
merged_communities = []
processed = set()
for i, comm_id in enumerate(self.communities.keys()):
if comm_id in processed:
continue
# Find similar communities
similar = []
for j, other_id in enumerate(self.communities.keys()):
if i != j and similarity_matrix[i][j] > threshold:
similar.append(other_id)
processed.add(other_id)
if similar:
# Merge communities
merged_ids = [comm_id] + similar
all_entities = []
for mid in merged_ids:
all_entities.extend(self.communities[mid].entity_ids)
# Create merged summary
summaries = [self.communities[mid].summary for mid in merged_ids]
merge_prompt = f"""
Summarize these related community summaries into a higher-level summary:
{chr(10).join(summaries)}
Provide a concise summary (max 200 words) of the overarching theme.
"""
try:
response = self.client.chat.completions.create(
model=self.config.summarization_model,
messages=[
{"role": "system", "content": "You are an expert at creating hierarchical summaries."},
{"role": "user", "content": merge_prompt}
],
max_tokens=250,
temperature=0.1
)
merged_summary = response.choices[0].message.content.strip()
except Exception as e:
logger.error(f"Error creating merged summary: {e}")
merged_summary = f"Higher-level community containing {len(all_entities)} entities"
# Create new community
merged_embedding = self.embedding_model.encode([merged_summary])[0]
merged_community = Community(
id=f"merged_community_{len(merged_communities)}",
entity_ids=all_entities,
summary=merged_summary,
embedding=merged_embedding,
level=1
)
self.communities[merged_community.id] = merged_community
merged_communities.append(merged_community)
logger.info(f"Created {len(merged_communities)} hierarchical communities")
def search(self, query: str, top_k: int = 5, search_type: str = "hybrid") -> List[Dict[str, Any]]:
"""
Search the knowledge graph.
Args:
query: Search query
top_k: Number of results to return
search_type: "entity", "community", or "hybrid"
"""
query_embedding = self.embedding_model.encode([query])[0]
results = []
if search_type in ["entity", "hybrid"]:
# Search entities
entity_scores = []
for entity_id, entity in self.entities.items():
if entity.embedding is not None:
score = cosine_similarity([query_embedding], [entity.embedding])[0][0]
entity_scores.append((entity_id, score))
entity_scores.sort(key=lambda x: x[1], reverse=True)
for entity_id, score in entity_scores[:top_k]:
entity = self.entities[entity_id]
# Get related entities
neighbors = list(self.graph.neighbors(entity_id)) if entity_id in self.graph else []
results.append({
"type": "entity",
"id": entity_id,
"name": entity.name,
"entity_type": entity.type,
"description": entity.description,
"score": float(score),
"related_entities": neighbors[:5]
})
if search_type in ["community", "hybrid"]:
# Search communities
community_scores = []
for comm_id, community in self.communities.items():
if community.embedding is not None:
score = cosine_similarity([query_embedding], [community.embedding])[0][0]
community_scores.append((comm_id, score))
community_scores.sort(key=lambda x: x[1], reverse=True)
for comm_id, score in community_scores[:top_k]:
community = self.communities[comm_id]
# Get sample entities from community
sample_entities = []
for entity_id in community.entity_ids[:5]:
if entity_id in self.entities:
entity = self.entities[entity_id]
sample_entities.append({
"name": entity.name,
"type": entity.type
})
results.append({
"type": "community",
"id": comm_id,
"summary": community.summary,
"level": community.level,
"score": float(score),
"entity_count": len(community.entity_ids),
"sample_entities": sample_entities
})
# Sort all results by score
results.sort(key=lambda x: x["score"], reverse=True)
return results[:top_k]
def multi_hop_search(self, start_entity: str, max_hops: int = 2,
relation_filter: Optional[str] = None,
top_k: int = 10) -> List[Dict[str, Any]]:
"""
多跳关系检索：沿知识图谱的关系边遍历，回答「A 通过什么与 B 相连」这类
扁平向量检索无法表达的关系性问题（对应书中「多跳关系推理」）。
与 search() 的区别：search() 只按嵌入相似度召回孤立的实体/社区，
而本方法真正利用图结构，返回从起始实体出发的**关系路径**。
Args:
start_entity: 起始实体名（不区分大小写，按子串匹配）。
max_hops: 最大跳数。
relation_filter: 若指定，只保留终点边为该关系类型的路径。
top_k: 返回的路径数上限。
Returns:
每条路径形如 {"target", "target_type", "hops", "path"}，
path 是若干 {"source", "relation", "target"} 步骤。
"""
# 按名字子串匹配定位起始节点
start_id = None
needle = start_entity.lower()
for entity_id, entity in self.entities.items():
if needle in entity.name.lower():
start_id = entity_id
break
if start_id is None or start_id not in self.graph:
logger.warning(f"multi_hop_search: 未找到起始实体 '{start_entity}'")
return []
# BFS 沿边遍历，收集 <= max_hops 跳的路径
results: List[Dict[str, Any]] = []
queue = [(start_id, [])]
while queue and len(results) < top_k * 4:
node_id, path = queue.pop(0)
if len(path) >= max_hops:
continue
for neighbor in self.graph.neighbors(node_id):
rel_type = self.graph[node_id][neighbor].get("type", "related")
src_name = self.entities[node_id].name if node_id in self.entities else node_id
dst_name = self.entities[neighbor].name if neighbor in self.entities else neighbor
step = {"source": src_name, "relation": rel_type, "target": dst_name}
new_path = path + [step]
if relation_filter is None or rel_type == relation_filter:
results.append({
"target": dst_name,
"target_type": self.entities[neighbor].type if neighbor in self.entities else "unknown",
"hops": len(new_path),
"path": new_path,
})
queue.append((neighbor, new_path))
results.sort(key=lambda r: r["hops"])
return results[:top_k]
def save_index(self, path: Optional[Path] = None):
"""Save the knowledge graph index to disk."""
save_path = path or self.config.index_dir / "graphrag_index.pkl"
# Convert to serializable format
index_data = {
'entities': {eid: asdict(e) for eid, e in self.entities.items()},
'relationships': [asdict(r) for r in self.relationships],
'communities': {cid: asdict(c) for cid, c in self.communities.items()},
'graph': nx.node_link_data(self.graph),
'config': asdict(self.config)
}
# Convert numpy arrays to lists
for entity in index_data['entities'].values():
if entity['embedding'] is not None:
entity['embedding'] = entity['embedding'].tolist()
for community in index_data['communities'].values():
if community['embedding'] is not None:
community['embedding'] = community['embedding'].tolist()
with open(save_path, 'wb') as f:
pickle.dump(index_data, f)
logger.info(f"Saved GraphRAG index to {save_path}")
def load_index(self, path: Optional[Path] = None):
"""Load knowledge graph index from disk."""
load_path = path or self.config.index_dir / "graphrag_index.pkl"
with open(load_path, 'rb') as f:
index_data = pickle.load(f)
# Reconstruct entities
self.entities = {}
for eid, entity_dict in index_data['entities'].items():
if entity_dict['embedding'] is not None:
entity_dict['embedding'] = np.array(entity_dict['embedding'])
self.entities[eid] = Entity(**entity_dict)
# Reconstruct relationships
self.relationships = [Relationship(**r) for r in index_data['relationships']]
# Reconstruct communities
self.communities = {}
for cid, comm_dict in index_data['communities'].items():
if comm_dict['embedding'] is not None:
comm_dict['embedding'] = np.array(comm_dict['embedding'])
self.communities[cid] = Community(**comm_dict)
# Reconstruct graph
self.graph = nx.node_link_graph(index_data['graph'])
logger.info(f"Loaded GraphRAG index from {load_path}")
def get_graph_statistics(self) -> Dict[str, Any]:
"""Get statistics about the knowledge graph."""
entity_types = defaultdict(int)
for entity in self.entities.values():
entity_types[entity.type] += 1
rel_types = defaultdict(int)
for rel in self.relationships:
rel_types[rel.type] += 1
return {
"total_entities": len(self.entities),
"total_relationships": len(self.relationships),
"total_communities": len(self.communities),
"entity_types": dict(entity_types),
"relationship_types": dict(rel_types),
"graph_density": nx.density(self.graph) if len(self.graph) > 0 else 0,
"average_degree": sum(dict(self.graph.degree()).values()) / max(1, len(self.graph.nodes))
}
main.py¶
"""
结构化索引工具的主入口：构建 / 查询 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-8。",
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()
raptor_indexer.py¶
"""
RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) implementation.
This creates a hierarchical tree structure with recursive summarization.
"""
import os
import json
import pickle
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, asdict
import numpy as np
from tqdm import tqdm
import tiktoken
from sklearn.mixture import GaussianMixture
from sklearn.metrics.pairwise import cosine_similarity
import umap
from openai import OpenAI
from sentence_transformers import SentenceTransformer
from loguru import logger
from config import RaptorConfig
@dataclass
class TreeNode:
"""Represents a node in the RAPTOR tree."""
id: str
level: int
text: str
summary: str
embedding: Optional[np.ndarray]
children: List[str] # IDs of child nodes
parent: Optional[str] # ID of parent node
class RaptorIndexer:
"""RAPTOR tree-based document indexer with recursive summarization."""
def __init__(self, config: RaptorConfig):
self.config = config
self.client = OpenAI(api_key=config.openai_api_key, base_url=config.base_url)
self.embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
try:
# OpenRouter-style ids (e.g. "openai/gpt-5.6-luna") aren't known to
# tiktoken; fall back to a general-purpose encoding for token counts.
self.tokenizer = tiktoken.encoding_for_model(config.model_name)
except KeyError:
self.tokenizer = tiktoken.get_encoding("cl100k_base")
# Tree structure
self.nodes: Dict[str, TreeNode] = {}
self.root_nodes: List[str] = []
# Ensure index directory exists
self.config.index_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Initialized RAPTOR indexer with model: {config.model_name}")
def chunk_text(self, text: str) -> List[str]:
"""Split text into chunks with overlap."""
words = text.split()
chunks = []
for i in range(0, len(words), self.config.chunk_size - self.config.chunk_overlap):
chunk = " ".join(words[i:i + self.config.chunk_size])
if chunk:
chunks.append(chunk)
logger.info(f"Created {len(chunks)} text chunks")
return chunks
def create_embeddings(self, texts: List[str]) -> np.ndarray:
"""Create embeddings for texts using sentence transformers."""
embeddings = self.embedding_model.encode(texts, show_progress_bar=True)
return np.array(embeddings)
def summarize_text(self, text: str, max_length: int = 200) -> str:
"""Summarize text using OpenAI API."""
try:
response = self.client.chat.completions.create(
model=self.config.model_name,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates concise summaries focusing on key technical information."},
{"role": "user", "content": f"Summarize the following text in {max_length} words or less, focusing on the main technical concepts and important details:\n\n{text}"}
],
max_tokens=max_length * 2,
temperature=self.config.temperature
)
return response.choices[0].message.content.strip()
except Exception as e:
logger.error(f"Error summarizing text: {e}")
# Return truncated text as fallback
words = text.split()[:max_length]
return " ".join(words) + "..."
def cluster_nodes(self, embeddings: np.ndarray, min_clusters: int = 2, max_clusters: int = 10) -> np.ndarray:
"""Cluster embeddings using Gaussian Mixture Model."""
n_samples = len(embeddings)
n_clusters = min(max(min_clusters, n_samples // 5), min(max_clusters, n_samples))
if n_samples < 2:
return np.zeros(n_samples)
# Use UMAP for dimensionality reduction if needed
if embeddings.shape[1] > 50:
reducer = umap.UMAP(n_components=50, n_neighbors=min(15, n_samples-1))
embeddings_reduced = reducer.fit_transform(embeddings)
else:
embeddings_reduced = embeddings
# Perform clustering
gmm = GaussianMixture(n_components=n_clusters, random_state=42)
cluster_labels = gmm.fit_predict(embeddings_reduced)
return cluster_labels
def build_tree_level(self, node_ids: List[str]) -> List[str]:
"""Build one level of the tree by clustering and summarizing nodes."""
if len(node_ids) <= 1:
return node_ids
# Get embeddings for nodes
texts = [self.nodes[nid].text for nid in node_ids]
embeddings = np.array([self.nodes[nid].embedding for nid in node_ids])
# Cluster nodes
cluster_labels = self.cluster_nodes(embeddings)
# Group nodes by cluster
clusters: Dict[int, List[str]] = {}
for i, label in enumerate(cluster_labels):
if label not in clusters:
clusters[label] = []
clusters[label].append(node_ids[i])
# Create parent nodes for each cluster
parent_ids = []
current_level = self.nodes[node_ids[0]].level + 1
for cluster_id, child_ids in clusters.items():
# Combine texts from child nodes
combined_text = "\n\n".join([self.nodes[cid].text for cid in child_ids])
# Create summary for parent node
summary = self.summarize_text(combined_text, self.config.summarization_length)
# Create embedding for summary
summary_embedding = self.embedding_model.encode([summary])[0]
# Create parent node
parent_id = f"level{current_level}_cluster{cluster_id}"
parent_node = TreeNode(
id=parent_id,
level=current_level,
text=summary,
summary=summary,
embedding=summary_embedding,
children=child_ids,
parent=None
)
# Update child nodes to reference parent
for child_id in child_ids:
self.nodes[child_id].parent = parent_id
self.nodes[parent_id] = parent_node
parent_ids.append(parent_id)
logger.info(f"Created {len(parent_ids)} parent nodes at level {current_level}")
return parent_ids
def build_index(self, text: str):
"""Build RAPTOR tree index from text."""
logger.info("Building RAPTOR tree index...")
# Chunk the text
chunks = self.chunk_text(text)
# Create leaf nodes from chunks
logger.info("Creating leaf nodes...")
leaf_ids = []
for i, chunk in enumerate(tqdm(chunks, desc="Processing chunks")):
# Create embedding
embedding = self.embedding_model.encode([chunk])[0]
# Create summary for chunk
summary = self.summarize_text(chunk, max_length=100)
# Create leaf node
node_id = f"leaf_{i}"
node = TreeNode(
id=node_id,
level=0,
text=chunk,
summary=summary,
embedding=embedding,
children=[],
parent=None
)
self.nodes[node_id] = node
leaf_ids.append(node_id)
# Build tree levels
current_level_ids = leaf_ids
for level in range(self.config.tree_depth):
if len(current_level_ids) <= 1:
break
logger.info(f"Building tree level {level + 1}...")
current_level_ids = self.build_tree_level(current_level_ids)
self.root_nodes = current_level_ids
logger.info(f"RAPTOR tree built with {len(self.nodes)} nodes and {len(self.root_nodes)} root nodes")
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
"""Search the RAPTOR tree for relevant information."""
# Create query embedding
query_embedding = self.embedding_model.encode([query])[0]
# Calculate similarity with all nodes
similarities = []
for node_id, node in self.nodes.items():
if node.embedding is not None:
sim = cosine_similarity([query_embedding], [node.embedding])[0][0]
similarities.append((node_id, sim))
# Sort by similarity
similarities.sort(key=lambda x: x[1], reverse=True)
# Get top-k results with different levels for diversity
results = []
levels_seen = set()
for node_id, score in similarities:
node = self.nodes[node_id]
# Add diversity by including nodes from different levels
if len(results) < top_k:
results.append({
"node_id": node_id,
"level": node.level,
"text": node.text,
"summary": node.summary,
"score": float(score)
})
levels_seen.add(node.level)
elif node.level not in levels_seen and len(results) < top_k * 2:
# Include some diverse results from other levels
results.append({
"node_id": node_id,
"level": node.level,
"text": node.text,
"summary": node.summary,
"score": float(score)
})
levels_seen.add(node.level)
return results[:top_k]
def save_index(self, path: Optional[Path] = None):
"""Save the RAPTOR tree index to disk."""
save_path = path or self.config.index_dir / "raptor_index.pkl"
# Convert nodes to serializable format
serializable_nodes = {}
for node_id, node in self.nodes.items():
node_dict = asdict(node)
# Convert numpy array to list for JSON serialization
if node.embedding is not None:
node_dict['embedding'] = node.embedding.tolist()
serializable_nodes[node_id] = node_dict
index_data = {
'nodes': serializable_nodes,
'root_nodes': self.root_nodes,
'config': asdict(self.config)
}
with open(save_path, 'wb') as f:
pickle.dump(index_data, f)
logger.info(f"Saved RAPTOR index to {save_path}")
def load_index(self, path: Optional[Path] = None):
"""Load RAPTOR tree index from disk."""
load_path = path or self.config.index_dir / "raptor_index.pkl"
with open(load_path, 'rb') as f:
index_data = pickle.load(f)
# Reconstruct nodes
self.nodes = {}
for node_id, node_dict in index_data['nodes'].items():
# Convert list back to numpy array
if node_dict['embedding'] is not None:
node_dict['embedding'] = np.array(node_dict['embedding'])
self.nodes[node_id] = TreeNode(**node_dict)
self.root_nodes = index_data['root_nodes']
logger.info(f"Loaded RAPTOR index from {load_path}")
def get_tree_statistics(self) -> Dict[str, Any]:
"""Get statistics about the RAPTOR tree."""
level_counts = {}
for node in self.nodes.values():
if node.level not in level_counts:
level_counts[node.level] = 0
level_counts[node.level] += 1
return {
"total_nodes": len(self.nodes),
"root_nodes": len(self.root_nodes),
"levels": len(level_counts),
"nodes_per_level": level_counts,
"average_children": sum(len(n.children) for n in self.nodes.values()) / max(1, len(self.nodes))
}
structured_vs_flat_demo.py¶
"""
结构化索引 vs 扁平检索：离线对比演示。
本模块不依赖 OpenAI / 向量模型 / 网络，纯 Python + networkx 即可运行。
它用一个手工整理的「Intel x86 SIMD 指令集」小知识库，直观对比两条检索路线：
* 扁平检索（Flat）：把每个知识点当成互相独立的文本块，按词面相似度打分召回。
这是传统 RAG「文档分块 + 向量检索」的抽象——只能返回零散片段。
* 结构化检索（Structured）：
- GraphRAG 式的实体-关系图：沿关系边做多跳遍历，能回答扁平检索答不了的
「A 通过什么和 B 相连」这类关系性问题（对应书中「多跳关系推理」）。
- RAPTOR 式的层次树：把细节聚合成上层摘要，能回答「概述某主题」这类
需要跨片段综合的宏观问题（对应书中「多层次导航」）。
这段演示对应实验 3-8（structured-index）中「知识表达哲学的对比研究」。
构建真实索引需要调用 LLM（见 main.py build），本演示则把索引结果预先手工写好，
让读者无需 API Key 也能看到「结构化索引到底解决了扁平检索的什么问题」。
"""
import json
import re
from collections import deque
from typing import Dict, List, Optional, Tuple
import networkx as nx
# ---------------------------------------------------------------------------
# 手工整理的小知识库（对应 test_indexing.py 中的 Intel x86 示例文档）
# 每个实体的 description 同时充当「扁平检索的一个文本块」。
# ---------------------------------------------------------------------------
ENTITIES: Dict[str, Dict[str, str]] = {
"ADDPS": {"type": "instruction",
"desc": "ADDPS：对打包的单精度浮点数执行并行加法，一次处理四路单精度浮点运算。"},
"MOVAPS": {"type": "instruction",
"desc": "MOVAPS：在向量寄存器与对齐内存之间搬运 128 位打包单精度浮点数据。"},
"VADDPS": {"type": "instruction",
"desc": "VADDPS：AVX 版本的打包单精度浮点加法，一次处理八路单精度浮点运算。"},
"CPUID": {"type": "instruction",
"desc": "CPUID：返回处理器标识与特性信息，用于探测处理器是否支持 SSE、AVX 等扩展。"},
"SSE": {"type": "extension",
"desc": "SSE（Streaming SIMD Extensions）：引入 128 位向量寄存器，支持打包单精度浮点并行运算。"},
"AVX": {"type": "extension",
"desc": "AVX（Advanced Vector Extensions）：把向量寄存器扩展到 256 位，进一步增强 SIMD 能力。"},
"XMM": {"type": "register",
"desc": "XMM0-XMM15：128 位向量寄存器，供 SSE 指令存放打包数据。"},
"YMM": {"type": "register",
"desc": "YMM0-YMM15：256 位向量寄存器，供 AVX 指令使用，低 128 位与 XMM 共享。"},
"CR4.OSFXSR": {"type": "control-bit",
"desc": "CR4.OSFXSR：操作系统支持 FXSAVE/FXRSTOR 的控制位，置 1 后才允许使用 SSE 指令。"},
"CR0.EM": {"type": "control-bit",
"desc": "CR0.EM：仿真标志位，为 1 时禁用 SIMD，必须清零才能执行 SSE / AVX 指令。"},
}
# 实体-关系三元组（主语, 关系, 宾语），构成 GraphRAG 的知识之网。
TRIPLES: List[Tuple[str, str, str]] = [
("ADDPS", "属于", "SSE"),
("MOVAPS", "属于", "SSE"),
("VADDPS", "属于", "AVX"),
("ADDPS", "操作", "XMM"),
("VADDPS", "操作", "YMM"),
("SSE", "使用寄存器", "XMM"),
("AVX", "使用寄存器", "YMM"),
("AVX", "扩展自", "SSE"),
("SSE", "需要启用", "CR4.OSFXSR"),
("AVX", "需要启用", "CR4.OSFXSR"),
("SSE", "要求清零", "CR0.EM"),
("CPUID", "探测", "SSE"),
("CPUID", "探测", "AVX"),
]
# RAPTOR 式层次树：把细粒度叶子聚合为上层摘要（父节点）。
TREE_SUMMARY = {
"id": "SIMD 指令集综述",
"summary": ("x86 的 SIMD 指令集自 MMX 起步，SSE 引入 128 位 XMM 向量寄存器并支持打包"
"单精度浮点运算，AVX 进一步把寄存器扩展到 256 位 YMM，逐代提升单指令多数据"
"的并行宽度；使用前需通过 CR0/CR4 控制位使能，并可用 CPUID 探测支持情况。"),
"children": ["ADDPS", "MOVAPS", "VADDPS", "SSE", "AVX", "XMM", "YMM"],
}
# ---------------------------------------------------------------------------
# 扁平检索：把每个实体描述当作独立文本块，按词面相似度（词频余弦）召回。
# 这是「向量检索」在离线场景下的一个确定性替身：无内在结构、只看片段本身。
# ---------------------------------------------------------------------------
def _tokenize(text: str) -> List[str]:
"""粗粒度分词：ASCII 词（如 ADDPS、CR4、XMM）整体保留，中文按单字切。"""
tokens = re.findall(r"[a-zA-Z0-9]+", text.lower())
tokens += re.findall(r"[一-鿿]", text)
return tokens
def _cosine(a: Dict[str, int], b: Dict[str, int]) -> float:
common = set(a) & set(b)
dot = sum(a[t] * b[t] for t in common)
na = sum(v * v for v in a.values()) ** 0.5
nb = sum(v * v for v in b.values()) ** 0.5
return dot / (na * nb) if na and nb else 0.0
class FlatRetriever:
"""按词面相似度召回独立文本块（模拟扁平向量检索）。"""
def __init__(self, entities: Dict[str, Dict[str, str]]):
self.docs = {name: e["desc"] for name, e in entities.items()}
self.types = {name: e["type"] for name, e in entities.items()}
self._vecs = {name: self._tf(text) for name, text in self.docs.items()}
@staticmethod
def _tf(text: str) -> Dict[str, int]:
vec: Dict[str, int] = {}
for tok in _tokenize(text):
vec[tok] = vec.get(tok, 0) + 1
return vec
def search(self, query: str, top_k: int = 3) -> List[Dict]:
qvec = self._tf(query)
scored = [
{"name": name, "type": self.types[name],
"desc": self.docs[name], "score": _cosine(qvec, self._vecs[name])}
for name in self.docs
]
scored.sort(key=lambda r: r["score"], reverse=True)
return scored[:top_k]
# ---------------------------------------------------------------------------
# 结构化检索：基于实体-关系图的多跳遍历（GraphRAG 的核心能力）。
# ---------------------------------------------------------------------------
def build_graph(triples: List[Tuple[str, str, str]]) -> nx.DiGraph:
g = nx.DiGraph()
for name, meta in ENTITIES.items():
g.add_node(name, **meta)
for src, rel, dst in triples:
g.add_edge(src, dst, rel=rel)
return g
def multi_hop_paths(graph: nx.DiGraph, start: str, max_hops: int = 3) -> List[List[Tuple[str, str, str]]]:
"""从 start 出发沿关系边做 BFS，返回所有 <= max_hops 跳的关系路径。
每条路径是若干 (源实体, 关系, 目标实体) 步骤的列表。这正是扁平检索无法表达的
「沿关系边遍历」——对应书中「知识图谱天然支持沿关系边遍历，使多跳查询高效可靠」。
"""
if start not in graph:
return []
paths: List[List[Tuple[str, str, str]]] = []
# 队列元素：(当前节点, 到达该节点的路径)
queue: deque = deque([(start, [])])
while queue:
node, path = queue.popleft()
if len(path) >= max_hops:
continue
for nbr in graph.successors(node):
step = (node, graph[node][nbr]["rel"], nbr)
new_path = path + [step]
paths.append(new_path)
queue.append((nbr, new_path))
return paths
def match_entity(graph: nx.DiGraph, query: str) -> Optional[str]:
"""在查询中找出出现的起始实体（按名字最长匹配，确定性）。"""
q = query.lower()
hits = [name for name in graph.nodes if name.lower() in q]
return max(hits, key=len) if hits else None
def format_path(path: List[Tuple[str, str, str]]) -> str:
if not path:
return ""
parts = [path[0][0]]
for src, rel, dst in path:
parts.append(f" --{rel}--> {dst}")
return "".join(parts)
# ---------------------------------------------------------------------------
# 三个演示查询：分别凸显扁平检索的三类短板。
# ---------------------------------------------------------------------------
def demo_multi_hop(flat: FlatRetriever, graph: nx.DiGraph, query: str, top_k: int) -> None:
print(f"\n【查询 1｜多跳关系推理】{query}")
print("-- 扁平检索（按词面相似度返回独立片段）--")
for i, r in enumerate(flat.search(query, top_k), 1):
print(f" {i}. [{r['type']}] {r['name']} (score={r['score']:.3f})")
print(" ✗ 只能召回词面相近的孤立片段，无法把 ADDPS 与某个控制位「连」起来——"
"缺少关系，就无法判断哪个控制位是 ADDPS 的答案。")
print("-- 结构化图检索（沿关系边多跳遍历）--")
start = match_entity(graph, query)
paths = multi_hop_paths(graph, start, max_hops=3)
# 只展示终点为控制位的路径（问题问的是「控制寄存器位」）
answers = [p for p in paths if graph.nodes[p[-1][2]]["type"] == "control-bit"]
for p in answers:
print(f" {format_path(p)}")
enable = [p for p in answers if p[-1][1] == "需要启用"]
if enable:
target = enable[0][-1][2]
print(f" ✓ 答案：{target}（从 {start} 经 {len(enable[0])} 跳可达）")
print(f" {graph.nodes[target]['desc']}")
def demo_compare(flat: FlatRetriever, graph: nx.DiGraph, query: str, top_k: int) -> None:
print(f"\n【查询 2｜跨节点综合对比】{query}")
print("-- 扁平检索 --")
for i, r in enumerate(flat.search(query, top_k), 1):
print(f" {i}. [{r['type']}] {r['name']} (score={r['score']:.3f})")
print(" ✗ SSE 与 AVX 各自的寄存器事实散落在不同片段里，扁平检索把它们分别召回，"
"却不会主动把「谁用哪种寄存器」对齐成一张对比表。")
print("-- 结构化图检索（顺着「使用寄存器」边取回两侧事实）--")
for ext in ("SSE", "AVX"):
regs = [dst for _, dst, d in graph.out_edges(ext, data=True) if d["rel"] == "使用寄存器"]
for reg in regs:
print(f" {ext} --使用寄存器--> {reg}：{graph.nodes[reg]['desc']}")
print(" ✓ 沿同一种关系边遍历两个实体，即可直接综合出「SSE=128 位 XMM，AVX=256 位 YMM」的对比。")
def demo_hierarchical(flat: FlatRetriever, query: str, top_k: int) -> None:
print(f"\n【查询 3｜多层次导航（RAPTOR 层次树）】{query}")
print("-- 扁平检索 --")
for i, r in enumerate(flat.search(query, top_k), 1):
print(f" {i}. [{r['type']}] {r['name']} (score={r['score']:.3f})")
print(" ✗ 召回的是零散的细节片段，过于细碎，答不了「概述」这种需要跨片段综合的宏观问题。")
print("-- 结构化树检索（返回上层摘要节点）--")
print(f" [父节点摘要] {TREE_SUMMARY['id']}")
print(f" {TREE_SUMMARY['summary']}")
print(f" ✓ 由宏观摘要切入，需要细节时再向下钻取到 {', '.join(TREE_SUMMARY['children'][:4])} 等叶子节点。")
def run_demo(top_k: int = 3, custom_query: Optional[str] = None,
output: Optional[str] = None) -> Dict:
"""运行离线对比演示；返回结构化结果（便于 --output 落盘）。"""
flat = FlatRetriever(ENTITIES)
graph = build_graph(TRIPLES)
print("=" * 68)
print("结构化索引 vs 扁平检索 · 离线对比演示（无需 API Key）")
print(f"知识库：Intel x86 SIMD 指令集 | 实体 {graph.number_of_nodes()} 个，"
f"关系 {graph.number_of_edges()} 条，层次树 1 棵")
print("=" * 68)
if custom_query:
# 自定义查询：同时给出扁平与图检索两种视角
print(f"\n【自定义查询】{custom_query}")
print("-- 扁平检索 --")
flat_hits = flat.search(custom_query, top_k)
for i, r in enumerate(flat_hits, 1):
print(f" {i}. [{r['type']}] {r['name']} (score={r['score']:.3f})")
print("-- 结构化图检索（从查询中识别到的实体多跳遍历）--")
start = match_entity(graph, custom_query)
if start is None:
print(" （未在查询中识别到已知实体，无法进行图遍历）")
paths = []
else:
paths = multi_hop_paths(graph, start, max_hops=3)
for p in paths:
print(f" {format_path(p)}")
result = {"query": custom_query,
"flat": [{"name": r["name"], "score": r["score"]} for r in flat_hits],
"graph_start": start,
"graph_paths": [format_path(p) for p in paths]}
else:
q1 = "运行 ADDPS 指令前，操作系统必须把哪个控制寄存器位置 1？"
q2 = "SSE 与 AVX 使用的向量寄存器有什么区别？"
q3 = "概述一下 x86 的 SIMD 指令集"
demo_multi_hop(flat, graph, q1, top_k)
demo_compare(flat, graph, q2, top_k)
demo_hierarchical(flat, q3, top_k)
start1 = match_entity(graph, q1)
result = {
"queries": [q1, q2, q3],
"multi_hop": {
"query": q1,
"start": start1,
"paths": [format_path(p) for p in multi_hop_paths(graph, start1, 3)
if graph.nodes[p[-1][2]]["type"] == "control-bit"],
},
}
print("\n" + "=" * 68)
print("结论：扁平检索擅长「找到含某信息的片段」，但一旦查询需要跨片段的关系推理或"
"多层次综合，就必须依赖结构化索引（图 / 层次树）。——对应书中实验 3-8 的核心观点。")
print("=" * 68)
if output:
with open(output, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False, indent=2)
print(f"\n结果已写入：{output}")
return result
if __name__ == "__main__":
run_demo()
test_indexing.py¶
"""
Test script for structured indexing with sample Intel x86 instruction documentation.
"""
import asyncio
from pathlib import Path
from loguru import logger
from config import get_raptor_config, get_graphrag_config
from raptor_indexer import RaptorIndexer
from graphrag_indexer import GraphRAGIndexer
from document_processor import DocumentProcessor
# Sample Intel x86/x64 instruction documentation text
SAMPLE_INTEL_DOC = """
Chapter 3: Basic Execution Environment
The Intel 64 and IA-32 architectures provide a comprehensive execution environment for running applications.
This chapter describes the basic elements of this environment including registers, memory organization, and instruction formats.
3.1 General-Purpose Registers
The general-purpose registers are used for arithmetic, logic, and memory operations. In 64-bit mode, there are 16 general-purpose registers:
- RAX, RBX, RCX, RDX: Traditional registers extended to 64 bits
- RSI, RDI, RBP, RSP: Index and pointer registers
- R8-R15: Additional registers available in 64-bit mode
Each register can be accessed as:
- 64-bit (RAX, RBX, etc.)
- 32-bit (EAX, EBX, etc.)
- 16-bit (AX, BX, etc.)
- 8-bit (AL/AH, BL/BH, etc.)
3.2 Instruction Format
Intel 64 and IA-32 instruction formats consist of:
1. Instruction prefixes (optional)
2. Primary opcode (1-3 bytes)
3. ModR/M byte (if required)
4. SIB byte (if required)
5. Displacement (if required)
6. Immediate data (if required)
MOV Instruction:
MOV - Move data between registers or between register and memory
The MOV instruction copies the source operand to the destination operand without affecting the source.
Syntax:
MOV destination, source
Examples:
MOV RAX, RBX ; Move RBX to RAX
MOV [RDI], RSI ; Move RSI to memory location pointed by RDI
MOV ECX, 42 ; Move immediate value 42 to ECX
ADD Instruction:
ADD - Add two operands
The ADD instruction adds the source operand to the destination operand and stores the result in the destination.
Syntax:
ADD destination, source
The instruction updates the following flags: OF, SF, ZF, AF, PF, CF
JMP Instruction:
JMP - Unconditional jump
The JMP instruction transfers program control to a different point in the code unconditionally.
Syntax:
JMP target
Chapter 4: SIMD Instructions
4.1 SSE Instructions
SSE (Streaming SIMD Extensions) provides 128-bit registers (XMM0-XMM15) for parallel operations on packed data.
MOVAPS - Move Aligned Packed Single-Precision Floating-Point Values
MOVAPS moves 128 bits of packed single-precision floating-point values from source to destination.
ADDPS - Add Packed Single-Precision Floating-Point Values
ADDPS performs parallel addition of four single-precision floating-point values.
4.2 AVX Instructions
AVX (Advanced Vector Extensions) extends SIMD capabilities with 256-bit registers (YMM0-YMM15).
VMOVAPS - Move Aligned Packed Single-Precision Floating-Point Values (AVX)
VMOVAPS moves 256 bits of packed single-precision floating-point values.
VADDPS - Add Packed Single-Precision Floating-Point Values (AVX)
VADDPS performs parallel addition of eight single-precision floating-point values.
Chapter 5: System Instructions
5.1 Control Registers
Control registers (CR0, CR2, CR3, CR4) control the operation mode and state of the processor:
- CR0: System control flags including protection enable and paging
- CR2: Page fault linear address
- CR3: Page directory base address
- CR4: Architecture extensions control
CPUID Instruction:
CPUID - CPU Identification
Returns processor identification and feature information in EAX, EBX, ECX, and EDX registers.
RDTSC Instruction:
RDTSC - Read Time-Stamp Counter
Reads the processor's time-stamp counter into EDX:EAX.
"""
async def test_indexing():
"""Test both RAPTOR and GraphRAG indexing with sample documentation."""
logger.info("Starting structured indexing test...")
# Test RAPTOR indexing
logger.info("\n" + "="*60)
logger.info("Testing RAPTOR Tree-Based Indexing")
logger.info("="*60)
raptor_config = get_raptor_config()
raptor = RaptorIndexer(raptor_config)
# Build index
raptor.build_index(SAMPLE_INTEL_DOC)
stats = raptor.get_tree_statistics()
logger.info(f"RAPTOR Statistics: {stats}")
# Test queries
test_queries = [
"What are the general-purpose registers?",
"How does the MOV instruction work?",
"What are SIMD instructions?",
"Explain control registers"
]
for query in test_queries:
logger.info(f"\nQuery: {query}")
results = raptor.search(query, top_k=3)
for i, result in enumerate(results, 1):
logger.info(f"{i}. Level {result['level']} (Score: {result['score']:.3f})")
logger.info(f" Summary: {result['summary'][:150]}...")
# Save index
raptor.save_index()
# Test GraphRAG indexing
logger.info("\n" + "="*60)
logger.info("Testing GraphRAG Knowledge Graph Indexing")
logger.info("="*60)
graphrag_config = get_graphrag_config()
graphrag = GraphRAGIndexer(graphrag_config)
# Build knowledge graph
graphrag.build_knowledge_graph(SAMPLE_INTEL_DOC)
graphrag.detect_communities()
graphrag.hierarchical_summarization()
stats = graphrag.get_graph_statistics()
logger.info(f"GraphRAG Statistics: {stats}")
# Test queries
for query in test_queries:
logger.info(f"\nQuery: {query}")
results = graphrag.search(query, top_k=3, search_type="hybrid")
for i, result in enumerate(results, 1):
if result['type'] == 'entity':
logger.info(f"{i}. Entity: {result['name']} ({result['entity_type']}) - Score: {result['score']:.3f}")
logger.info(f" Description: {result['description'][:150]}...")
else:
logger.info(f"{i}. Community (Level {result['level']}) - Score: {result['score']:.3f}")
logger.info(f" Summary: {result['summary'][:150]}...")
# Save index
graphrag.save_index()
logger.info("\n" + "="*60)
logger.info("Test completed successfully!")
logger.info("="*60)
if __name__ == "__main__":
# Set up logging
logger.add("test_indexing.log", rotation="10 MB")
# Run the test
asyncio.run(test_indexing())