dalgona039/TreeRAG

# TreeRAG TreeRAG 是一个层次化的 RAG 系统，它将 PDF 文档索引为树状结构的 JSON，并通过可溯源到页码的引用来回答问题。 该项目包含： - 用于上传、索引、聊天、树结构、缓存和图谱 API 的 FastAPI 后端 - 具备 DFS/Beam 遍历、可选上下文压缩以及可选交叉引用解析的 TreeRAG 推理器 - 用于多文档聊天和树结构探索的 Next.js 前端 - 用于 BM25、FlatRAG、TreeRAG-DFS、TreeRAG-Beam 及相关基线的基准/评估脚本 ## 为什么选择 TreeRAG 传统的扁平检索通常会丢失章节层级。TreeRAG 保留了文档结构，并且仅遍历相关的分支。 核心理念： - 从 PDF 构建具备页码感知的层次化索引 - 将查询路由到一个或多个已索引的文档 - 遍历树节点（DFS 或 Beam Search） - 通过来源依据和引用提取生成答案 - 缓存响应并执行 API 速率限制，以确保稳定的服务 ## 主要特性 - 具备自动文档路由的多文档查询 - 深度遍历模式（TreeNavigator 或 BeamSearchNavigator） - 在最终生成前进行可选的上下文压缩 - 从生成的答案中提取引用 - 基于节点依据置信度的幻觉警告层 - 带有统计信息和清除端点的内存响应缓存 - API 速率限制（SlowAPI） - 使用 Celery/Redis 的可选异步索引任务端点 - 推理图谱构建/加载端点 ## 技术栈 后端： - Python - FastAPI - google-genai (Gemini API) - pypdf - slowapi - Redis + Celery（可选的异步任务队列） 前端： - Next.js 16 - React 19 - TypeScript 评估： - pytest - benchmarks/ 下的基准测试脚本 ## 仓库结构 - src/ - api/: REST 路由及请求/响应模型 - core/: 索引器、推理器、遍历、Beam Search、压缩器、基线 - middleware/: 安全标头 - repositories/: 会话持久化 - utils/: 缓存、文件验证、幻觉检测 - benchmarks/: 评估运行器、数据集、分析、指标 - frontend/: Next.js 应用 - data/ - raw/: 上传的 PDF - indices/: 生成的索引 JSON 文件和图谱文件 - benchmark_reports/: 基准测试输出 ## 前置条件 - Python 3.11+（本仓库中已使用 Python 3.13 测试） - 用于前端的 Node.js 20+ - Gemini API key ## 快速开始（本地） ### 1) 创建环境并安装后端依赖 ``` cd /Volumes/a3122a1/TreeRAG python3 -m venv .venv source .venv/bin/activate pip install -U pip pip install -r requirements.txt ``` 对于需要额外指标包的基准测试脚本： ``` pip install rouge-score rank-bm25 scipy ``` ### 2) 配置环境变量 在仓库根目录创建 .env： ``` GOOGLE_API_KEY=your_key_here ``` 可选的 Gemini 弹性控制（建议在免费层使用）： ``` GEMINI_MIN_INTERVAL_S=13 GEMINI_MAX_RETRIES=5 ``` 注意事项： - GEMINI_MIN_INTERVAL_S 通过限制请求间隔来减少 429 爆发错误 - 代码同样会在 src/config.py 中通过退避机制重试配额类错误 ### 3) 运行后端 ``` source .venv/bin/activate python main.py ``` 后端端点： - API 基础地址: http://localhost:8000/api - OpenAPI 文档: http://localhost:8000/docs - 健康检查: http://localhost:8000/health 和 /health/deep ### 4) 运行前端 ``` cd frontend npm install npm run dev ``` 前端默认 URL: http://localhost:3000 ## 快速开始（Docker） ``` cd /Volumes/a3122a1/TreeRAG # 确保 .env 包含 GOOGLE_API_KEY docker-compose up --build -d ``` 服务： - frontend: 3000 - backend: 8000 - redis: 6379 - celery-worker（compose 中包含的可选 worker 服务） 有关操作命令，请参阅 DOCKER.md。 ## 基本 API 流程 1. 上传 PDF ``` curl -F "file=@/absolute/path/to/doc.pdf" http://localhost:8000/api/upload ``` 2. 创建索引 ``` curl -X POST http://localhost:8000/api/index \ -H "Content-Type: application/json" \ -d '{"filename":".pdf"}' ``` 3. 查询聊天 ``` curl -X POST http://localhost:8000/api/chat \ -H "Content-Type: application/json" \ -d '{ "question":"핵심 내용을 요약해줘", "index_filenames":["_index.json"], "use_deep_traversal":true, "max_depth":5, "max_branches":3, "domain_template":"general", "language":"ko" }' ``` ## 主要 API 端点 文档和聊天： - GET /api/ - POST /api/upload - POST /api/index - POST /api/chat - GET /api/indices - GET /api/pdfs - GET /api/tree/{index_filename} - GET /api/pdf/{filename} 会话： - GET /api/sessions - PUT /api/sessions 缓存： - GET /api/cache/stats - POST /api/cache/clear 推理图谱： - POST /api/graph/build/{document_name} - GET /api/graph/{document_name} 异步任务（需要 Celery 可用）： - POST /api/tasks/index - POST /api/tasks/index/batch - GET /api/tasks/{task_id} - DELETE /api/tasks/{task_id} - GET /api/tasks/ ## 遍历与推理行为 推理器类: src/core/reasoner.py - 遍历算法: - DFS (TreeNavigator) - Beam Search (BeamSearchNavigator) - 可选的上下文压缩 - 用于节/章引用的可选引用解析器 - 领域 prompt 模板：通用、医疗、法律、金融、学术 - 语言指令：ko/en/ja - 缓存 key 包含 prompt 缓存版本和遍历设置 ## 基准测试与评估 主要运行器: - benchmarks/run_real_evaluation.py 示例（低成本的方向检查）： ``` source .venv/bin/activate export GEMINI_MIN_INTERVAL_S=13 python benchmarks/run_real_evaluation.py \ --systems bm25,flatrag,treerag_beam \ --limit 3 --mode online \ --output data/benchmark_reports/direction_check.json ``` 脚本辅助工具： ``` bash scripts/run_online_direction_check.sh ``` 注意事项： - 在线模式会快速消耗 Gemini 配额，尤其是 TreeRAG-Beam - 免费层的限制可能会导致 429/503 错误以及部分内容为空的答案 - 如果由于配额耗尽导致结果行为空，请在进行模型质量比较时将该次运行视为无效 ## 测试 运行所有测试： ``` source .venv/bin/activate pytest ``` 运行特定测试： ``` pytest tests/test_config_resilient.py -v ``` pytest 配置位于 pytest.ini 中。 ## 安全与验证 已实施的安全防护措施包括： - 上传文件验证（扩展名、MIME、大小、路径遍历） - 安全标头中间件 - API 速率限制 - 通过 repository 层实现会话持久化 为了保持本地开发环境的整洁，请运行： ``` bash setup-git-hooks.sh ``` ## 故障排除 ### 1) 429 RESOURCE_EXHAUSTED 症状： - 基准测试中出现空响应或失败响应 - 日志中频繁出现配额错误 措施： - 设置 GEMINI_MIN_INTERVAL_S=13（或更大） - 减少 benchmark limit 和/或测试系统 - 在低成本检查中禁用 LLM judge - 如果超出每日限额，请在配额重置后重试 ### 2) 503 UNAVAILABLE 症状： - 瞬时的模型过载错误 措施： - 稍后重试（服务端拥堵） - 保持开启重试（GEMINI_MAX_RETRIES） ### 3) 未找到已索引的文档 措施： - 验证文件是否成功上传至 data/raw - 为每个上传的 PDF 调用 /api/index - 检查 data/indices 下的文件 ### 4) 前端无法调用后端 措施： - 确保后端正在 8000 端口上运行 - 在使用 Docker/生产环境时验证 NEXT_PUBLIC_API_BASE_URL ## 许可证 MIT（详见 LICENSE）。 ## 致谢 本项目受树状结构和无向量 RAG 方法的影响，并作为用于层次化文档推理与评估的工程/研究代码库进行开发。

标签：AV绕过, DLL 劫持, FastAPI, RAG系统, 人工智能, 大语言模型, 安全规则引擎, 搜索引擎查询, 文档检索, 用户模式Hook绕过, 请求拦截, 逆向工具