ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,313 @@
|
||||
# Structured Index Implementation Guide
|
||||
|
||||
## Overview
|
||||
|
||||
This project implements two advanced document indexing approaches for handling large technical documentation:
|
||||
|
||||
1. **RAPTOR** (Recursive Abstractive Processing for Tree-Organized Retrieval)
|
||||
2. **GraphRAG** (Graph-based Retrieval Augmented Generation)
|
||||
|
||||
Both approaches are designed to handle complex technical documentation like the Intel® 64 and IA-32 Architectures Software Developer's Manual (5000+ pages).
|
||||
|
||||
## Architecture
|
||||
|
||||
### RAPTOR Tree-Based Index
|
||||
|
||||
RAPTOR creates a hierarchical tree structure through recursive summarization:
|
||||
|
||||
```
|
||||
Document
|
||||
↓
|
||||
[Chunks] → [Embeddings] → [Clusters]
|
||||
↓ ↓ ↓
|
||||
Level 0: Leaf nodes (original chunks with summaries)
|
||||
↓
|
||||
Level 1: Parent nodes (cluster summaries)
|
||||
↓
|
||||
Level 2: Higher-level summaries
|
||||
↓
|
||||
Root: Top-level abstraction
|
||||
```
|
||||
|
||||
**Key Features:**
|
||||
- Multi-level abstraction hierarchy
|
||||
- Gaussian Mixture Model clustering
|
||||
- Recursive summarization at each level
|
||||
- Cross-level search capability
|
||||
|
||||
### GraphRAG Knowledge Graph
|
||||
|
||||
GraphRAG builds a knowledge graph with entities and relationships:
|
||||
|
||||
```
|
||||
Document
|
||||
↓
|
||||
[Chunks] → [Entity Extraction] → [Relationship Discovery]
|
||||
↓ ↓ ↓
|
||||
Entities ←→ Relationships → Knowledge Graph
|
||||
↓
|
||||
[Community Detection]
|
||||
↓
|
||||
Community Summaries
|
||||
↓
|
||||
Hierarchical Communities
|
||||
```
|
||||
|
||||
**Key Features:**
|
||||
- LLM-based entity and relationship extraction
|
||||
- Community detection (Leiden/Louvain algorithms)
|
||||
- Hierarchical community summarization
|
||||
- Graph-based search across entities and communities
|
||||
|
||||
## Implementation Details
|
||||
|
||||
### Core Components
|
||||
|
||||
1. **Document Processor** (`document_processor.py`)
|
||||
- Handles multiple file formats (PDF, TXT, MD, HTML)
|
||||
- Optimized for technical documentation
|
||||
- Special handling for Intel manual format
|
||||
- Table extraction and formatting
|
||||
|
||||
2. **RAPTOR Indexer** (`raptor_indexer.py`)
|
||||
- Text chunking with configurable overlap
|
||||
- Embedding generation using sentence transformers
|
||||
- GMM clustering for node grouping
|
||||
- Recursive summarization using LLMs
|
||||
- Tree persistence and loading
|
||||
|
||||
3. **GraphRAG Indexer** (`graphrag_indexer.py`)
|
||||
- Entity extraction focused on technical concepts
|
||||
- Relationship discovery between entities
|
||||
- NetworkX graph construction
|
||||
- Community detection and summarization
|
||||
- Graph persistence and querying
|
||||
|
||||
4. **API Service** (`api_service.py`)
|
||||
- RESTful API using FastAPI
|
||||
- Asynchronous processing for large documents
|
||||
- Support for file uploads
|
||||
- Unified interface for both indexing approaches
|
||||
- Real-time status and statistics
|
||||
|
||||
### Processing Pipeline
|
||||
|
||||
#### Building Indexes
|
||||
|
||||
1. **Document Processing**
|
||||
```python
|
||||
processor = DocumentProcessor()
|
||||
text = await processor.process_file(Path("intel_manual.pdf"))
|
||||
```
|
||||
|
||||
2. **RAPTOR Indexing**
|
||||
```python
|
||||
raptor = RaptorIndexer(config)
|
||||
raptor.build_index(text) # Creates tree structure
|
||||
raptor.save_index() # Persists to disk
|
||||
```
|
||||
|
||||
3. **GraphRAG Indexing**
|
||||
```python
|
||||
graphrag = GraphRAGIndexer(config)
|
||||
graphrag.build_knowledge_graph(text) # Extract entities
|
||||
graphrag.detect_communities() # Find communities
|
||||
graphrag.hierarchical_summarization() # Create hierarchies
|
||||
graphrag.save_index() # Persist graph
|
||||
```
|
||||
|
||||
#### Querying
|
||||
|
||||
1. **RAPTOR Search**
|
||||
- Creates query embedding
|
||||
- Searches across all tree levels
|
||||
- Returns nodes with different abstraction levels
|
||||
- Includes level-specific summaries
|
||||
|
||||
2. **GraphRAG Search**
|
||||
- Supports entity, community, or hybrid search
|
||||
- Returns entities with relationships
|
||||
- Includes community summaries
|
||||
- Provides graph context
|
||||
|
||||
### API Endpoints
|
||||
|
||||
| Endpoint | Method | Purpose |
|
||||
|----------|--------|---------|
|
||||
| `/build` | POST | Build index from text/file |
|
||||
| `/upload` | POST | Upload and index document |
|
||||
| `/query` | POST | Query indexes |
|
||||
| `/status` | GET | Check index status |
|
||||
| `/statistics` | GET | Get index statistics |
|
||||
| `/indexes` | DELETE | Clear indexes |
|
||||
|
||||
### Integration with Agentic RAG
|
||||
|
||||
The structured indexes integrate seamlessly with the Agentic RAG system:
|
||||
|
||||
1. **Configuration** (`agentic-rag/config.py`)
|
||||
```python
|
||||
KnowledgeBaseType.RAPTOR # Tree-based backend
|
||||
KnowledgeBaseType.GRAPHRAG # Graph-based backend
|
||||
```
|
||||
|
||||
2. **Tool Integration** (`agentic-rag/tools.py`)
|
||||
- `_search_raptor()`: Queries RAPTOR API
|
||||
- `_search_graphrag()`: Queries GraphRAG API
|
||||
- Unified search interface for agents
|
||||
|
||||
3. **Agent Usage**
|
||||
```python
|
||||
config.knowledge_base.type = KnowledgeBaseType.RAPTOR
|
||||
agent = AgenticRAG(config)
|
||||
response = agent.query("What are x86 registers?")
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### Command Line Interface
|
||||
|
||||
```bash
|
||||
# Build both indexes
|
||||
python main.py build intel_manual.pdf --type both
|
||||
|
||||
# Query RAPTOR
|
||||
python main.py query "MOV instruction syntax" --type raptor
|
||||
|
||||
# Query GraphRAG
|
||||
python main.py query "CPU register relationships" --type graphrag
|
||||
|
||||
# Start API server
|
||||
python main.py serve
|
||||
```
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from config import get_raptor_config, get_graphrag_config
|
||||
from raptor_indexer import RaptorIndexer
|
||||
from graphrag_indexer import GraphRAGIndexer
|
||||
|
||||
# RAPTOR Example
|
||||
raptor_config = get_raptor_config()
|
||||
raptor = RaptorIndexer(raptor_config)
|
||||
raptor.build_index(document_text)
|
||||
results = raptor.search("SSE instructions", top_k=5)
|
||||
|
||||
# GraphRAG Example
|
||||
graphrag_config = get_graphrag_config()
|
||||
graphrag = GraphRAGIndexer(graphrag_config)
|
||||
graphrag.build_knowledge_graph(document_text)
|
||||
results = graphrag.search("instruction relationships", top_k=5)
|
||||
```
|
||||
|
||||
### HTTP API
|
||||
|
||||
```bash
|
||||
# Build index
|
||||
curl -X POST http://localhost:4242/build \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"file_path": "intel_manual.pdf", "index_type": "both"}'
|
||||
|
||||
# Query
|
||||
curl -X POST http://localhost:4242/query \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"query": "vector instructions", "index_type": "hybrid"}'
|
||||
```
|
||||
|
||||
## Performance Considerations
|
||||
|
||||
### RAPTOR
|
||||
- **Indexing Time**: O(n log n) for clustering
|
||||
- **Memory**: Stores embeddings for all nodes
|
||||
- **Query Time**: Fast similarity search
|
||||
- **Best For**: Hierarchical information, long documents
|
||||
|
||||
### GraphRAG
|
||||
- **Indexing Time**: O(n²) for relationship extraction
|
||||
- **Memory**: Graph structure can be large
|
||||
- **Query Time**: Graph traversal overhead
|
||||
- **Best For**: Complex relationships, entity-centric queries
|
||||
|
||||
## Configuration Options
|
||||
|
||||
### RAPTOR Settings
|
||||
```python
|
||||
chunk_size: 1000 # Words per chunk
|
||||
chunk_overlap: 200 # Overlap between chunks
|
||||
tree_depth: 3 # Maximum tree levels
|
||||
summarization_length: 200 # Summary word count
|
||||
```
|
||||
|
||||
### GraphRAG Settings
|
||||
```python
|
||||
chunk_size: 1200 # Words per chunk
|
||||
max_knowledge_triples: 10 # Triples per chunk
|
||||
community_detection: "leiden" # Algorithm choice
|
||||
summarization_model: "gpt-5.6-luna"
|
||||
```
|
||||
|
||||
## Extending the System
|
||||
|
||||
### Adding New Document Types
|
||||
1. Extend `DocumentProcessor` with new format handlers
|
||||
2. Add format-specific extraction logic
|
||||
3. Update supported_formats dictionary
|
||||
|
||||
### Custom Entity Extraction
|
||||
1. Modify prompt in `extract_entities_relationships()`
|
||||
2. Add domain-specific entity types
|
||||
3. Customize relationship types
|
||||
|
||||
### Alternative Clustering
|
||||
1. Replace GMM in RAPTOR with other algorithms
|
||||
2. Implement custom similarity metrics
|
||||
3. Add dimensionality reduction options
|
||||
|
||||
### Graph Algorithms
|
||||
1. Add new community detection algorithms
|
||||
2. Implement graph embedding techniques
|
||||
3. Add path-finding for relationship queries
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Out of Memory**
|
||||
- Reduce chunk_size
|
||||
- Process documents in sections
|
||||
- Use smaller embedding models
|
||||
|
||||
2. **Slow Indexing**
|
||||
- Use faster/smaller LLMs
|
||||
- Reduce tree_depth or max_triples
|
||||
- Enable caching
|
||||
|
||||
3. **Poor Search Results**
|
||||
- Adjust chunk_size and overlap
|
||||
- Fine-tune clustering parameters
|
||||
- Improve entity extraction prompts
|
||||
|
||||
4. **API Errors**
|
||||
- Check API keys in .env
|
||||
- Monitor rate limits
|
||||
- Verify index exists before querying
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
1. **Hybrid Indexing**: Combine RAPTOR and GraphRAG
|
||||
2. **Incremental Updates**: Add documents without rebuilding
|
||||
3. **Multi-modal Support**: Handle images and tables
|
||||
4. **Cross-lingual**: Support multiple languages
|
||||
5. **Active Learning**: Improve extraction with feedback
|
||||
6. **Distributed Processing**: Scale to larger documents
|
||||
7. **Query Optimization**: Cache frequent queries
|
||||
8. **Visualization**: Interactive graph/tree exploration
|
||||
|
||||
## References
|
||||
|
||||
- [RAPTOR Paper](https://arxiv.org/abs/2401.18059)
|
||||
- [GraphRAG by Microsoft](https://github.com/microsoft/graphrag)
|
||||
- [Intel SDM](https://www.intel.com/content/www/us/en/developer/articles/technical/intel-sdm.html)
|
||||
- [FastAPI Documentation](https://fastapi.tiangolo.com/)
|
||||
- [NetworkX Documentation](https://networkx.org/)
|
||||
@@ -0,0 +1,288 @@
|
||||
# Structured Indexing: RAPTOR & GraphRAG / 结构化索引：RAPTOR 与 GraphRAG
|
||||
|
||||
> Companion material for *AI Agents in Depth*, Chapter 3 — **Experiment 3-7**: hierarchical RAPTOR trees vs GraphRAG knowledge graphs, plus offline structured-vs-flat demo.
|
||||
> 配套《深入理解 AI Agent》第 3 章 **实验 3-7**：RAPTOR 层次树 vs GraphRAG 知识图谱，含离线「结构化 vs 扁平」演示。
|
||||
|
||||
← [Chapter 3 index / 返回第 3 章目录](../README.md)
|
||||
|
||||
---
|
||||
|
||||
## English
|
||||
|
||||
### Overview
|
||||
|
||||
Two advanced approaches for large technical documents (e.g. Intel® SDM-style manuals):
|
||||
|
||||
1. **RAPTOR** — hierarchical tree with recursive abstractive summarization
|
||||
2. **GraphRAG** — entities, relations, communities, multi-hop traversal
|
||||
|
||||
### Features
|
||||
|
||||
**RAPTOR:** multi-level abstraction; recursive summaries; leaf→root search; GMM clustering; UMAP.
|
||||
|
||||
**GraphRAG:** LLM entity/relation extract; community detection; community summaries; multi-strategy search; **`GraphRAGIndexer.multi_hop_search`** for “how is A connected to B” questions flat vector search cannot express.
|
||||
|
||||
**HTTP API:** build/query, uploads, async large docs, hybrid search, stats.
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
# From the repository root: use the shared Chapter 3 environment
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# Activate it before changing directories:
|
||||
# macOS/Linux:
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell: .venv\Scripts\Activate.ps1
|
||||
# Windows cmd: .venv\Scripts\activate.bat
|
||||
|
||||
# pip fallback when uv is not installed:
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/structured-index
|
||||
|
||||
# Exact legacy parity path, including optional RAPTOR/GraphRAG/Azure packages:
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
# API keys and preferences
|
||||
```
|
||||
|
||||
### CLI
|
||||
|
||||
Chinese `--help` on all subcommands: `python main.py --help`, `python main.py demo --help`, etc.
|
||||
|
||||
```
|
||||
usage: main.py [-h] {build,query,demo,serve} ...
|
||||
build Build structured indexes (needs OPENAI_API_KEY)
|
||||
query Query existing indexes (needs key + built indexes)
|
||||
demo Offline structured vs flat compare (no API key)
|
||||
serve Start HTTP API
|
||||
```
|
||||
|
||||
#### 0. Offline demo (no API key — recommended first)
|
||||
|
||||
Hand-curated small Intel x86 SIMD knowledge base; three query types: multi-hop, cross-node synthesis, multi-level navigation.
|
||||
|
||||
```bash
|
||||
python main.py demo
|
||||
python main.py demo --query "VADDPS 用到哪个寄存器"
|
||||
python main.py demo --output demo_result.json
|
||||
```
|
||||
|
||||
Example (multi-hop; flat fails, graph succeeds):
|
||||
|
||||
```
|
||||
【查询 1｜多跳关系推理】运行 ADDPS 指令前，操作系统必须把哪个控制寄存器位置 1？
|
||||
-- 扁平检索（按词面相似度返回独立片段）--
|
||||
1. [control-bit] CR4.OSFXSR (score=0.459)
|
||||
...
|
||||
✗ 只能召回词面相近的孤立片段，无法把 ADDPS 与某个控制位「连」起来。
|
||||
-- 结构化图检索（沿关系边多跳遍历）--
|
||||
ADDPS --属于--> SSE --需要启用--> CR4.OSFXSR
|
||||
✓ 答案：CR4.OSFXSR（从 ADDPS 经 2 跳可达）
|
||||
```
|
||||
|
||||
> `build` / `query` need real indexes (LLM for entities/summaries) → `OPENAI_API_KEY` (embeddings: local SentenceTransformers). `demo` uses hand-authored structure so readers see the point without keys.
|
||||
|
||||
#### 1. Build (needs OPENAI_API_KEY)
|
||||
|
||||
```bash
|
||||
python main.py build path/to/document.pdf
|
||||
python main.py build path/to/document.pdf --type raptor
|
||||
python main.py build path/to/document.pdf --type graphrag
|
||||
python main.py build path/to/document.pdf --output stats.json
|
||||
```
|
||||
|
||||
#### 2. Query
|
||||
|
||||
```bash
|
||||
python main.py query "What are the MOV instruction variants?"
|
||||
python main.py query "explain SSE instructions" --type raptor --top-k 10
|
||||
python main.py query "SSE registers" --type graphrag --multi-hop 2
|
||||
python main.py query "control registers" --output result.json
|
||||
```
|
||||
|
||||
#### 3. Serve
|
||||
|
||||
```bash
|
||||
python main.py serve
|
||||
# http://localhost:4242
|
||||
```
|
||||
|
||||
### HTTP API examples
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:4242/upload" \
|
||||
-F "file=@path/to/intel_manual.pdf" \
|
||||
-F "index_type=both"
|
||||
|
||||
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}'
|
||||
|
||||
curl -X POST "http://localhost:4242/query" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"query": "What are vector instructions?", "index_type": "hybrid", "top_k": 5}'
|
||||
|
||||
curl http://localhost:4242/status
|
||||
curl http://localhost:4242/statistics
|
||||
```
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|----------|--------|-------------|
|
||||
| `/` | GET | API info |
|
||||
| `/build` | POST | Build from text/file |
|
||||
| `/upload` | POST | Upload + build |
|
||||
| `/query` | POST | Query indexes |
|
||||
| `/status` | GET | Status |
|
||||
| `/statistics` | GET | Stats |
|
||||
| `/indexes` | DELETE | Clear |
|
||||
|
||||
### Project structure
|
||||
|
||||
```
|
||||
structured-index/
|
||||
├── config.py, raptor_indexer.py, graphrag_indexer.py
|
||||
├── document_processor.py, api_service.py
|
||||
├── structured_vs_flat_demo.py # offline demo
|
||||
├── main.py, requirements.txt, env.example
|
||||
├── indexes/{raptor,graphrag}/, cache/
|
||||
```
|
||||
|
||||
### How it works
|
||||
|
||||
**RAPTOR:** chunk → embed → leaves → GMM cluster → parent summaries → multi-level tree → multi-level search.
|
||||
|
||||
**GraphRAG:** entity extract → relations → NetworkX graph → communities → summaries → hierarchical merge → entity/community search (+ multi-hop).
|
||||
|
||||
### Advanced params (see `config.py`)
|
||||
|
||||
RAPTOR: `chunk_size`, `chunk_overlap`, `tree_depth`, `summarization_length`.
|
||||
GraphRAG: `chunk_size`, `max_knowledge_triples`, community algorithm, summarization model.
|
||||
|
||||
### Performance / troubleshooting
|
||||
|
||||
Large manuals take time; watch API rate limits and memory. Cache speeds re-queries. OOM → smaller chunks; check keys; start with smaller models for tests.
|
||||
|
||||
### Integration
|
||||
|
||||
Backend for agentic-rag style projects; see related chapter labs.
|
||||
|
||||
### References
|
||||
|
||||
- [RAPTOR](https://arxiv.org/abs/2401.18059)
|
||||
- [GraphRAG](https://github.com/microsoft/graphrag)
|
||||
- [Intel SDM](https://www.intel.com/content/www/us/en/developer/articles/technical/intel-sdm.html)
|
||||
|
||||
---
|
||||
|
||||
## 中文
|
||||
|
||||
### 概述
|
||||
|
||||
面向大型技术文档的两种结构化索引：
|
||||
|
||||
1. **RAPTOR** — 递归摘要的层次树
|
||||
2. **GraphRAG** — 实体/关系/社区与多跳遍历
|
||||
|
||||
### 功能
|
||||
|
||||
**RAPTOR：** 多层抽象、递归摘要、自叶到根检索、GMM 聚类、UMAP。
|
||||
**GraphRAG：** LLM 抽实体关系、社区发现、社区摘要、多策略检索、**多跳关系遍历**（扁平向量难以表达的「A 与 B 如何相连」）。
|
||||
**HTTP API：** 构建/查询、上传、异步大文档、混合检索、状态统计。
|
||||
|
||||
### 安装
|
||||
|
||||
```bash
|
||||
# 在仓库根目录使用统一的第 3 章环境
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# 切换目录前先激活环境：
|
||||
# macOS/Linux：
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell：.venv\Scripts\Activate.ps1
|
||||
# Windows cmd：.venv\Scripts\activate.bat
|
||||
|
||||
# 未安装 uv 时可用 pip 兜底：
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/structured-index
|
||||
|
||||
# 精确复现旧版单项目环境，含可选 RAPTOR/GraphRAG/Azure 依赖：
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
```
|
||||
|
||||
### 命令行
|
||||
|
||||
所有子命令有中文 `--help`：
|
||||
|
||||
```
|
||||
usage: main.py [-h] {build,query,demo,serve} ...
|
||||
build 从文档构建结构化索引（需要 OPENAI_API_KEY）
|
||||
query 查询已构建的索引
|
||||
demo 离线对比：结构化 vs 扁平（无需 API Key）
|
||||
serve 启动 HTTP API
|
||||
```
|
||||
|
||||
#### 0. 离线对比演示（推荐先跑）
|
||||
|
||||
```bash
|
||||
python main.py demo
|
||||
python main.py demo --query "VADDPS 用到哪个寄存器"
|
||||
python main.py demo --output demo_result.json
|
||||
```
|
||||
|
||||
示例输出见 English 节：扁平只能召回词面片段；图检索可经 `ADDPS → SSE → CR4.OSFXSR` 多跳得到答案。
|
||||
|
||||
#### 1–3. 构建 / 查询 / 服务
|
||||
|
||||
```bash
|
||||
python main.py build path/to/document.pdf
|
||||
python main.py build path/to/document.pdf --type raptor
|
||||
python main.py build path/to/document.pdf --type graphrag
|
||||
python main.py build path/to/document.pdf --output stats.json
|
||||
|
||||
python main.py query "What are the MOV instruction variants?"
|
||||
python main.py query "explain SSE instructions" --type raptor --top-k 10
|
||||
python main.py query "SSE registers" --type graphrag --multi-hop 2
|
||||
python main.py query "control registers" --output result.json
|
||||
|
||||
python main.py serve
|
||||
```
|
||||
|
||||
HTTP 示例与端点表与 English 节相同。
|
||||
|
||||
### 项目结构
|
||||
|
||||
```
|
||||
structured-index/
|
||||
├── config.py, raptor_indexer.py, graphrag_indexer.py
|
||||
├── document_processor.py, api_service.py
|
||||
├── structured_vs_flat_demo.py
|
||||
├── main.py, requirements.txt, env.example
|
||||
├── indexes/{raptor,graphrag}/, cache/
|
||||
```
|
||||
|
||||
### 工作原理
|
||||
|
||||
**RAPTOR：** 分块 → 嵌入 → 叶节点 → 聚类 → 父节点摘要 → 多层树 → 多层检索。
|
||||
**GraphRAG：** 实体 → 关系 → 图 → 社区 → 摘要 → 层次聚合 → 实体/社区检索（+ 多跳）。
|
||||
|
||||
### 性能与排错
|
||||
|
||||
大文档耗时；注意限流与内存。OOM 减小 chunk；检查 API Key。
|
||||
|
||||
### 参考
|
||||
|
||||
[RAPTOR](https://arxiv.org/abs/2401.18059) · [GraphRAG](https://github.com/microsoft/graphrag) · [Intel SDM](https://www.intel.com/content/www/us/en/developer/articles/technical/intel-sdm.html)
|
||||
|
||||
---
|
||||
|
||||
## Notes / 说明
|
||||
|
||||
### OpenRouter 通用回退 / Universal OpenRouter fallback
|
||||
|
||||
Chat LLM for RAPTOR summarization and GraphRAG entity extraction can use OpenRouter when `OPENROUTER_API_KEY` is set. **Embeddings stay local SentenceTransformers (all-MiniLM-L6-v2)** and are unaffected.
|
||||
@@ -0,0 +1,391 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,560 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Real Intel SDM RAPTOR-vs-GraphRAG campaign for Experiment 3-7.
|
||||
|
||||
This campaign deliberately does not use the hand-authored offline demo. It
|
||||
extracts a bounded, pinned set of pages from Intel's current Volume 1 PDF,
|
||||
builds hierarchical summaries and entity relationships with live Ark calls,
|
||||
answers concept/detail and relationship/multi-hop questions through both
|
||||
indexes, and has Moonshot judge the grounded answers. Every provider call is
|
||||
checkpointed and reused on restart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
import networkx as nx
|
||||
import numpy as np
|
||||
import torch
|
||||
from openai import OpenAI
|
||||
from sklearn.cluster import KMeans
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
PROJECT_DIR = Path(__file__).resolve().parent
|
||||
CHAPTER_DIR = PROJECT_DIR.parent
|
||||
sys.path.insert(0, str(CHAPTER_DIR))
|
||||
|
||||
from experiment_utils import ChatRecorder, sha256_file, write_campaign_evidence # noqa: E402
|
||||
|
||||
ARK_ENDPOINT = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
MOONSHOT_ENDPOINT = "https://api.moonshot.cn/v1"
|
||||
INTEL_URL = "https://cdrdv2.intel.com/v1/dl/getContent/671436"
|
||||
PAGES = [254, 255, 256, 257, 258, 259, 260, 323, 360, 361, 362, 363, 364, 365]
|
||||
SEED = 37
|
||||
QUERIES = [
|
||||
{
|
||||
"id": "concept_sse_environment",
|
||||
"category": "concept-detail",
|
||||
"question": "What architectural state and data model define the Intel SSE programming environment?",
|
||||
"reference": "SSE adds eight 128-bit XMM0-XMM7 registers and the 32-bit MXCSR control/status register, and operates on packed or scalar single-precision floating-point data; 64-bit mode exposes XMM8-XMM15.",
|
||||
"gold_pages": [254, 255, 256, 257],
|
||||
},
|
||||
{
|
||||
"id": "detail_xmm64",
|
||||
"category": "concept-detail",
|
||||
"question": "In 64-bit mode, which additional XMM registers become accessible and how are they encoded?",
|
||||
"reference": "XMM8 through XMM15 become accessible and are selected with REX prefixes.",
|
||||
"gold_pages": [255],
|
||||
},
|
||||
{
|
||||
"id": "detail_mxcsr",
|
||||
"category": "concept-detail",
|
||||
"question": "Which MXCSR bits form the SIMD floating-point rounding-control field?",
|
||||
"reference": "MXCSR bits 13 and 14 form the rounding-control (RC) field.",
|
||||
"gold_pages": [256],
|
||||
},
|
||||
{
|
||||
"id": "concept_avx_features",
|
||||
"category": "concept-detail",
|
||||
"question": "What broad capabilities distinguish the AVX programming model described here?",
|
||||
"reference": "AVX uses VEX-encoded instructions, extends vector processing including 256-bit YMM state, and adds flexible data fetching, manipulation, and branch-support primitives.",
|
||||
"gold_pages": [361, 362, 363],
|
||||
},
|
||||
{
|
||||
"id": "relation_avx_detection",
|
||||
"category": "relationship-multi-hop",
|
||||
"question": "What complete processor-and-operating-system checks must an application perform before using AVX?",
|
||||
"reference": "Check CPUID OSXSAVE bit 27 and AVX bit 28, execute XGETBV with ECX=0, and verify XCR0 bits 2:1 are 11b so both XMM and YMM state are enabled by the OS.",
|
||||
"gold_pages": [363, 364],
|
||||
"path_hints": ["AVX", "XCR0"],
|
||||
},
|
||||
{
|
||||
"id": "relation_cpuid_insufficient",
|
||||
"category": "relationship-multi-hop",
|
||||
"question": "Why is CPUID.AVX alone insufficient proof that AVX instructions can execute?",
|
||||
"reference": "The operating system must enable XSAVE/XGETBV and XMM/YMM state management in XCR0; otherwise AVX instructions raise #UD even when CPUID.AVX is set.",
|
||||
"gold_pages": [323, 363, 364],
|
||||
"path_hints": ["CPUID", "YMM"],
|
||||
},
|
||||
{
|
||||
"id": "relation_cr4_xcr0",
|
||||
"category": "relationship-multi-hop",
|
||||
"question": "How do CR4.OSXSAVE, XGETBV, XCR0, and AVX state availability depend on one another?",
|
||||
"reference": "CR4.OSXSAVE enables the XSAVE feature set and application use of XGETBV; XGETBV reads XCR0, whose XMM/YMM bits must be enabled for AVX state and instructions to be available.",
|
||||
"gold_pages": [323, 363, 364],
|
||||
"path_hints": ["CR4.OSXSAVE", "AVX"],
|
||||
},
|
||||
{
|
||||
"id": "relation_xcr0_ud",
|
||||
"category": "relationship-multi-hop",
|
||||
"question": "What happens when an XSAVE-enabled feature is not fully enabled in XCR0, and how does that explain AVX #UD behavior?",
|
||||
"reference": "Instructions for a feature not fully enabled in XCR0 raise invalid-opcode #UD; AVX likewise #UDs when the OS has not enabled both XMM and YMM state even if the processor advertises AVX.",
|
||||
"gold_pages": [323, 364],
|
||||
"path_hints": ["XCR0", "#UD"],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_json(text: str) -> dict[str, Any]:
|
||||
match = re.search(r"\{[\s\S]*\}", text or "")
|
||||
if not match:
|
||||
raise ValueError("provider response contained no JSON object")
|
||||
return json.loads(match.group())
|
||||
|
||||
|
||||
class CachedCalls:
|
||||
def __init__(self, client: OpenAI, provider: str, endpoint: str, checkpoint: Path):
|
||||
self.recorder = ChatRecorder(client, provider, endpoint)
|
||||
self.checkpoint = checkpoint
|
||||
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
||||
if checkpoint.exists():
|
||||
self.recorder.calls = json.loads(checkpoint.read_text(encoding="utf-8"))
|
||||
|
||||
def complete(self, purpose: str, **request: Any) -> str:
|
||||
for call in reversed(self.recorder.calls):
|
||||
choices = (call.get("response") or {}).get("choices") or []
|
||||
if call.get("purpose") == purpose and choices and choices[0].get("finish_reason") != "length":
|
||||
return choices[0]["message"]["content"] or ""
|
||||
try:
|
||||
response = self.recorder.create(purpose=purpose, **request)
|
||||
return response.choices[0].message.content or ""
|
||||
finally:
|
||||
self.checkpoint.write_text(
|
||||
json.dumps(self.recorder.calls, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
class Encoder:
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(name)
|
||||
self.model = AutoModel.from_pretrained(name).eval()
|
||||
|
||||
def encode(self, texts: Iterable[str]) -> np.ndarray:
|
||||
values = list(texts)
|
||||
batches = []
|
||||
for start in range(0, len(values), 16):
|
||||
tokens = self.tokenizer(
|
||||
values[start : start + 16], padding=True, truncation=True,
|
||||
max_length=384, return_tensors="pt",
|
||||
)
|
||||
with torch.no_grad():
|
||||
hidden = self.model(**tokens).last_hidden_state
|
||||
mask = tokens["attention_mask"].unsqueeze(-1).float()
|
||||
pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
|
||||
pooled = torch.nn.functional.normalize(pooled.float(), p=2, dim=1)
|
||||
batches.append(pooled.numpy())
|
||||
return np.concatenate(batches).astype("float32")
|
||||
|
||||
|
||||
def extract_pages(pdf: Path) -> list[dict[str, Any]]:
|
||||
pages = []
|
||||
for page in PAGES:
|
||||
proc = subprocess.run(
|
||||
["pdftotext", "-f", str(page), "-l", str(page), "-layout", str(pdf), "-"],
|
||||
text=True, capture_output=True, check=True,
|
||||
)
|
||||
text = proc.stdout.strip()
|
||||
pages.append({"page": page, "id": f"physical-page-{page}", "text": text})
|
||||
return pages
|
||||
|
||||
|
||||
def chat_json(calls: CachedCalls, purpose: str, model: str, system: str, user: str,
|
||||
max_tokens: int = 2400) -> dict[str, Any]:
|
||||
return parse_json(
|
||||
calls.complete(
|
||||
purpose,
|
||||
model=model,
|
||||
seed=SEED,
|
||||
temperature=0,
|
||||
max_tokens=max_tokens,
|
||||
response_format={"type": "json_object"},
|
||||
messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def build_raptor(pages: list[dict[str, Any]], ark: CachedCalls, model: str,
|
||||
encoder: Encoder) -> list[dict[str, Any]]:
|
||||
nodes = []
|
||||
for page in pages:
|
||||
data = chat_json(
|
||||
ark, f"raptor-leaf-summary-page-{page['page']}", model,
|
||||
"Summarize Intel technical documentation without adding facts. Return JSON {summary,key_terms}.",
|
||||
f"Source physical PDF page {page['page']}:\n{page['text'][:14000]}",
|
||||
)
|
||||
nodes.append(
|
||||
{"id": f"leaf-{page['page']}", "level": 0, "summary": data["summary"],
|
||||
"key_terms": data.get("key_terms", []), "source_pages": [page["page"]],
|
||||
"children": [], "text_preview": page["text"][:1200]}
|
||||
)
|
||||
vectors = encoder.encode(node["summary"] for node in nodes)
|
||||
labels = KMeans(n_clusters=3, random_state=SEED, n_init=10).fit_predict(vectors)
|
||||
parents = []
|
||||
for cluster in sorted(set(labels.tolist())):
|
||||
children = [node for node, label in zip(nodes, labels) if int(label) == cluster]
|
||||
data = chat_json(
|
||||
ark, f"raptor-parent-summary-cluster-{cluster}", model,
|
||||
"Create a faithful cross-page technical summary. Return JSON {summary,key_relationships}.",
|
||||
"\n\n".join(f"PAGES {c['source_pages']}: {c['summary']}" for c in children),
|
||||
)
|
||||
parents.append(
|
||||
{"id": f"parent-{cluster}", "level": 1, "summary": data["summary"],
|
||||
"key_relationships": data.get("key_relationships", []),
|
||||
"source_pages": sorted({p for c in children for p in c["source_pages"]}),
|
||||
"children": [c["id"] for c in children]}
|
||||
)
|
||||
root_data = chat_json(
|
||||
ark, "raptor-root-summary", model,
|
||||
"Create the root summary of a technical hierarchy. Return JSON {summary,major_themes}.",
|
||||
"\n\n".join(f"PAGES {p['source_pages']}: {p['summary']}" for p in parents),
|
||||
)
|
||||
root = {
|
||||
"id": "root", "level": 2, "summary": root_data["summary"],
|
||||
"major_themes": root_data.get("major_themes", []),
|
||||
"source_pages": sorted({p for parent in parents for p in parent["source_pages"]}),
|
||||
"children": [parent["id"] for parent in parents],
|
||||
}
|
||||
return nodes + parents + [root]
|
||||
|
||||
|
||||
def entity_key(name: str) -> str:
|
||||
return re.sub(r"[^a-z0-9#]+", "_", name.casefold()).strip("_")
|
||||
|
||||
|
||||
def build_graph(pages: list[dict[str, Any]], ark: CachedCalls, model: str,
|
||||
encoder: Encoder) -> tuple[dict[str, Any], list[dict[str, Any]], list[dict[str, Any]], nx.Graph]:
|
||||
entities: dict[str, Any] = {}
|
||||
relationships = []
|
||||
for page in pages:
|
||||
data = chat_json(
|
||||
ark, f"graphrag-extract-page-{page['page']}", model,
|
||||
(
|
||||
"Extract a technical knowledge graph only from the source. Return JSON with entities "
|
||||
"[{name,type,description}] and relationships "
|
||||
"[{source,target,type,description}]. Preserve exact register/feature names "
|
||||
"and explicitly connect prerequisites, state components, instructions and failure outcomes."
|
||||
" Select at most 8 high-value entities and 8 explicit relationships per page; "
|
||||
"each description must be at most 18 words; avoid aliases and repetition."
|
||||
),
|
||||
f"Physical PDF page {page['page']}:\n{page['text'][:9000]}",
|
||||
max_tokens=1400,
|
||||
)
|
||||
for raw in data.get("entities", []):
|
||||
name = str(raw.get("name", "")).strip()
|
||||
if not name:
|
||||
continue
|
||||
key = entity_key(name)
|
||||
item = entities.setdefault(
|
||||
key, {"id": key, "name": name, "type": raw.get("type", "concept"),
|
||||
"descriptions": [], "evidence": [], "source_pages": []},
|
||||
)
|
||||
if raw.get("description") and raw["description"] not in item["descriptions"]:
|
||||
item["descriptions"].append(raw["description"])
|
||||
if raw.get("evidence_quote"):
|
||||
item["evidence"].append({"page": page["page"], "quote": raw["evidence_quote"]})
|
||||
item["source_pages"] = sorted(set(item["source_pages"] + [page["page"]]))
|
||||
for raw in data.get("relationships", []):
|
||||
source_name = str(raw.get("source", "")).strip()
|
||||
target_name = str(raw.get("target", "")).strip()
|
||||
if not source_name or not target_name:
|
||||
continue
|
||||
for name in (source_name, target_name):
|
||||
key = entity_key(name)
|
||||
entities.setdefault(
|
||||
key, {"id": key, "name": name, "type": "concept", "descriptions": [],
|
||||
"evidence": [], "source_pages": [page["page"]]},
|
||||
)
|
||||
relationships.append(
|
||||
{"source": entity_key(source_name), "target": entity_key(target_name),
|
||||
"type": raw.get("type", "related_to"), "description": raw.get("description", ""),
|
||||
"evidence_quote": raw.get("evidence_quote", ""), "source_page": page["page"]}
|
||||
)
|
||||
graph = nx.Graph()
|
||||
for key, entity in entities.items():
|
||||
graph.add_node(key, **entity)
|
||||
for rel in relationships:
|
||||
graph.add_edge(rel["source"], rel["target"], **rel)
|
||||
raw_communities = list(nx.community.greedy_modularity_communities(graph)) if graph.number_of_edges() else []
|
||||
raw_communities = sorted(raw_communities, key=len, reverse=True)[:8]
|
||||
communities = []
|
||||
for idx, members in enumerate(raw_communities):
|
||||
member_list = sorted(members)
|
||||
sub_relationships = [r for r in relationships if r["source"] in members and r["target"] in members]
|
||||
data = chat_json(
|
||||
ark, f"graphrag-community-summary-{idx}", model,
|
||||
"Summarize this graph community and its technical relationships. Return JSON {summary,key_relationships}.",
|
||||
json.dumps(
|
||||
{"entities": [entities[m] for m in member_list], "relationships": sub_relationships},
|
||||
ensure_ascii=False,
|
||||
)[:16000],
|
||||
)
|
||||
communities.append(
|
||||
{"id": f"community-{idx}", "entity_ids": member_list, "summary": data["summary"],
|
||||
"key_relationships": data.get("key_relationships", []),
|
||||
"source_pages": sorted({p for m in member_list for p in entities[m]["source_pages"]})}
|
||||
)
|
||||
return entities, relationships, communities, graph
|
||||
|
||||
|
||||
def top_indices(scores: np.ndarray, k: int) -> list[int]:
|
||||
return np.argsort(-scores)[: min(k, len(scores))].tolist()
|
||||
|
||||
|
||||
def raptor_context(query: str, nodes: list[dict[str, Any]], encoder: Encoder,
|
||||
matrix: np.ndarray) -> tuple[str, list[int], list[dict[str, Any]]]:
|
||||
scores = matrix @ encoder.encode([query])[0]
|
||||
selected = [{**nodes[i], "score": float(scores[i])} for i in top_indices(scores, 5)]
|
||||
pages = sorted({page for node in selected for page in node["source_pages"]})
|
||||
text = "\n\n".join(
|
||||
f"NODE {node['id']} LEVEL {node['level']} SOURCE PAGES {node['source_pages']}: {node['summary']}"
|
||||
for node in selected
|
||||
)
|
||||
return text, pages, selected
|
||||
|
||||
|
||||
def resolve_hint(graph: nx.Graph, hint: str) -> str | None:
|
||||
needle = hint.casefold()
|
||||
for node, attrs in graph.nodes(data=True):
|
||||
if needle in attrs.get("name", "").casefold() or needle in node.casefold():
|
||||
return node
|
||||
return None
|
||||
|
||||
|
||||
def graph_context(query: str, spec: dict[str, Any], entities: dict[str, Any],
|
||||
relationships: list[dict[str, Any]], communities: list[dict[str, Any]],
|
||||
graph: nx.Graph, encoder: Encoder, entity_matrix: np.ndarray,
|
||||
community_matrix: np.ndarray) -> tuple[str, list[int], dict[str, Any]]:
|
||||
q = encoder.encode([query])[0]
|
||||
entity_ids = list(entities)
|
||||
selected_ids = [entity_ids[i] for i in top_indices(entity_matrix @ q, 7)]
|
||||
expanded = set(selected_ids)
|
||||
for node in selected_ids:
|
||||
expanded.update(graph.neighbors(node))
|
||||
selected_relationships = [
|
||||
rel for rel in relationships if rel["source"] in expanded and rel["target"] in expanded
|
||||
]
|
||||
selected_communities = []
|
||||
if communities:
|
||||
selected_communities = [communities[i] for i in top_indices(community_matrix @ q, 2)]
|
||||
pages = sorted(
|
||||
{p for node in expanded for p in entities[node]["source_pages"]}
|
||||
| {r["source_page"] for r in selected_relationships}
|
||||
| {p for c in selected_communities for p in c["source_pages"]}
|
||||
)
|
||||
paths = []
|
||||
hints = spec.get("path_hints") or []
|
||||
if len(hints) == 2:
|
||||
source, target = resolve_hint(graph, hints[0]), resolve_hint(graph, hints[1])
|
||||
if source and target and nx.has_path(graph, source, target):
|
||||
path = nx.shortest_path(graph, source, target)
|
||||
paths.append(
|
||||
{
|
||||
"hints": hints, "nodes": [entities[node]["name"] for node in path],
|
||||
"hops": len(path) - 1,
|
||||
"edges": [graph[path[i]][path[i + 1]] for i in range(len(path) - 1)],
|
||||
}
|
||||
)
|
||||
context = {
|
||||
"entities": [entities[node] for node in sorted(expanded)],
|
||||
"relationships": selected_relationships,
|
||||
"communities": selected_communities,
|
||||
"explicit_paths": paths,
|
||||
}
|
||||
return json.dumps(context, ensure_ascii=False), pages, context
|
||||
|
||||
|
||||
def answer(calls: CachedCalls, purpose: str, model: str, question: str, context: str) -> str:
|
||||
return calls.complete(
|
||||
purpose,
|
||||
model=model,
|
||||
seed=SEED,
|
||||
temperature=0,
|
||||
max_tokens=1000,
|
||||
messages=[
|
||||
{"role": "system", "content": "Answer only from the retrieved Intel manual evidence. Cite physical PDF pages in brackets. If evidence is incomplete, say so."},
|
||||
{"role": "user", "content": f"QUESTION: {question}\n\nRETRIEVED EVIDENCE:\n{context}"},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def call_totals(calls: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
||||
for call in calls:
|
||||
for key in usage:
|
||||
usage[key] += int((call.get("usage") or {}).get(key) or 0)
|
||||
return {"calls": len(calls), "usage": usage, "latency_ms": sum(float(c.get("latency_ms") or 0) for c in calls)}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description="Experiment 3-7 real Intel structured-index campaign")
|
||||
parser.add_argument("--pdf", type=Path, default=PROJECT_DIR / "data" / "intel-sdm-volume-1.pdf")
|
||||
parser.add_argument("--model", default=os.getenv("ARK_MODEL", "doubao-seed-1-6-250615"))
|
||||
parser.add_argument("--judge-model", default=os.getenv("STRUCTURED_INDEX_JUDGE", "moonshot-v1-32k"))
|
||||
parser.add_argument("--embedding-model", default="sentence-transformers/all-MiniLM-L6-v2")
|
||||
args = parser.parse_args()
|
||||
ark_key = os.getenv("ARK_API_KEY") or os.getenv("DOUBAO_API_KEY")
|
||||
moonshot_key = os.getenv("MOONSHOT_API_KEY")
|
||||
if not ark_key or not moonshot_key:
|
||||
raise RuntimeError("ARK_API_KEY and MOONSHOT_API_KEY are required")
|
||||
if not args.pdf.exists():
|
||||
args.pdf.parent.mkdir(parents=True, exist_ok=True)
|
||||
subprocess.run(["curl", "-L", "--fail", "-o", str(args.pdf), INTEL_URL], check=True)
|
||||
source_hash = sha256_file(args.pdf)
|
||||
checkpoint = PROJECT_DIR / "validation" / "checkpoints" / f"intel-{source_hash[:12]}"
|
||||
ark = CachedCalls(
|
||||
OpenAI(api_key=ark_key, base_url=ARK_ENDPOINT, timeout=180, max_retries=3),
|
||||
"volcengine-ark", ARK_ENDPOINT, checkpoint / "ark.json",
|
||||
)
|
||||
judge = CachedCalls(
|
||||
OpenAI(api_key=moonshot_key, base_url=MOONSHOT_ENDPOINT, timeout=180, max_retries=3),
|
||||
"moonshot", MOONSHOT_ENDPOINT, checkpoint / "moonshot.json",
|
||||
)
|
||||
|
||||
pages = extract_pages(args.pdf)
|
||||
cover_text = subprocess.run(
|
||||
["pdftotext", "-f", "1", "-l", "1", "-layout", str(args.pdf), "-"],
|
||||
text=True, capture_output=True, check=True,
|
||||
).stdout
|
||||
encoder_started = time.perf_counter()
|
||||
encoder = Encoder(args.embedding_model)
|
||||
embedding_identity = {
|
||||
"provider": "local Hugging Face transformers", "model": args.embedding_model,
|
||||
"class": type(encoder.model).__name__,
|
||||
"parameters": sum(p.numel() for p in encoder.model.parameters()),
|
||||
"load_latency_ms": round((time.perf_counter() - encoder_started) * 1000, 3),
|
||||
}
|
||||
raptor_start = time.perf_counter()
|
||||
raptor_nodes = build_raptor(pages, ark, args.model, encoder)
|
||||
raptor_build_ms = (time.perf_counter() - raptor_start) * 1000
|
||||
raptor_matrix = encoder.encode(node["summary"] for node in raptor_nodes)
|
||||
|
||||
graph_start = time.perf_counter()
|
||||
entities, relationships, communities, graph = build_graph(pages, ark, args.model, encoder)
|
||||
graph_build_ms = (time.perf_counter() - graph_start) * 1000
|
||||
entity_ids = list(entities)
|
||||
entity_matrix = encoder.encode(
|
||||
f"{entities[key]['name']}: {' '.join(entities[key]['descriptions'])}" for key in entity_ids
|
||||
)
|
||||
community_matrix = encoder.encode(c["summary"] for c in communities) if communities else np.empty((0, entity_matrix.shape[1]))
|
||||
|
||||
results = []
|
||||
for spec in QUERIES:
|
||||
for method in ("raptor", "graphrag"):
|
||||
started = time.perf_counter()
|
||||
if method == "raptor":
|
||||
context, retrieved_pages, trace = raptor_context(spec["question"], raptor_nodes, encoder, raptor_matrix)
|
||||
else:
|
||||
context, retrieved_pages, trace = graph_context(
|
||||
spec["question"], spec, entities, relationships, communities, graph,
|
||||
encoder, entity_matrix, community_matrix,
|
||||
)
|
||||
response = answer(
|
||||
ark, f"answer-{method}-{spec['id']}", args.model, spec["question"], context
|
||||
)
|
||||
gold = set(spec["gold_pages"])
|
||||
recall = len(gold & set(retrieved_pages)) / len(gold)
|
||||
results.append(
|
||||
{"id": f"{method}:{spec['id']}", "method": method, "query_id": spec["id"],
|
||||
"category": spec["category"], "question": spec["question"],
|
||||
"reference": spec["reference"], "gold_pages": spec["gold_pages"],
|
||||
"retrieved_pages": retrieved_pages, "citation_recall": recall,
|
||||
"retrieval_trace": trace, "answer": response,
|
||||
"query_latency_ms": round((time.perf_counter() - started) * 1000, 3)}
|
||||
)
|
||||
|
||||
judgement_payload = [
|
||||
{k: row[k] for k in ("id", "question", "reference", "answer")} for row in results
|
||||
]
|
||||
judgements = chat_json(
|
||||
judge, "external-judge-all-answers", args.judge_model,
|
||||
"Judge Intel technical answers. Return JSON {items:[{id,score,correct,reason}]}; score 0-4. Require all material conditions and no contradiction.",
|
||||
json.dumps(judgement_payload, ensure_ascii=False), max_tokens=3500,
|
||||
)["items"]
|
||||
judged = {item["id"]: item for item in judgements}
|
||||
for row in results:
|
||||
row["external_judge"] = judged[row["id"]]
|
||||
|
||||
summary: dict[str, Any] = {}
|
||||
for method in ("raptor", "graphrag"):
|
||||
summary[method] = {}
|
||||
for category in ("concept-detail", "relationship-multi-hop", "overall"):
|
||||
selected = [
|
||||
row for row in results
|
||||
if row["method"] == method and (category == "overall" or row["category"] == category)
|
||||
]
|
||||
summary[method][category] = {
|
||||
"n": len(selected),
|
||||
"mean_citation_recall": sum(r["citation_recall"] for r in selected) / len(selected),
|
||||
"mean_judge_score": sum(float(r["external_judge"]["score"]) for r in selected) / len(selected),
|
||||
"mean_query_latency_ms": sum(r["query_latency_ms"] for r in selected) / len(selected),
|
||||
}
|
||||
explicit_paths = [
|
||||
path for row in results if row["method"] == "graphrag"
|
||||
for path in (row["retrieval_trace"].get("explicit_paths") or [])
|
||||
]
|
||||
acceptance = {
|
||||
"official_intel_pdf_pinned": args.pdf.stat().st_size > 1_000_000 and "June 2026" in cover_text,
|
||||
"bounded_real_pages_extracted": len(pages) == len(PAGES) and all(page["text"] for page in pages),
|
||||
"live_hierarchical_leaf_parent_root_summaries": len(raptor_nodes) == len(PAGES) + 4,
|
||||
"live_entity_relationship_extraction": len(entities) >= 10 and len(relationships) >= 5,
|
||||
"graph_communities_summarized": bool(communities),
|
||||
"concept_detail_and_relationship_multihop_sets": {q["category"] for q in QUERIES}
|
||||
== {"concept-detail", "relationship-multi-hop"},
|
||||
"both_indexes_answered_identical_queries": len(results) == len(QUERIES) * 2,
|
||||
"actual_graph_paths_retained": bool(explicit_paths),
|
||||
"external_judge_complete": len(judgements) == len(results),
|
||||
"raw_live_receipts_checkpointed": ark.checkpoint.exists() and judge.checkpoint.exists(),
|
||||
}
|
||||
evidence = {
|
||||
"status": "passed" if all(acceptance.values()) else "failed",
|
||||
"source": {
|
||||
"publisher": "Intel Corporation", "url": INTEL_URL,
|
||||
"title": "Intel 64 and IA-32 Architectures Software Developer's Manual, Volume 1: Basic Architecture",
|
||||
"order_number_revision": "253665-092US", "publication": "June 2026",
|
||||
"pdf_path": str(args.pdf.resolve()), "pdf_sha256": source_hash,
|
||||
"pdf_bytes": args.pdf.stat().st_size, "physical_pages_selected": PAGES,
|
||||
"cover_metadata_text": cover_text, "extracted_pages": pages,
|
||||
},
|
||||
"providers": {
|
||||
"builder_answerer": {"provider": "Volcengine Ark", "endpoint": ARK_ENDPOINT, "model": args.model, "seed": SEED},
|
||||
"judge": {"provider": "Moonshot", "endpoint": MOONSHOT_ENDPOINT, "model": args.judge_model, "seed": SEED},
|
||||
"embedding": embedding_identity,
|
||||
},
|
||||
"raptor": {"build_latency_ms": round(raptor_build_ms, 3), "nodes": raptor_nodes,
|
||||
"statistics": {"levels": 3, "leaves": len(PAGES), "parents": 3, "roots": 1}},
|
||||
"graphrag": {
|
||||
"build_latency_ms": round(graph_build_ms, 3), "entities": entities,
|
||||
"relationships": relationships, "communities": communities,
|
||||
"statistics": {"entities": len(entities), "relationships": len(relationships),
|
||||
"communities": len(communities), "density": nx.density(graph)},
|
||||
},
|
||||
"queries": QUERIES, "results": results, "explicit_multi_hop_paths": explicit_paths,
|
||||
"call_totals": {"ark": call_totals(ark.recorder.calls), "moonshot": call_totals(judge.recorder.calls)},
|
||||
"summary": summary, "acceptance": acceptance,
|
||||
"checkpoint_files": [str(ark.checkpoint), str(judge.checkpoint)],
|
||||
}
|
||||
manifest = write_campaign_evidence(
|
||||
PROJECT_DIR, "3-7", evidence,
|
||||
receipts=ark.recorder.calls + judge.recorder.calls,
|
||||
input_paths=[__file__, args.pdf, PROJECT_DIR / "config.py", PROJECT_DIR / "raptor_indexer.py", PROJECT_DIR / "graphrag_indexer.py"],
|
||||
)
|
||||
print(json.dumps(summary, indent=2))
|
||||
print(json.dumps(acceptance, indent=2))
|
||||
print(f"evidence: {manifest['run_dir']}")
|
||||
return 0 if all(acceptance.values()) else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,180 @@
|
||||
"""
|
||||
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."""
|
||||
provider = os.getenv("LLM_PROVIDER", "openai").lower()
|
||||
provider = {"qwen": "dashscope", "bailian": "dashscope"}.get(provider, provider)
|
||||
openai_key = os.getenv("OPENAI_API_KEY", "")
|
||||
dashscope_key = os.getenv("DASHSCOPE_API_KEY", "")
|
||||
ark_key = os.getenv("ARK_API_KEY") or os.getenv("DOUBAO_API_KEY", "")
|
||||
direct_key = dashscope_key if provider == "dashscope" else (openai_key or ark_key)
|
||||
default_model = (
|
||||
os.getenv("RAPTOR_MODEL", "qwen3.7-plus")
|
||||
if provider == "dashscope"
|
||||
else (
|
||||
os.getenv("ARK_MODEL", "doubao-seed-1-6-250615")
|
||||
if ark_key and not openai_key
|
||||
else "gpt-5.6-luna")
|
||||
)
|
||||
api_key, base_url, model_name = _resolve_llm(
|
||||
direct_key,
|
||||
os.getenv("RAPTOR_MODEL", default_model),
|
||||
)
|
||||
if base_url is None and provider == "dashscope":
|
||||
base_url = os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
|
||||
elif base_url is None and ark_key and not openai_key:
|
||||
base_url = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
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."""
|
||||
provider = os.getenv("LLM_PROVIDER", "openai").lower()
|
||||
provider = {"qwen": "dashscope", "bailian": "dashscope"}.get(provider, provider)
|
||||
openai_key = os.getenv("OPENAI_API_KEY", "")
|
||||
dashscope_key = os.getenv("DASHSCOPE_API_KEY", "")
|
||||
ark_key = os.getenv("ARK_API_KEY") or os.getenv("DOUBAO_API_KEY", "")
|
||||
direct_key = dashscope_key if provider == "dashscope" else (openai_key or ark_key)
|
||||
default_model = (
|
||||
os.getenv("GRAPHRAG_MODEL", "qwen3.7-plus")
|
||||
if provider == "dashscope"
|
||||
else (
|
||||
os.getenv("ARK_MODEL", "doubao-seed-1-6-250615")
|
||||
if ark_key and not openai_key
|
||||
else "gpt-5.6-luna")
|
||||
)
|
||||
api_key, base_url, llm_model, summ_model = _resolve_llm(
|
||||
direct_key,
|
||||
os.getenv("GRAPHRAG_MODEL", default_model),
|
||||
os.getenv("GRAPHRAG_SUMMARY_MODEL", default_model),
|
||||
)
|
||||
if base_url is None and provider == "dashscope":
|
||||
base_url = os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
|
||||
elif base_url is None and ark_key and not openai_key:
|
||||
base_url = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
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"))
|
||||
)
|
||||
@@ -0,0 +1,291 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,297 @@
|
||||
"""
|
||||
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")
|
||||
@@ -0,0 +1,43 @@
|
||||
# LLM provider: openai (default) or dashscope/qwen/bailian
|
||||
LLM_PROVIDER=openai
|
||||
|
||||
# OpenAI API Configuration
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
|
||||
# Alibaba Cloud Model Studio / Bailian (Qwen)
|
||||
# DASHSCOPE_API_KEY=your-dashscope-api-key-here
|
||||
# DASHSCOPE_BASE_URL=https://dashscope-intl.aliyuncs.com/compatible-mode/v1
|
||||
|
||||
# OpenRouter universal fallback (optional): the chat LLM used for RAPTOR
|
||||
# summarization and GraphRAG entity extraction is routed through OpenRouter
|
||||
# when OPENAI_API_KEY is missing but OPENROUTER_API_KEY is set (gpt-5.6-luna ->
|
||||
# openai/gpt-5.6-luna; OPENROUTER_MODEL overrides). Embeddings here are local
|
||||
# SentenceTransformers, so they are unaffected.
|
||||
OPENROUTER_API_KEY=your-openrouter-api-key-here
|
||||
# OPENROUTER_MODEL=openai/gpt-5.6-luna
|
||||
|
||||
# RAPTOR Configuration
|
||||
RAPTOR_MODEL=gpt-5.6-luna
|
||||
RAPTOR_EMBEDDING_MODEL=text-embedding-3-small
|
||||
RAPTOR_MAX_TOKENS=2048
|
||||
RAPTOR_TEMPERATURE=0.1
|
||||
RAPTOR_CHUNK_SIZE=1000
|
||||
RAPTOR_CHUNK_OVERLAP=200
|
||||
RAPTOR_TREE_DEPTH=3
|
||||
RAPTOR_SUMMARY_LENGTH=200
|
||||
|
||||
# GraphRAG Configuration
|
||||
GRAPHRAG_MODEL=gpt-5.6-luna
|
||||
GRAPHRAG_EMBEDDING_MODEL=text-embedding-3-small
|
||||
GRAPHRAG_CHUNK_SIZE=1200
|
||||
GRAPHRAG_CHUNK_OVERLAP=100
|
||||
GRAPHRAG_MAX_TRIPLES=10
|
||||
GRAPHRAG_COMMUNITY_ALG=leiden
|
||||
GRAPHRAG_SUMMARY_MODEL=gpt-5.6-luna
|
||||
|
||||
# API Service Configuration
|
||||
API_HOST=127.0.0.1
|
||||
API_PORT=4242
|
||||
API_RELOAD=true
|
||||
API_MAX_RESULTS=10
|
||||
API_TIMEOUT=30
|
||||
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
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. chunk_overlap is a WORD budget
|
||||
# (the same unit as chunk_size, which current_size is measured
|
||||
# in); len(current_chunk) is a SENTENCE count, so using it here
|
||||
# kept the whole previous chunk and the window never advanced.
|
||||
overlap: List[str] = []
|
||||
overlap_size = 0
|
||||
for prev in reversed(current_chunk):
|
||||
prev_size = len(prev.split())
|
||||
if overlap_size + prev_size > self.config.chunk_overlap:
|
||||
break
|
||||
overlap.insert(0, prev)
|
||||
overlap_size += prev_size
|
||||
current_chunk = overlap
|
||||
current_size = overlap_size
|
||||
|
||||
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. Snapshot the ids up front: the loop
|
||||
# below inserts the merged communities into self.communities, and
|
||||
# iterating the live dict raised "RuntimeError: dictionary changed size
|
||||
# during iteration". The snapshot also keeps i/j aligned with
|
||||
# similarity_matrix, which is built once from these same communities.
|
||||
community_ids = list(self.communities.keys())
|
||||
community_embeddings = np.array([self.communities[cid].embedding for cid in community_ids])
|
||||
similarity_matrix = cosine_similarity(community_embeddings)
|
||||
|
||||
# Simple hierarchical clustering
|
||||
threshold = 0.7
|
||||
merged_communities = []
|
||||
processed = set()
|
||||
|
||||
for i, comm_id in enumerate(community_ids):
|
||||
if comm_id in processed:
|
||||
continue
|
||||
|
||||
# Find similar communities
|
||||
similar = []
|
||||
for j, other_id in enumerate(community_ids):
|
||||
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"
|
||||
"""
|
||||
if top_k <= 0:
|
||||
return []
|
||||
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))
|
||||
}
|
||||
@@ -0,0 +1,514 @@
|
||||
"""Hybrid Structured Retriever for RAPTOR Hierarchical Trees and GraphRAG Knowledge Graphs.
|
||||
|
||||
Merges RAPTOR tree summary nodes and GraphRAG entity-relation summaries into a unified retrieval index.
|
||||
Performs Reciprocal Rank Fusion (RRF) scoring and evidence citation tracking across hierarchical and graph chunks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvidenceCitation:
|
||||
"""Represents evidence citation metadata tracking hierarchical and graph provenance."""
|
||||
source_type: str # "raptor_tree", "graphrag_entity", "graphrag_relation", "graphrag_community"
|
||||
node_id: str
|
||||
citation_label: str
|
||||
hierarchical_level: Optional[int] = None
|
||||
entity_type: Optional[str] = None
|
||||
relation_type: Optional[str] = None
|
||||
community_level: Optional[int] = None
|
||||
lineage: List[str] = field(default_factory=list)
|
||||
snippet: str = ""
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert evidence citation to dictionary format."""
|
||||
return {
|
||||
"source_type": self.source_type,
|
||||
"node_id": self.node_id,
|
||||
"citation_label": self.citation_label,
|
||||
"hierarchical_level": self.hierarchical_level,
|
||||
"entity_type": self.entity_type,
|
||||
"relation_type": self.relation_type,
|
||||
"community_level": self.community_level,
|
||||
"lineage": self.lineage,
|
||||
"snippet": self.snippet,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Represents a unified hybrid search result item with RRF score and citation."""
|
||||
node_id: str
|
||||
text: str
|
||||
summary: str
|
||||
score: float # Fused RRF score
|
||||
source_type: str
|
||||
citation: EvidenceCitation
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert search result to dictionary format."""
|
||||
return {
|
||||
"node_id": self.node_id,
|
||||
"text": self.text,
|
||||
"summary": self.summary,
|
||||
"score": self.score,
|
||||
"source_type": self.source_type,
|
||||
"citation": self.citation.to_dict(),
|
||||
"metadata": self.metadata,
|
||||
}
|
||||
|
||||
|
||||
class HybridStructuredRetriever:
|
||||
"""Retriever that merges RAPTOR tree summaries and GraphRAG graph summaries into a hybrid index.
|
||||
|
||||
Supports Reciprocal Rank Fusion (RRF) across hierarchical (RAPTOR) and knowledge graph (GraphRAG)
|
||||
indexes with evidence citation tracking for full auditability.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rrf_k: int = 60,
|
||||
embedding_fn: Optional[Callable[[str], np.ndarray]] = None,
|
||||
) -> None:
|
||||
"""Initialize the HybridStructuredRetriever.
|
||||
|
||||
Args:
|
||||
rrf_k: Smoothing constant for Reciprocal Rank Fusion (RRF). Default 60.
|
||||
embedding_fn: Optional callable to convert text into vector embeddings.
|
||||
"""
|
||||
self.rrf_k = max(1, int(rrf_k))
|
||||
self.embedding_fn = embedding_fn
|
||||
|
||||
# Internal node stores
|
||||
self.raptor_nodes: Dict[str, Dict[str, Any]] = {}
|
||||
self.graphrag_entities: Dict[str, Dict[str, Any]] = {}
|
||||
self.graphrag_relations: Dict[str, Dict[str, Any]] = {}
|
||||
self.graphrag_communities: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
# Unified document registry
|
||||
self.unified_nodes: Dict[str, Dict[str, Any]] = {}
|
||||
|
||||
def add_raptor_node(
|
||||
self,
|
||||
node_id: str,
|
||||
level: int,
|
||||
text: str,
|
||||
summary: str = "",
|
||||
embedding: Optional[np.ndarray] = None,
|
||||
children: Optional[List[str]] = None,
|
||||
parent: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Add a RAPTOR tree summary node to the index."""
|
||||
record = {
|
||||
"id": str(node_id),
|
||||
"level": int(level),
|
||||
"text": str(text),
|
||||
"summary": str(summary or text),
|
||||
"embedding": embedding,
|
||||
"children": [str(c) for c in (children or [])],
|
||||
"parent": str(parent) if parent is not None else None,
|
||||
"source_type": "raptor_tree",
|
||||
}
|
||||
self.raptor_nodes[str(node_id)] = record
|
||||
self.unified_nodes[f"raptor_{node_id}"] = record
|
||||
|
||||
def add_graphrag_entity(
|
||||
self,
|
||||
entity_id: str,
|
||||
name: str,
|
||||
type: str = "GENERIC",
|
||||
description: str = "",
|
||||
embedding: Optional[np.ndarray] = None,
|
||||
attributes: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""Add a GraphRAG entity node to the index."""
|
||||
record = {
|
||||
"id": str(entity_id),
|
||||
"name": str(name),
|
||||
"type": str(type),
|
||||
"description": str(description),
|
||||
"embedding": embedding,
|
||||
"attributes": dict(attributes or {}),
|
||||
"source_type": "graphrag_entity",
|
||||
}
|
||||
self.graphrag_entities[str(entity_id)] = record
|
||||
self.unified_nodes[f"entity_{entity_id}"] = record
|
||||
|
||||
def add_graphrag_relationship(
|
||||
self,
|
||||
relation_id: str,
|
||||
source: str,
|
||||
target: str,
|
||||
type: str = "RELATED_TO",
|
||||
description: str = "",
|
||||
weight: float = 1.0,
|
||||
) -> None:
|
||||
"""Add a GraphRAG relationship summary node to the index."""
|
||||
record = {
|
||||
"id": str(relation_id),
|
||||
"source": str(source),
|
||||
"target": str(target),
|
||||
"type": str(type),
|
||||
"description": str(description),
|
||||
"weight": float(weight),
|
||||
"source_type": "graphrag_relation",
|
||||
}
|
||||
self.graphrag_relations[str(relation_id)] = record
|
||||
self.unified_nodes[f"rel_{relation_id}"] = record
|
||||
|
||||
def add_graphrag_community(
|
||||
self,
|
||||
community_id: str,
|
||||
entity_ids: List[str],
|
||||
summary: str,
|
||||
level: int = 0,
|
||||
embedding: Optional[np.ndarray] = None,
|
||||
) -> None:
|
||||
"""Add a GraphRAG community summary node to the index."""
|
||||
record = {
|
||||
"id": str(community_id),
|
||||
"entity_ids": [str(e) for e in (entity_ids or [])],
|
||||
"summary": str(summary),
|
||||
"level": int(level),
|
||||
"embedding": embedding,
|
||||
"source_type": "graphrag_community",
|
||||
}
|
||||
self.graphrag_communities[str(community_id)] = record
|
||||
self.unified_nodes[f"community_{community_id}"] = record
|
||||
|
||||
def index_raptor_nodes(self, nodes: Sequence[Any]) -> None:
|
||||
"""Bulk ingest RAPTOR tree nodes (objects or dicts)."""
|
||||
for node in nodes:
|
||||
if isinstance(node, dict):
|
||||
n_id = node.get("id") if node.get("id") is not None else node.get("node_id")
|
||||
level = node.get("level", 0)
|
||||
text = node.get("text", "")
|
||||
summary = node.get("summary", text)
|
||||
embedding = node.get("embedding")
|
||||
children = node.get("children", [])
|
||||
parent = node.get("parent")
|
||||
else:
|
||||
n_id = getattr(node, "id", None)
|
||||
if n_id is None:
|
||||
n_id = getattr(node, "node_id", None)
|
||||
level = getattr(node, "level", 0)
|
||||
text = getattr(node, "text", "")
|
||||
summary = getattr(node, "summary", text)
|
||||
embedding = getattr(node, "embedding", None)
|
||||
children = getattr(node, "children", [])
|
||||
parent = getattr(node, "parent", None)
|
||||
if n_id is not None:
|
||||
self.add_raptor_node(
|
||||
node_id=str(n_id),
|
||||
level=level,
|
||||
text=text,
|
||||
summary=summary,
|
||||
embedding=embedding,
|
||||
children=children,
|
||||
parent=parent,
|
||||
)
|
||||
|
||||
def index_graphrag_data(
|
||||
self,
|
||||
entities: Optional[Sequence[Any]] = None,
|
||||
relationships: Optional[Sequence[Any]] = None,
|
||||
communities: Optional[Sequence[Any]] = None,
|
||||
) -> None:
|
||||
"""Bulk ingest GraphRAG entities, relationships, and communities."""
|
||||
if entities:
|
||||
for item in entities:
|
||||
if isinstance(item, dict):
|
||||
e_id = item.get("id") if item.get("id") is not None else item.get("entity_id")
|
||||
name = item.get("name") if item.get("name") is not None else e_id
|
||||
e_type = item.get("type", "GENERIC")
|
||||
desc = item.get("description", "")
|
||||
emb = item.get("embedding")
|
||||
attrs = item.get("attributes", {})
|
||||
else:
|
||||
e_id = getattr(item, "id", None)
|
||||
if e_id is None:
|
||||
e_id = getattr(item, "entity_id", None)
|
||||
name = getattr(item, "name", None)
|
||||
if name is None:
|
||||
name = str(e_id)
|
||||
e_type = getattr(item, "type", "GENERIC")
|
||||
desc = getattr(item, "description", "")
|
||||
emb = getattr(item, "embedding", None)
|
||||
attrs = getattr(item, "attributes", {})
|
||||
if e_id is not None:
|
||||
self.add_graphrag_entity(e_id, name, e_type, desc, emb, attrs)
|
||||
|
||||
if relationships:
|
||||
for item in relationships:
|
||||
if isinstance(item, dict):
|
||||
r_id = item.get("id") if item.get("id") is not None else item.get("relation_id")
|
||||
src = item.get("source", "")
|
||||
tgt = item.get("target", "")
|
||||
r_type = item.get("type", "RELATED_TO")
|
||||
desc = item.get("description", "")
|
||||
wt = item.get("weight", 1.0)
|
||||
else:
|
||||
r_id = getattr(item, "id", None)
|
||||
if r_id is None:
|
||||
r_id = getattr(item, "relation_id", None)
|
||||
src = getattr(item, "source", "")
|
||||
tgt = getattr(item, "target", "")
|
||||
r_type = getattr(item, "type", "RELATED_TO")
|
||||
desc = getattr(item, "description", "")
|
||||
wt = getattr(item, "weight", 1.0)
|
||||
if r_id is not None:
|
||||
self.add_graphrag_relationship(r_id, src, tgt, r_type, desc, wt)
|
||||
|
||||
if communities:
|
||||
for item in communities:
|
||||
if isinstance(item, dict):
|
||||
c_id = item.get("id") if item.get("id") is not None else item.get("community_id")
|
||||
e_ids = item.get("entity_ids", [])
|
||||
summ = item.get("summary", "")
|
||||
lvl = item.get("level", 0)
|
||||
emb = item.get("embedding")
|
||||
else:
|
||||
c_id = getattr(item, "id", None)
|
||||
if c_id is None:
|
||||
c_id = getattr(item, "community_id", None)
|
||||
e_ids = getattr(item, "entity_ids", [])
|
||||
summ = getattr(item, "summary", "")
|
||||
lvl = getattr(item, "level", 0)
|
||||
emb = getattr(item, "embedding", None)
|
||||
if c_id is not None:
|
||||
self.add_graphrag_community(c_id, e_ids, summ, lvl, emb)
|
||||
|
||||
def _compute_scores(
|
||||
self, query: str, query_terms: set[str], query_vector: Optional[np.ndarray], node: Dict[str, Any]
|
||||
) -> Tuple[float, float, float]:
|
||||
"""Compute (final_score, lexical_score, semantic_score) for a node."""
|
||||
if not query_terms:
|
||||
return 0.0, 0.0, 0.0
|
||||
|
||||
# Construct textual content for evaluation
|
||||
text_content = ""
|
||||
src_type = node.get("source_type")
|
||||
if src_type == "raptor_tree":
|
||||
text_content = f"{node.get('summary', '')} {node.get('text', '')}"
|
||||
elif src_type == "graphrag_entity":
|
||||
text_content = f"{node.get('name', '')} {node.get('type', '')} {node.get('description', '')}"
|
||||
elif src_type == "graphrag_relation":
|
||||
text_content = f"{node.get('source', '')} {node.get('type', '')} {node.get('target', '')} {node.get('description', '')}"
|
||||
elif src_type == "graphrag_community":
|
||||
text_content = f"{node.get('summary', '')}"
|
||||
|
||||
if not text_content.strip():
|
||||
return 0.0, 0.0, 0.0
|
||||
|
||||
words = re.findall(r"\w+", text_content.lower())
|
||||
if not words:
|
||||
return 0.0, 0.0, 0.0
|
||||
|
||||
word_counts = defaultdict(int)
|
||||
for w in words:
|
||||
word_counts[w] += 1
|
||||
|
||||
matched_terms = [qt for qt in query_terms if qt in word_counts]
|
||||
lexical_score = len(matched_terms) / len(query_terms) if query_terms else 0.0
|
||||
matches = sum(word_counts[qt] for qt in matched_terms)
|
||||
coverage_score = min(1.0, matches / len(words)) if words else 0.0
|
||||
|
||||
semantic_score = 0.0
|
||||
has_vector = False
|
||||
if query_vector is not None:
|
||||
try:
|
||||
n_emb = node.get("embedding")
|
||||
if n_emb is None and self.embedding_fn is not None:
|
||||
n_emb = self.embedding_fn(text_content)
|
||||
node["embedding"] = n_emb
|
||||
if n_emb is not None:
|
||||
q_norm = np.linalg.norm(query_vector)
|
||||
n_norm = np.linalg.norm(n_emb)
|
||||
if q_norm > 0 and n_norm > 0:
|
||||
cos_sim = float(np.dot(query_vector, n_emb) / (q_norm * n_norm))
|
||||
semantic_score = max(0.0, cos_sim)
|
||||
has_vector = True
|
||||
except Exception:
|
||||
has_vector = False
|
||||
semantic_score = 0.0
|
||||
|
||||
if query_vector is not None:
|
||||
if has_vector:
|
||||
final_score = float(semantic_score * 0.7 + lexical_score * 0.3)
|
||||
else:
|
||||
# No vector available for this item: use lexical score at the
|
||||
# same weight as the no-query-vector path (0.8) so textually matching
|
||||
# items are not penalized below non-matching vectorized items.
|
||||
semantic_score = 0.0
|
||||
final_score = float(lexical_score * 0.8)
|
||||
else:
|
||||
semantic_score = coverage_score
|
||||
final_score = float(lexical_score * 0.8 + coverage_score * 0.2)
|
||||
return final_score, lexical_score, semantic_score
|
||||
|
||||
def _compute_relevance_score(
|
||||
self, query: str, query_terms: set[str], query_vector: Optional[np.ndarray], node: Dict[str, Any]
|
||||
) -> float:
|
||||
return self._compute_scores(query, query_terms, query_vector, node)[0]
|
||||
def _build_citation(self, node: Dict[str, Any]) -> EvidenceCitation:
|
||||
"""Construct structured evidence citation tracking provenance for a node."""
|
||||
src_type = node.get("source_type", "unknown")
|
||||
n_id = str(node.get("id", ""))
|
||||
|
||||
if src_type == "raptor_tree":
|
||||
lvl = node.get("level", 0)
|
||||
label = f"[RAPTOR Tree Level {lvl} Node: {n_id}]"
|
||||
lineage = []
|
||||
if node.get("parent") is not None:
|
||||
lineage.append(f"Parent: {node['parent']}")
|
||||
if node.get("children"):
|
||||
lineage.append(f"Children: {', '.join(str(c) for c in node['children'])}")
|
||||
snippet = node.get("summary") or node.get("text") or ""
|
||||
return EvidenceCitation(
|
||||
source_type="raptor_tree",
|
||||
node_id=n_id,
|
||||
citation_label=label,
|
||||
hierarchical_level=lvl,
|
||||
lineage=lineage,
|
||||
snippet=snippet[:200],
|
||||
)
|
||||
|
||||
elif src_type == "graphrag_entity":
|
||||
e_type = node.get("type", "GENERIC")
|
||||
e_name = node.get("name", n_id)
|
||||
label = f"[GraphRAG Entity: {e_name} (Type: {e_type})]"
|
||||
snippet = node.get("description", "")
|
||||
return EvidenceCitation(
|
||||
source_type="graphrag_entity",
|
||||
node_id=n_id,
|
||||
citation_label=label,
|
||||
entity_type=e_type,
|
||||
lineage=[f"EntityName: {e_name}"],
|
||||
snippet=snippet[:200],
|
||||
)
|
||||
|
||||
elif src_type == "graphrag_relation":
|
||||
r_type = node.get("type", "RELATED_TO")
|
||||
src = node.get("source", "")
|
||||
tgt = node.get("target", "")
|
||||
label = f"[GraphRAG Relation: {src} --({r_type})--> {tgt}]"
|
||||
snippet = node.get("description", "")
|
||||
return EvidenceCitation(
|
||||
source_type="graphrag_relation",
|
||||
node_id=n_id,
|
||||
citation_label=label,
|
||||
relation_type=r_type,
|
||||
lineage=[f"Source: {src}", f"Target: {tgt}"],
|
||||
snippet=snippet[:200],
|
||||
)
|
||||
|
||||
elif src_type == "graphrag_community":
|
||||
lvl = node.get("level", 0)
|
||||
e_ids = node.get("entity_ids", [])
|
||||
label = f"[GraphRAG Community Level {lvl}: {n_id}]"
|
||||
snippet = node.get("summary", "")
|
||||
return EvidenceCitation(
|
||||
source_type="graphrag_community",
|
||||
node_id=n_id,
|
||||
citation_label=label,
|
||||
community_level=lvl,
|
||||
lineage=[f"Entities: {', '.join(str(e) for e in e_ids[:5])}"],
|
||||
snippet=snippet[:200],
|
||||
)
|
||||
|
||||
return EvidenceCitation(
|
||||
source_type=src_type,
|
||||
node_id=n_id,
|
||||
citation_label=f"[Source: {src_type} Node: {n_id}]",
|
||||
)
|
||||
|
||||
def retrieve(self, query: str, top_k: int = 5, rrf_k: Optional[int] = None) -> List[SearchResult]:
|
||||
"""Retrieve and rank hybrid results using Reciprocal Rank Fusion (RRF).
|
||||
|
||||
Args:
|
||||
query: The search query string.
|
||||
top_k: Number of top ranked results to return.
|
||||
rrf_k: Optional override for RRF k constant.
|
||||
|
||||
Returns:
|
||||
List of SearchResult items ordered descending by fused RRF score.
|
||||
"""
|
||||
k_val = max(1, int(rrf_k)) if rrf_k is not None else self.rrf_k
|
||||
top_k = max(0, int(top_k))
|
||||
if top_k == 0:
|
||||
return []
|
||||
|
||||
if not query or not query.strip():
|
||||
return []
|
||||
|
||||
query_terms = set(re.findall(r"\w+", query.lower()))
|
||||
query_vector = None
|
||||
if self.embedding_fn is not None:
|
||||
try:
|
||||
query_vector = self.embedding_fn(query)
|
||||
except Exception:
|
||||
query_vector = None
|
||||
|
||||
candidates: Dict[str, Tuple[float, float, float]] = {}
|
||||
for key, node in self.unified_nodes.items():
|
||||
final_sc, lex_sc, sem_sc = self._compute_scores(query, query_terms, query_vector, node)
|
||||
if final_sc > 0:
|
||||
candidates[key] = (final_sc, lex_sc, sem_sc)
|
||||
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
# 1. Lexical ranking across all candidates
|
||||
lexical_sorted = sorted(candidates.keys(), key=lambda k: (candidates[k][1], candidates[k][0], k), reverse=True)
|
||||
lexical_ranks = {k: r + 1 for r, k in enumerate(lexical_sorted)}
|
||||
|
||||
# 2. Semantic/Coverage ranking across all candidates
|
||||
semantic_sorted = sorted(candidates.keys(), key=lambda k: (candidates[k][2], candidates[k][0], k), reverse=True)
|
||||
semantic_ranks = {k: r + 1 for r, k in enumerate(semantic_sorted)}
|
||||
|
||||
# 3. Reciprocal Rank Fusion (RRF) scoring across identical candidate set
|
||||
all_candidate_keys = sorted(candidates.keys())
|
||||
rrf_scores: Dict[str, float] = {}
|
||||
for key in all_candidate_keys:
|
||||
rrf_scores[key] = (1.0 / (k_val + lexical_ranks[key])) + (1.0 / (k_val + semantic_ranks[key]))
|
||||
|
||||
sorted_keys = sorted(
|
||||
all_candidate_keys,
|
||||
key=lambda k: (rrf_scores[k], candidates[k][0], k),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
# 4. Construct final SearchResult objects with citations
|
||||
results: List[SearchResult] = []
|
||||
for key in sorted_keys[:top_k]:
|
||||
node = self.unified_nodes[key]
|
||||
citation = self._build_citation(node)
|
||||
|
||||
src_type = node.get("source_type", "unknown")
|
||||
text = node.get("text") or node.get("description") or node.get("summary") or ""
|
||||
summary = node.get("summary") or node.get("description") or text
|
||||
|
||||
res = SearchResult(
|
||||
node_id=str(node.get("id")),
|
||||
text=text,
|
||||
summary=summary,
|
||||
score=rrf_scores[key],
|
||||
source_type=src_type,
|
||||
citation=citation,
|
||||
metadata={
|
||||
"lexical_rank": lexical_ranks.get(key),
|
||||
"semantic_rank": semantic_ranks.get(key),
|
||||
"raw_node": {k: v for k, v in node.items() if k != "embedding"},
|
||||
},
|
||||
)
|
||||
results.append(res)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,202 @@
|
||||
"""
|
||||
结构化索引工具的主入口：构建 / 查询 RAPTOR 与 GraphRAG 索引，或运行离线对比演示。
|
||||
|
||||
说明：RAPTOR、GraphRAG 的**索引构建**需要调用 LLM（实体抽取、递归摘要），因此
|
||||
build / query 依赖 OPENAI_API_KEY 及相应重型依赖（umap、sentence-transformers 等）。
|
||||
若只想直观理解「结构化索引解决了扁平检索的什么问题」，可运行无需 API 的 `demo` 子命令。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
import json
|
||||
import sys
|
||||
|
||||
from loguru import logger
|
||||
|
||||
|
||||
async def build_indexes(file_path: Path, index_type: str = "both",
|
||||
output: str = None):
|
||||
"""Build RAPTOR and/or GraphRAG indexes from a document."""
|
||||
# 重型依赖延迟导入：保证 --help / demo 在缺少 umap 等依赖时仍可用
|
||||
from config import get_raptor_config, get_graphrag_config
|
||||
from raptor_indexer import RaptorIndexer
|
||||
from graphrag_indexer import GraphRAGIndexer
|
||||
from document_processor import DocumentProcessor
|
||||
|
||||
logger.info(f"Building {index_type} index(es) from {file_path}")
|
||||
|
||||
# Process document
|
||||
processor = DocumentProcessor()
|
||||
text = await processor.process_file(file_path)
|
||||
logger.info(f"Processed document: {len(text)} characters")
|
||||
|
||||
all_stats = {}
|
||||
|
||||
# Build RAPTOR index
|
||||
if index_type in ["raptor", "both"]:
|
||||
logger.info("Building RAPTOR tree index...")
|
||||
raptor_config = get_raptor_config()
|
||||
raptor = RaptorIndexer(raptor_config)
|
||||
raptor.build_index(text)
|
||||
raptor.save_index()
|
||||
stats = raptor.get_tree_statistics()
|
||||
all_stats["raptor"] = stats
|
||||
logger.info(f"RAPTOR index built: {stats}")
|
||||
|
||||
# Build GraphRAG index
|
||||
if index_type in ["graphrag", "both"]:
|
||||
logger.info("Building GraphRAG knowledge graph...")
|
||||
graphrag_config = get_graphrag_config()
|
||||
graphrag = GraphRAGIndexer(graphrag_config)
|
||||
graphrag.build_knowledge_graph(text)
|
||||
graphrag.detect_communities()
|
||||
graphrag.hierarchical_summarization()
|
||||
graphrag.save_index()
|
||||
stats = graphrag.get_graph_statistics()
|
||||
all_stats["graphrag"] = stats
|
||||
logger.info(f"GraphRAG index built: {stats}")
|
||||
|
||||
if output:
|
||||
with open(output, "w", encoding="utf-8") as f:
|
||||
json.dump(all_stats, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"索引统计已写入：{output}")
|
||||
|
||||
logger.info("Indexing complete!")
|
||||
|
||||
|
||||
async def query_indexes(query: str, index_type: str = "both", top_k: int = 5,
|
||||
multi_hop: int = 0):
|
||||
"""Query RAPTOR and/or GraphRAG indexes."""
|
||||
from config import get_raptor_config, get_graphrag_config
|
||||
from raptor_indexer import RaptorIndexer
|
||||
from graphrag_indexer import GraphRAGIndexer
|
||||
|
||||
results = {}
|
||||
|
||||
# Query RAPTOR
|
||||
if index_type in ["raptor", "both"]:
|
||||
try:
|
||||
raptor_config = get_raptor_config()
|
||||
raptor = RaptorIndexer(raptor_config)
|
||||
raptor.load_index()
|
||||
raptor_results = raptor.search(query, top_k)
|
||||
results["raptor"] = raptor_results
|
||||
logger.info(f"RAPTOR returned {len(raptor_results)} results")
|
||||
except Exception as e:
|
||||
logger.error(f"Error querying RAPTOR: {e}")
|
||||
|
||||
# Query GraphRAG
|
||||
if index_type in ["graphrag", "both"]:
|
||||
try:
|
||||
graphrag_config = get_graphrag_config()
|
||||
graphrag = GraphRAGIndexer(graphrag_config)
|
||||
graphrag.load_index()
|
||||
graphrag_results = graphrag.search(query, top_k)
|
||||
results["graphrag"] = graphrag_results
|
||||
logger.info(f"GraphRAG returned {len(graphrag_results)} results")
|
||||
|
||||
# 多跳关系检索：以召回的最佳实体为起点，沿关系边遍历
|
||||
if multi_hop > 0 and graphrag_results:
|
||||
start = next((r.get("name") for r in graphrag_results
|
||||
if r.get("type") == "entity"), None)
|
||||
if start:
|
||||
paths = graphrag.multi_hop_search(start, max_hops=multi_hop)
|
||||
results["graphrag_multi_hop"] = paths
|
||||
logger.info(f"GraphRAG multi-hop from '{start}' "
|
||||
f"returned {len(paths)} paths")
|
||||
except Exception as e:
|
||||
logger.error(f"Error querying GraphRAG: {e}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="结构化索引工具：在统一框架下构建并查询 RAPTOR（树状层次）与 "
|
||||
"GraphRAG（实体关系图）索引，对应本书实验 3-7。",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command", help="要执行的子命令")
|
||||
|
||||
# Build command
|
||||
build_parser = subparsers.add_parser(
|
||||
"build", help="从文档构建结构化索引（需要 OPENAI_API_KEY）")
|
||||
build_parser.add_argument("file", type=str,
|
||||
help="待索引的文档路径（支持 .pdf/.txt/.md/.html）")
|
||||
build_parser.add_argument("--type", choices=["raptor", "graphrag", "both"],
|
||||
default="both", help="要构建的索引类型（默认 both）")
|
||||
build_parser.add_argument("--output", type=str, default=None,
|
||||
help="将索引统计信息写入指定 JSON 文件")
|
||||
|
||||
# Query command
|
||||
query_parser = subparsers.add_parser(
|
||||
"query", help="查询已构建的索引（需要 OPENAI_API_KEY 及已有索引）")
|
||||
query_parser.add_argument("query", type=str, help="检索查询语句")
|
||||
query_parser.add_argument("--type", choices=["raptor", "graphrag", "both"],
|
||||
default="both", help="要查询的索引类型（默认 both）")
|
||||
query_parser.add_argument("--top-k", type=int, default=5,
|
||||
help="返回结果条数（默认 5）")
|
||||
query_parser.add_argument("--multi-hop", type=int, default=0, metavar="N",
|
||||
help="对 GraphRAG 额外执行 N 跳关系遍历（0 表示关闭）")
|
||||
query_parser.add_argument("--output", type=str, default=None,
|
||||
help="将查询结果写入指定 JSON 文件")
|
||||
|
||||
# Demo command（离线，无需 API）
|
||||
demo_parser = subparsers.add_parser(
|
||||
"demo", help="离线对比演示：结构化索引 vs 扁平检索（无需 API Key）")
|
||||
demo_parser.add_argument("--query", type=str, default=None,
|
||||
help="自定义查询；缺省时运行内置的三组对比查询")
|
||||
demo_parser.add_argument("--top-k", type=int, default=3,
|
||||
help="扁平检索展示的结果条数（默认 3）")
|
||||
demo_parser.add_argument("--output", type=str, default=None,
|
||||
help="将演示结果写入指定 JSON 文件")
|
||||
|
||||
# Server command
|
||||
subparsers.add_parser("serve", help="启动 HTTP API 服务")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "build":
|
||||
asyncio.run(build_indexes(Path(args.file), args.type, args.output))
|
||||
elif args.command == "query":
|
||||
results = asyncio.run(query_indexes(args.query, args.type, args.top_k,
|
||||
args.multi_hop))
|
||||
|
||||
# Display results
|
||||
for index_type, index_results in results.items():
|
||||
print(f"\n{index_type.upper()} Results:")
|
||||
print("-" * 50)
|
||||
if index_type == "graphrag_multi_hop":
|
||||
for i, r in enumerate(index_results, 1):
|
||||
chain = r["path"][0]["source"]
|
||||
for step in r["path"]:
|
||||
chain += f" --{step['relation']}--> {step['target']}"
|
||||
print(f"\n{i}. [{r['hops']} 跳] {chain}")
|
||||
continue
|
||||
for i, result in enumerate(index_results, 1):
|
||||
print(f"\n{i}. Score: {result.get('score', 'N/A'):.3f}")
|
||||
if 'summary' in result:
|
||||
print(f" Summary: {result['summary'][:200]}...")
|
||||
elif 'description' in result:
|
||||
print(f" Description: {result['description'][:200]}...")
|
||||
if 'level' in result:
|
||||
print(f" Level: {result['level']}")
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2, default=str)
|
||||
print(f"\n查询结果已写入：{args.output}")
|
||||
elif args.command == "demo":
|
||||
from structured_vs_flat_demo import run_demo
|
||||
run_demo(top_k=args.top_k, custom_query=args.query, output=args.output)
|
||||
elif args.command == "serve":
|
||||
from api_service import run_server
|
||||
run_server()
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,324 @@
|
||||
"""
|
||||
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
|
||||
try:
|
||||
import umap
|
||||
except ImportError: # Optional: deterministic PCA keeps the index usable.
|
||||
umap = None
|
||||
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 = []
|
||||
step = max(1, self.config.chunk_size - self.config.chunk_overlap)
|
||||
|
||||
for i in range(0, len(words), step):
|
||||
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 and umap is not None:
|
||||
reducer = umap.UMAP(n_components=50, n_neighbors=min(15, n_samples-1))
|
||||
embeddings_reduced = reducer.fit_transform(embeddings)
|
||||
elif embeddings.shape[1] > 50:
|
||||
from sklearn.decomposition import PCA
|
||||
dimensions = max(1, min(50, n_samples - 1, embeddings.shape[1]))
|
||||
embeddings_reduced = PCA(n_components=dimensions, random_state=42).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))
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
# RAPTOR dependencies
|
||||
raptor-rag==0.3.0
|
||||
openai>=1.12.0
|
||||
scikit-learn>=1.3.0
|
||||
numpy>=1.24.0
|
||||
tiktoken>=0.5.0
|
||||
umap-learn>=0.5.4
|
||||
|
||||
# GraphRAG dependencies
|
||||
graphrag>=0.3.0
|
||||
azure-search-documents>=11.4.0
|
||||
azure-storage-blob>=12.19.0
|
||||
networkx>=3.0
|
||||
pyarrow>=15.0.0
|
||||
pyyaml>=6.0
|
||||
rich>=13.0.0
|
||||
|
||||
# Common dependencies
|
||||
python-dotenv>=1.0.0
|
||||
requests>=2.31.0
|
||||
pandas>=2.0.0
|
||||
faiss-cpu>=1.7.4
|
||||
sentence-transformers>=2.2.2
|
||||
transformers>=4.36.0
|
||||
|
||||
# API service dependencies
|
||||
fastapi>=0.109.0
|
||||
uvicorn>=0.25.0
|
||||
pydantic>=2.5.0
|
||||
httpx>=0.25.0
|
||||
|
||||
# Document processing
|
||||
aiofiles>=24.1.0
|
||||
pypdf>=3.17.0
|
||||
beautifulsoup4>=4.12.0
|
||||
lxml>=5.0.0
|
||||
markdown>=3.5.0
|
||||
|
||||
# For Intel manual specific parsing
|
||||
pdfplumber>=0.10.3
|
||||
tabulate>=0.9.0
|
||||
|
||||
# Logging and monitoring
|
||||
loguru>=0.7.2
|
||||
tqdm>=4.66.0
|
||||
|
||||
# Testing
|
||||
pytest>=7.4.0
|
||||
pytest-asyncio>=0.21.0
|
||||
@@ -0,0 +1,299 @@
|
||||
"""
|
||||
结构化索引 vs 扁平检索：离线对比演示。
|
||||
|
||||
本模块不依赖 OpenAI / 向量模型 / 网络，纯 Python + networkx 即可运行。
|
||||
它用一个手工整理的「Intel x86 SIMD 指令集」小知识库，直观对比两条检索路线：
|
||||
|
||||
* 扁平检索（Flat）：把每个知识点当成互相独立的文本块，按词面相似度打分召回。
|
||||
这是传统 RAG「文档分块 + 向量检索」的抽象——只能返回零散片段。
|
||||
* 结构化检索（Structured）：
|
||||
- GraphRAG 式的实体-关系图：沿关系边做多跳遍历，能回答扁平检索答不了的
|
||||
「A 通过什么和 B 相连」这类关系性问题（对应书中「多跳关系推理」）。
|
||||
- RAPTOR 式的层次树：把细节聚合成上层摘要，能回答「概述某主题」这类
|
||||
需要跨片段综合的宏观问题（对应书中「多层次导航」）。
|
||||
|
||||
这段演示对应实验 3-7（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-7 的核心观点。")
|
||||
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()
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Regression test: GraphRAGIndexer.search must return empty list for non-positive top_k."""
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class STStub:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.encode_calls = 0
|
||||
|
||||
def encode(self, texts, **kwargs):
|
||||
self.encode_calls += 1
|
||||
return np.array([[0.1, 0.2, 0.3]])
|
||||
|
||||
|
||||
@dataclass
|
||||
class GraphRAGConfig:
|
||||
llm_api_key: str = "test"
|
||||
base_url: str = "test"
|
||||
llm_model: str = "test"
|
||||
|
||||
|
||||
_MISSING = object()
|
||||
_STUBBED_MODULES = (
|
||||
"openai",
|
||||
"sentence_transformers",
|
||||
"pandas",
|
||||
"sklearn",
|
||||
"sklearn.metrics",
|
||||
"sklearn.metrics.pairwise",
|
||||
"loguru",
|
||||
"tqdm",
|
||||
"config",
|
||||
"networkx",
|
||||
)
|
||||
|
||||
|
||||
class GraphStub:
|
||||
def __init__(self):
|
||||
self._neighbors = {}
|
||||
|
||||
def add_node(self, node):
|
||||
self._neighbors.setdefault(node, set())
|
||||
|
||||
def __contains__(self, node):
|
||||
return node in self._neighbors
|
||||
|
||||
def neighbors(self, node):
|
||||
return iter(self._neighbors[node])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _isolated_graphrag_module():
|
||||
modules = {name: types.ModuleType(name) for name in _STUBBED_MODULES}
|
||||
modules["openai"].OpenAI = object
|
||||
modules["sentence_transformers"].SentenceTransformer = STStub
|
||||
modules["sklearn"].__path__ = []
|
||||
modules["sklearn"].metrics = modules["sklearn.metrics"]
|
||||
modules["sklearn.metrics"].__path__ = []
|
||||
modules["sklearn.metrics"].pairwise = modules["sklearn.metrics.pairwise"]
|
||||
modules["sklearn.metrics.pairwise"].cosine_similarity = (
|
||||
lambda a, b: np.array([[0.95]])
|
||||
)
|
||||
modules["loguru"].logger = types.SimpleNamespace(
|
||||
info=lambda *a, **k: None,
|
||||
warning=lambda *a, **k: None,
|
||||
error=lambda *a, **k: None,
|
||||
)
|
||||
modules["tqdm"].tqdm = lambda x, **k: x
|
||||
modules["config"].GraphRAGConfig = GraphRAGConfig
|
||||
modules["networkx"].Graph = GraphStub
|
||||
|
||||
previous_module = sys.modules.pop("graphrag_indexer", _MISSING)
|
||||
try:
|
||||
with pytest.MonkeyPatch.context() as monkeypatch:
|
||||
for name, module in modules.items():
|
||||
monkeypatch.setitem(sys.modules, name, module)
|
||||
yield importlib.import_module("graphrag_indexer")
|
||||
finally:
|
||||
sys.modules.pop("graphrag_indexer", None)
|
||||
if previous_module is not _MISSING:
|
||||
sys.modules["graphrag_indexer"] = previous_module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def graphrag_module():
|
||||
with _isolated_graphrag_module() as module:
|
||||
yield module
|
||||
|
||||
|
||||
def _make_indexer(graphrag_module):
|
||||
indexer = graphrag_module.GraphRAGIndexer.__new__(
|
||||
graphrag_module.GraphRAGIndexer
|
||||
)
|
||||
indexer.config = graphrag_module.GraphRAGConfig()
|
||||
indexer.embedding_model = graphrag_module.SentenceTransformer()
|
||||
indexer.entities = {
|
||||
"e1": graphrag_module.Entity(
|
||||
"e1",
|
||||
"intel x86",
|
||||
"instruction",
|
||||
"intel x86 instruction",
|
||||
np.array([0.1, 0.2, 0.3]),
|
||||
{},
|
||||
),
|
||||
"e2": graphrag_module.Entity(
|
||||
"e2",
|
||||
"registers",
|
||||
"register",
|
||||
"intel registers",
|
||||
np.array([0.1, 0.2, 0.3]),
|
||||
{},
|
||||
),
|
||||
"e3": graphrag_module.Entity(
|
||||
"e3",
|
||||
"cpu flags",
|
||||
"feature",
|
||||
"cpu status flags",
|
||||
np.array([0.1, 0.2, 0.3]),
|
||||
{},
|
||||
),
|
||||
}
|
||||
indexer.communities = {}
|
||||
indexer.graph = graphrag_module.nx.Graph()
|
||||
for entity_id in indexer.entities:
|
||||
indexer.graph.add_node(entity_id)
|
||||
return indexer
|
||||
|
||||
|
||||
def test_search_nonpositive_top_k_returns_empty(graphrag_module):
|
||||
"""Non-positive result limits return before query encoding."""
|
||||
indexer = _make_indexer(graphrag_module)
|
||||
assert indexer.search("intel", top_k=0) == []
|
||||
assert indexer.search("intel", top_k=-1) == []
|
||||
assert indexer.search("intel", top_k=-5) == []
|
||||
assert indexer.embedding_model.encode_calls == 0
|
||||
|
||||
|
||||
def test_search_positive_top_k_returns_results(graphrag_module):
|
||||
"""Positive result limits still run retrieval and cap the results."""
|
||||
indexer = _make_indexer(graphrag_module)
|
||||
results = indexer.search("intel", top_k=2)
|
||||
assert len(results) == 2
|
||||
assert results[0]["id"] in ("e1", "e2", "e3")
|
||||
assert results[1]["id"] in ("e1", "e2", "e3")
|
||||
|
||||
|
||||
def test_dependency_stubs_are_restored():
|
||||
"""Scoped dependency replacements leave neighboring collection unchanged."""
|
||||
tracked_modules = (*_STUBBED_MODULES, "graphrag_indexer")
|
||||
before = {
|
||||
name: sys.modules.get(name, _MISSING)
|
||||
for name in tracked_modules
|
||||
}
|
||||
|
||||
with _isolated_graphrag_module() as module:
|
||||
assert sys.modules["graphrag_indexer"] is module
|
||||
for name in _STUBBED_MODULES:
|
||||
assert sys.modules[name] is not before[name]
|
||||
|
||||
for name, previous_module in before.items():
|
||||
assert sys.modules.get(name, _MISSING) is previous_module
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Unit tests for HybridStructuredRetriever covering core requirements and edge cases."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from hybrid_retriever import HybridStructuredRetriever, SearchResult
|
||||
|
||||
|
||||
def test_relation_target_included_in_text_content():
|
||||
"""Verify that relation matching includes target entity name."""
|
||||
retriever = HybridStructuredRetriever()
|
||||
retriever.add_graphrag_relationship(
|
||||
relation_id="rel_1",
|
||||
source="Attention",
|
||||
target="Transformer",
|
||||
type="USED_IN",
|
||||
description="Core mechanism for neural architecture",
|
||||
)
|
||||
|
||||
results = retriever.retrieve("Transformer", top_k=5)
|
||||
assert len(results) == 1
|
||||
assert results[0].node_id == "rel_1"
|
||||
assert "Transformer" in results[0].citation.citation_label or "Transformer" in results[0].text or "Transformer" in results[0].citation.lineage[1]
|
||||
|
||||
|
||||
def test_precision_calculation_unique_matched_words():
|
||||
"""Verify precision calculation counts unique matched query terms rather than token frequencies."""
|
||||
retriever = HybridStructuredRetriever()
|
||||
# Node content repeats "python" 5 times
|
||||
retriever.add_raptor_node(
|
||||
node_id="n1",
|
||||
level=0,
|
||||
text="python python python python python tutorial",
|
||||
)
|
||||
|
||||
# Query has 2 terms: python, fast
|
||||
query_terms = {"python", "fast"}
|
||||
node = retriever.unified_nodes["raptor_n1"]
|
||||
final_sc, lex_sc, sem_sc = retriever._compute_scores("python fast", query_terms, None, node)
|
||||
|
||||
# Lexical score should be 1 matched query term / 2 total query terms = 0.5
|
||||
assert lex_sc == 0.5
|
||||
|
||||
|
||||
def test_stringify_numeric_ids_in_citation():
|
||||
"""Verify that numeric node IDs, children, parents, and entity_ids do not cause TypeError during citation building."""
|
||||
retriever = HybridStructuredRetriever()
|
||||
|
||||
# Numeric IDs in RAPTOR node
|
||||
retriever.add_raptor_node(
|
||||
node_id=101,
|
||||
level=1,
|
||||
text="Hierarchical summary text",
|
||||
children=[201, 202],
|
||||
parent=50,
|
||||
)
|
||||
|
||||
# Numeric IDs in GraphRAG community
|
||||
retriever.add_graphrag_community(
|
||||
community_id=99,
|
||||
entity_ids=[1, 2, 3],
|
||||
summary="Community summary text",
|
||||
)
|
||||
|
||||
results = retriever.retrieve("summary text", top_k=5)
|
||||
assert len(results) == 2
|
||||
for res in results:
|
||||
assert isinstance(res.node_id, str)
|
||||
assert isinstance(res.citation.node_id, str)
|
||||
assert all(isinstance(lin, str) for lin in res.citation.lineage)
|
||||
|
||||
|
||||
def test_cache_node_embeddings():
|
||||
"""Verify that node embeddings computed via embedding_fn are cached in node['embedding']."""
|
||||
embed_count = 0
|
||||
|
||||
def mock_embed(text: str) -> np.ndarray:
|
||||
nonlocal embed_count
|
||||
embed_count += 1
|
||||
return np.array([0.1, 0.2, 0.3])
|
||||
|
||||
retriever = HybridStructuredRetriever(embedding_fn=mock_embed)
|
||||
retriever.add_raptor_node(
|
||||
node_id="n1",
|
||||
level=0,
|
||||
text="Machine learning models",
|
||||
)
|
||||
|
||||
node = retriever.unified_nodes["raptor_n1"]
|
||||
assert node["embedding"] is None
|
||||
|
||||
# First retrieval computes and caches embedding
|
||||
retriever.retrieve("Machine learning", top_k=5)
|
||||
assert node["embedding"] is not None
|
||||
assert isinstance(node["embedding"], np.ndarray)
|
||||
initial_count = embed_count
|
||||
|
||||
# Second retrieval reuses cached embedding
|
||||
retriever.retrieve("Machine learning", top_k=5)
|
||||
# embed_count should increase by 1 for query vector only, not for node embedding
|
||||
assert embed_count == initial_count + 1
|
||||
|
||||
|
||||
def test_top_k_zero_returns_empty_list():
|
||||
"""Verify that top_k == 0 returns an empty list immediately."""
|
||||
retriever = HybridStructuredRetriever()
|
||||
retriever.add_raptor_node(node_id="n1", level=0, text="Sample text")
|
||||
|
||||
assert retriever.retrieve("Sample", top_k=0) == []
|
||||
assert retriever.retrieve("Sample", top_k=-1) == []
|
||||
|
||||
|
||||
def test_clamp_rrf_k():
|
||||
"""Verify that rrf_k is clamped with max(1, int(rrf_k))."""
|
||||
retriever = HybridStructuredRetriever(rrf_k=0)
|
||||
assert retriever.rrf_k == 1
|
||||
|
||||
retriever_neg = HybridStructuredRetriever(rrf_k=-10)
|
||||
assert retriever_neg.rrf_k == 1
|
||||
|
||||
retriever.add_raptor_node(node_id="n1", level=0, text="Sample text")
|
||||
results = retriever.retrieve("Sample", rrf_k=-5)
|
||||
assert len(results) == 1
|
||||
assert results[0].score > 0
|
||||
|
||||
|
||||
def test_parent_id_zero_in_citation_lineage():
|
||||
"""Verify that parent ID 0 is preserved in citation lineage, not dropped by truthiness (Finding 12)."""
|
||||
retriever = HybridStructuredRetriever()
|
||||
retriever.add_raptor_node(
|
||||
node_id="child_node",
|
||||
level=1,
|
||||
text="Child content for retrieval",
|
||||
summary="Child summary for retrieval",
|
||||
parent=0,
|
||||
)
|
||||
|
||||
results = retriever.retrieve("Child content", top_k=1)
|
||||
assert len(results) == 1
|
||||
parent_entries = [lin for lin in results[0].citation.lineage if lin.startswith("Parent:")]
|
||||
assert len(parent_entries) == 1
|
||||
assert "0" in parent_entries[0]
|
||||
@@ -0,0 +1,192 @@
|
||||
"""
|
||||
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())
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Regression: equal chunk_size/overlap must not crash range() with step 0."""
|
||||
import sys
|
||||
import types
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
def _stub_raptor_deps() -> None:
|
||||
mods = [
|
||||
"tiktoken",
|
||||
"tqdm",
|
||||
"umap",
|
||||
"openai",
|
||||
"sentence_transformers",
|
||||
"loguru",
|
||||
"sklearn",
|
||||
"sklearn.mixture",
|
||||
"sklearn.metrics",
|
||||
"sklearn.metrics.pairwise",
|
||||
"config",
|
||||
]
|
||||
for name in mods:
|
||||
sys.modules.setdefault(name, types.ModuleType(name))
|
||||
sys.modules["sklearn.mixture"].GaussianMixture = object
|
||||
sys.modules["sklearn.metrics.pairwise"].cosine_similarity = lambda *a, **k: None
|
||||
sys.modules["openai"].OpenAI = object
|
||||
sys.modules["sentence_transformers"].SentenceTransformer = object
|
||||
sys.modules["loguru"].logger = types.SimpleNamespace(
|
||||
info=lambda *a, **k: None,
|
||||
error=lambda *a, **k: None,
|
||||
)
|
||||
sys.modules["tqdm"].tqdm = lambda x, **k: x
|
||||
|
||||
@dataclass
|
||||
class RaptorConfig:
|
||||
pass
|
||||
|
||||
sys.modules["config"].RaptorConfig = RaptorConfig
|
||||
|
||||
|
||||
_stub_raptor_deps()
|
||||
|
||||
from raptor_indexer import RaptorIndexer # noqa: E402
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Cfg:
|
||||
chunk_size: int = 1000
|
||||
chunk_overlap: int = 1000
|
||||
|
||||
|
||||
def test_chunk_text_equal_size_and_overlap():
|
||||
indexer = RaptorIndexer.__new__(RaptorIndexer)
|
||||
indexer.config = _Cfg()
|
||||
words = ("alpha beta gamma " * 200).strip()
|
||||
chunks = indexer.chunk_text(words)
|
||||
assert len(chunks) >= 1
|
||||
assert all(isinstance(c, str) and c for c in chunks)
|
||||
|
||||
|
||||
def test_chunk_text_normal_overlap_still_advances():
|
||||
indexer = RaptorIndexer.__new__(RaptorIndexer)
|
||||
indexer.config = _Cfg(chunk_size=10, chunk_overlap=2)
|
||||
chunks = indexer.chunk_text(" ".join(f"w{i}" for i in range(30)))
|
||||
assert len(chunks) > 1
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,94 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-7",
|
||||
"run_id": "20260729T200642Z-3_7-4d2d5f9c",
|
||||
"created_at": "2026-07-29T20:06:42.648599+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/structured-index/validation/runs/20260729T200642Z-3_7-4d2d5f9c",
|
||||
"artifacts": {
|
||||
"evidence.json": "89b3c03fbbd8e32304ed2639537a229c006ecf422bd8e9292bedf9201ee5810e",
|
||||
"receipts.json": "907dff07a0ee5496cfc2693cbd75ece8a7dc6f78eb76c928127cc98fba04c5e4",
|
||||
"manifest.json": "438f2a73543dd019ac047623b73ec94646ccd20b7dd31ae7b636fedca4e70c25"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/campaign.py",
|
||||
"sha256": "2d525c2b2ccfaf6a5a1cc53be3517b0a2dcb7658d5831365f71929ad38f9365c",
|
||||
"bytes": 27215
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/data/intel-sdm-volume-1.pdf",
|
||||
"sha256": "9d862bd7592d9fdd9f747c91d5e85be23ae3103f77185d1dcf7c5eb7277e5bdb",
|
||||
"bytes": 3645716
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/config.py",
|
||||
"sha256": "429a8b8991ce4b7548b3475009febc21710d846b707432f1b71b177b621834ff",
|
||||
"bytes": 6045
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/raptor_indexer.py",
|
||||
"sha256": "a8169f187a4324a7aab74dcdc9dd28c42a40f9414fb20e5cfa35e3b6802483b2",
|
||||
"bytes": 12699
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/graphrag_indexer.py",
|
||||
"sha256": "20fafbdb26ae1755c03af0af0bbcad1fd1663ae2e9235bdaad669721366c4862",
|
||||
"bytes": 25225
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"raptor": {
|
||||
"concept-detail": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 1.0,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 10825.52175
|
||||
},
|
||||
"relationship-multi-hop": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 0.9166666666666666,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 16744.6045
|
||||
},
|
||||
"overall": {
|
||||
"n": 8,
|
||||
"mean_citation_recall": 0.9583333333333334,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 13785.063125
|
||||
}
|
||||
},
|
||||
"graphrag": {
|
||||
"concept-detail": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 1.0,
|
||||
"mean_judge_score": 3.0,
|
||||
"mean_query_latency_ms": 11719.119749999998
|
||||
},
|
||||
"relationship-multi-hop": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 0.9166666666666666,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 15248.592
|
||||
},
|
||||
"overall": {
|
||||
"n": 8,
|
||||
"mean_citation_recall": 0.9583333333333334,
|
||||
"mean_judge_score": 3.5,
|
||||
"mean_query_latency_ms": 13483.855875000001
|
||||
}
|
||||
}
|
||||
},
|
||||
"acceptance": {
|
||||
"official_intel_pdf_pinned": true,
|
||||
"bounded_real_pages_extracted": true,
|
||||
"live_hierarchical_leaf_parent_root_summaries": true,
|
||||
"live_entity_relationship_extraction": true,
|
||||
"graph_communities_summarized": true,
|
||||
"concept_detail_and_relationship_multihop_sets": true,
|
||||
"both_indexes_answered_identical_queries": true,
|
||||
"actual_graph_paths_retained": true,
|
||||
"external_judge_complete": true,
|
||||
"raw_live_receipts_checkpointed": true
|
||||
}
|
||||
}
|
||||
+10268
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,93 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-7",
|
||||
"run_id": "20260729T200642Z-3_7-4d2d5f9c",
|
||||
"created_at": "2026-07-29T20:06:42.648599+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/structured-index/validation/runs/20260729T200642Z-3_7-4d2d5f9c",
|
||||
"artifacts": {
|
||||
"evidence.json": "89b3c03fbbd8e32304ed2639537a229c006ecf422bd8e9292bedf9201ee5810e",
|
||||
"receipts.json": "907dff07a0ee5496cfc2693cbd75ece8a7dc6f78eb76c928127cc98fba04c5e4"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/campaign.py",
|
||||
"sha256": "2d525c2b2ccfaf6a5a1cc53be3517b0a2dcb7658d5831365f71929ad38f9365c",
|
||||
"bytes": 27215
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/data/intel-sdm-volume-1.pdf",
|
||||
"sha256": "9d862bd7592d9fdd9f747c91d5e85be23ae3103f77185d1dcf7c5eb7277e5bdb",
|
||||
"bytes": 3645716
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/config.py",
|
||||
"sha256": "429a8b8991ce4b7548b3475009febc21710d846b707432f1b71b177b621834ff",
|
||||
"bytes": 6045
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/raptor_indexer.py",
|
||||
"sha256": "a8169f187a4324a7aab74dcdc9dd28c42a40f9414fb20e5cfa35e3b6802483b2",
|
||||
"bytes": 12699
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/structured-index/graphrag_indexer.py",
|
||||
"sha256": "20fafbdb26ae1755c03af0af0bbcad1fd1663ae2e9235bdaad669721366c4862",
|
||||
"bytes": 25225
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"raptor": {
|
||||
"concept-detail": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 1.0,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 10825.52175
|
||||
},
|
||||
"relationship-multi-hop": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 0.9166666666666666,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 16744.6045
|
||||
},
|
||||
"overall": {
|
||||
"n": 8,
|
||||
"mean_citation_recall": 0.9583333333333334,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 13785.063125
|
||||
}
|
||||
},
|
||||
"graphrag": {
|
||||
"concept-detail": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 1.0,
|
||||
"mean_judge_score": 3.0,
|
||||
"mean_query_latency_ms": 11719.119749999998
|
||||
},
|
||||
"relationship-multi-hop": {
|
||||
"n": 4,
|
||||
"mean_citation_recall": 0.9166666666666666,
|
||||
"mean_judge_score": 4.0,
|
||||
"mean_query_latency_ms": 15248.592
|
||||
},
|
||||
"overall": {
|
||||
"n": 8,
|
||||
"mean_citation_recall": 0.9583333333333334,
|
||||
"mean_judge_score": 3.5,
|
||||
"mean_query_latency_ms": 13483.855875000001
|
||||
}
|
||||
}
|
||||
},
|
||||
"acceptance": {
|
||||
"official_intel_pdf_pinned": true,
|
||||
"bounded_real_pages_extracted": true,
|
||||
"live_hierarchical_leaf_parent_root_summaries": true,
|
||||
"live_entity_relationship_extraction": true,
|
||||
"graph_communities_summarized": true,
|
||||
"concept_detail_and_relationship_multihop_sets": true,
|
||||
"both_indexes_answered_identical_queries": true,
|
||||
"actual_graph_paths_retained": true,
|
||||
"external_judge_complete": true,
|
||||
"raw_live_receipts_checkpointed": true
|
||||
}
|
||||
}
|
||||
+4751
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user