16-pdf-rag

PDF RAG 示例

本示例演示如何使用 GoAgent 构建一个基于 PDF 文档的 RAG（检索增强生成）系统。

功能特点

• PDF 文档加载：支持从本地文件或 URL 加载 PDF 文档
• 文本分块：使用递归字符分割器将长文档分割成小块
• 向量索引：使用内存向量存储进行文档索引
• 语义检索：基于余弦相似度检索相关文档片段
• 增强生成：将检索到的上下文与 LLM 结合生成准确回答

系统架构

┌─────────────┐    ┌─────────────┐    ┌─────────────┐
│  PDF 文档   │───▶│  文本分块   │───▶│  向量存储   │
└─────────────┘    └─────────────┘    └─────────────┘
                                            │
                                            ▼
┌─────────────┐    ┌─────────────┐    ┌─────────────┐
│  用户查询   │───▶│  语义检索   │───▶│  RAG 链     │
└─────────────┘    └─────────────┘    └─────────────┘
                                            │
                                            ▼
                                    ┌─────────────┐
                                    │  LLM 生成   │
                                    └─────────────┘

前置条件

1. 设置 DeepSeek API Key：

   export DEEPSEEK_API_KEY=your-api-key
   

2. 准备一个 PDF 文档用于测试

使用方法

基本用法

# 使用 RAG 模式查询 PDF 文档
go run main.go -pdf /path/to/document.pdf -prompt "文档的主要内容是什么？"

# 不使用 RAG，直接查询 LLM
go run main.go -prompt "什么是 Go 语言？" -rag=false

命令行参数

示例

# 示例 1: 查询技术文档
go run main.go \
    -pdf ./tech-manual.pdf \
    -prompt "如何配置系统参数？" \
    -chunk-size 400 \
    -topk 5

# 示例 2: 查询产品手册
go run main.go \
    -pdf ./product-guide.pdf \
    -prompt "产品的主要功能有哪些？"

# 示例 3: 对比有无 RAG 的效果
# 启用 RAG
go run main.go -pdf ./doc.pdf -prompt "文档中提到的关键技术是什么？"

# 禁用 RAG（直接使用 LLM 知识）
go run main.go -prompt "文档中提到的关键技术是什么？" -rag=false

输出示例

========================================
PDF RAG 演示
========================================

[RAG 模式] 使用 PDF 文档增强回答

步骤 1: 加载 PDF 文档...
  已加载 PDF，共 12 个文档片段
  总字符数: 8456

步骤 2: 创建向量存储并索引文档...
  已索引 12 个文档片段

步骤 3: 创建 RAG 检索器...
  TopK: 4, 阈值: 0.30

步骤 4: 创建 RAG 链...

步骤 5: 执行 RAG 查询...
问题: 文档的主要内容是什么？

检索到的相关文档:
----------------------------------------
[1] (相似度: 0.8523)
    本文档介绍了系统的核心功能和配置方法...

[2] (相似度: 0.7891)
    主要功能包括数据处理、报表生成、权限管理...

回答:
----------------------------------------
根据文档内容，系统的主要功能包括：
1. 数据处理和分析
2. 自动报表生成
3. 用户权限管理
...
----------------------------------------

总耗时: 2.34s

工作原理

1. 文档加载：使用 document.PDFLoader 加载 PDF 文件并提取文本内容

2. 文本分块：使用 document.RecursiveCharacterTextSplitter 将长文本分割成适合检索的小块

   • 默认块大小：500 字符
   • 默认重叠：100 字符

3. 向量索引：使用 retrieval.MemoryVectorStore 存储文档向量

   • 使用简单嵌入器生成向量（演示用途）
   • 支持余弦相似度检索

4. 语义检索：使用 retrieval.RAGRetriever 检索相关文档

   • 基于查询向量检索 Top-K 最相似的文档片段
   • 支持相似度阈值过滤

5. 增强生成：使用 retrieval.RAGChain 组合检索和生成

   • 将检索到的文档作为上下文
   • 调用 LLM 生成基于上下文的回答

生产环境建议

1. 使用专业嵌入模型：

   • 替换 SimpleEmbedder 为 OpenAI Embeddings 或其他专业模型
   • 可以获得更准确的语义检索结果

2. 使用持久化向量存储：

   • 替换 MemoryVectorStore 为 Qdrant 或其他向量数据库
   • 支持大规模文档索引和持久化存储

3. 优化分块策略：

   • 根据文档类型调整分块大小和重叠
   • 考虑使用语义分块而非固定字符分块

4. 添加重排序：

   • 使用 retrieval.MMRReranker 或 CrossEncoderReranker 提升检索质量
   • 平衡相关性和多样性

参考资料

• RAG PDF Annoy 示例
• GoAgent Retrieval 模块文档