Rag cn
本文档深入解析 Jianmu 的检索增强生成(RAG)模块——一个基于协议驱动、可插拔组件架构的知识检索流水线。从文档摄取到向量化检索,从多策略混合排序到 Agent 工具封装,我们将逐层揭开其设计全貌。
架构分层:协议契约驱动的组件矩阵¶
Jianmu 的 RAG 体系构建在一组核心协议之上,每一个关键环节都通过 Protocol 定义接口契约,使得实现可以自由替换。整体分为六个层次:
graph TB
subgraph 配置层
CFG[RagEmbeddingConfig<br/>provider / model / base_url]
LIM[LimitsConfig<br/>chunk_size / chunk_overlap / rerank_alpha]
end
subgraph 协议层
EMB_P[EmbedderProtocol]
STORE_P[VectorStoreProtocol]
RET_P[RetrieverProtocol]
RERANK_P[RerankerProtocol]
LOAD_P[DocumentLoaderProtocol]
NORM_P[TextNormalizerProtocol]
CHUNK_P[TextChunkerProtocol]
end
subgraph 实现层
EMB[SimpleEmbedder / OpenAIEmbedder / GeminiEmbedder]
STORE[InMemoryStore / SQLiteStore / FaissStore / ChromaStore]
RET[BM25Retriever / DenseRetriever / HybridRetriever]
RERANK[PassthroughReranker]
INGEST[DefaultDocumentLoader / WhitespaceNormalizer / SimpleTextChunker]
end
subgraph 管线层
IP[IngestPipeline]
RP[RAGPipeline]
KB[KnowledgeBase 外观]
end
subgraph 工具层
KST[KnowledgeSearchTool]
end
CFG --> EMB
CFG --> LIM
LIM --> INGEST
EMB_P -.-> EMB
STORE_P -.-> STORE
RET_P -.-> RET
RERANK_P -.-> RERANK
LOAD_P -.-> INGEST
NORM_P -.-> INGEST
CHUNK_P -.-> INGEST
EMB --> RP
STORE --> RP
RET --> RP
RERANK --> RP
INGEST --> IP
IP --> KB
RP --> KB
KB --> KST
协议层定义了 7 个核心接口,分别对应嵌入、存储、检索、重排序、加载、归一化和分块。实现层提供了每个协议的多种具体实现,而管线层则将它们组合为可用的工作流。最顶层的 KnowledgeBase 是一个外观类,将摄取管道和检索管道统一封装。
数据模型:Document、SearchResult 与 SearchOptions¶
整个 RAG 系统围绕三个核心数据结构运转。
Document 是一个 @dataclass,包含五个字段:id(稳定标识符)、text(文档正文)、metadata(任意键值对元数据,可用于过滤)、created_at(UTC 创建时间戳)和 embedding(可选的缓存稠密向量)。向量缓存机制允许文档被检索时无需反复调用嵌入服务。
SearchResult 将匹配到的 Document 与一个 score 分数绑定,分数语义取决于检索器类型——对稠密检索来说是余弦相似度,对 BM25 来说是统计相关性分数,对混合检索来说是归一化后的融合分数。
SearchOptions 提供检索的运行时配置:k(返回数量,默认 5)、mode(检索模式:"vector"、"keyword" 或 "hybrid")、alpha(混合检索的稠密-稀疏加权系数,默认从配置读取 limits.rerank_alpha,出厂值 0.6)、filter(元数据等值过滤字典)。
嵌入器:SimpleEmbedder → OpenAIEmbedder → GeminiEmbedder¶
嵌入器负责将文本转换为稠密向量。所有嵌入器均实现 EmbedderProtocol 的两个方法:embed_documents(批量嵌入文档)和 embed_query(嵌入单条查询)。
SimpleEmbedder 是一个无需任何外部依赖的确定性嵌入器。它基于字符编码的模运算生成向量,每个字符的位置加权累加到对应维度,最终做单位归一化。默认维度为 64,适合测试或作为回退方案。
OpenAIEmbedder 同时兼容新版 OpenAI SDK(openai.OpenAI)和旧版 openai 包。API Key 的读取优先级为:构造函数参数 → OPENAI_API_KEY → API_KEY;base_url 的优先级为:构造函数参数 → 配置 rag.embedding.base_url → OPENAI_BASE_URL → OPENAI_API_BASE → BASE_URL;模型名的优先级为:构造函数参数 → 配置 rag.embedding.model → OPENAI_EMBED_MODEL → EMBEDDING_MODEL,最终回退到 text-embedding-3-small。
GeminiEmbedder 基于 Google genai SDK,模型默认使用 text-embedding-004,API Key 的查找顺序为 GOOGLE_API_KEY 或 GEMINI_API_KEY。
resolve_embedder() 是一个自动解析函数,按偏好列表(默认 ["gemini", "openai", "simple"])检查可用的 API Key,返回第一个可实例化的嵌入器。如果所有外部提供商都不可用,最终回退到 SimpleEmbedder,确保系统在任何环境下都能运行。
向量存储:四种后端的统一接口¶
所有向量存储均实现 VectorStoreProtocol 的三个方法:add(插入文档与嵌入向量)、search(按查询向量检索)、delete(按 ID 删除)。
| 存储后端 | 依赖 | 持久化 | 适用场景 |
|---|---|---|---|
InMemoryStore |
无 | 内存 | 测试、快速原型、临时知识库 |
SQLiteStore |
sqlite3(标准库) |
文件 | 小型生产、单机部署 |
FaissStore |
faiss-cpu/gpu |
内存 + 重建 | 大规模向量检索、高性能场景 |
ChromaStore |
chromadb |
文件 | 全功能向量数据库、元数据过滤 |
InMemoryStore 支持 max_size 参数做 FIFO 淘汰,当文档数超出上限时自动移除最早插入的文档。其 search 方法在查询向量为空时返回最新的 k 个文档(分数统一为 1.0),在元数据过滤方面使用简单的等值匹配。
SQLiteStore 将文档以 text、metadata(JSON 字符串)、embedding(JSON BLOB)格式存入单表,路径默认从配置 paths.knowledge_db 读取,父目录自动创建。其检索在 Python 侧计算余弦相似度。
FaissStore 使用 Facebook 的 FAISS 库进行近邻搜索,支持通过 index_factory 字符串自定义索引结构(默认 "Flat"),支持内积(IP)和 L2 两种距离度量,自动做向量归一化。删除操作通过重建索引实现:保留未被删除的文档,重建 IndexIDMap2 并重新添加。
ChromaStore 封装 ChromaDB 的 PersistentClient,集合名默认为 "jianmu"。内置的 query 方法直接支持元数据 where 过滤,距离分数转换为相似度分数(1.0 - distance)。
检索器:BM25 → Dense → Hybrid 的三级递进¶
检索器实现 RetrieverProtocol 的 search 方法,接收查询字符串、k 值和可选的 SearchOptions。
BM25Retriever 实现了经典的 BM25 概率检索模型。它在搜索时先通过 list_all() 获取所有文档,构建语料级别的 IDF 统计,然后对每条查询词计算 tf * idf / (tf + k1 * (1 - b + b * dl/avgdl)) * (k1 + 1) 的 BM25 分数。分词采用小写字母数字正则 \w+。
DenseRetriever 使用嵌入器对查询向量化,然后优先调用存储后端的 search 方法(利用向量索引的近邻搜索能力);如果后端不支持,则回退到 Python 侧遍历所有文档计算余弦相似度。查询向量经 _normalize_vector 做单位归一化,并通过 _coerce_embedding 统一处理各种嵌入返回格式(list、numpy array 等)。
HybridRetriever 是核心的混合检索器,也是 KnowledgeBase 默认使用的检索器。其工作流程分为四步:
- 并行检索:同时运行 BM25 和 Dense 检索,各取
k * 2个候选结果 - 分数归一化:将两路分数各自缩放到 [0, 1] 区间(除以最大值)
- 加权融合:按
alpha * dense_score + (1 - alpha) * sparse_score计算最终分数 - 排序截断:按融合分数降序排列,返回前 k 个结果
模式控制通过 SearchOptions.mode 实现:"keyword" 模式仅使用 BM25 分数(alpha 等效为 0),"semantic" 模式仅使用稠密分数(alpha 等效为 1),"hybrid" 模式使用配置的 alpha 值。
重排序器:可插拔的 PassthroughReranker¶
RerankerProtocol 定义 rerank(query, results) -> List[SearchResult] 接口,接收原始检索结果并返回重排后的列表。
当前唯一的实现是 PassthroughReranker——它原封不动地返回输入结果。这是有意为之的设计:重排序接口已经就位,为未来接入 Cross-Encoder(如 Cohere Rerank、BGE-Reranker 等)预留了扩展点。RAGPipeline.search() 方法中预留了 reranker 的两段式逻辑:当 reranker 存在时,先以 k * 2 的数量检索,重排后再截取前 k 个。
摄取管道:Load → Normalize → Chunk → Embed → Store¶
IngestPipeline 将五个步骤串联为一个流水线。
DefaultDocumentLoader 支持三种文件格式:.txt(直接读取)、.pdf(通过 pypdf.PdfReader 逐页提取文本)、.docx(通过 python-docx 按段落提取)。对 .doc 格式会显式报错提示转换。编码默认 UTF-8。
WhitespaceNormalizer 执行基础的空白字符规范化:去首尾空白,将连续空白符(含换行)压缩为单个空格。
SimpleTextChunker 使用固定大小的滑动窗口分块。两个控制参数 chunk_size(默认 500 字符)和 overlap(默认 50 字符)均从配置 limits 读取。分块逻辑为:从文本起始位置开始,每次取 chunk_size 个字符,然后回退 overlap 个字符作为下一个窗口的起点。
ingest_text() 方法是核心入口:归一化文本 → 分块 → 为每块创建带 chunk_index 元数据的 Document → 批量嵌入 → 将嵌入缓存到每个 Document 的 embedding 字段 → 存入向量存储。
ingest_file() 是文件入口,在调用 ingest_text 前先通过 loader 加载文件内容,并自动添加 source 元数据标记来源路径。
RAGPipeline:检索 + 可选重排序¶
RAGPipeline 是查询侧的流水线,构造时注入 RetrieverProtocol 和可选的 RerankerProtocol。其 search() 方法实现了两步策略:
- 候选召回:如果配置了 reranker,则检索
k * 2条结果作为宽召回池;否则直接检索 k 条 - 精排截断:如果有 reranker 则调用
rerank()重排后取前 k 条
这使得在不增加 reranker 时零开销(直通检索器),在接入 reranker 时有足够的候选项供精排模型发挥作用。
KnowledgeBase:一站式知识库外观¶
KnowledgeBase 是面向用户的最高层抽象,将 IngestPipeline、RAGPipeline、向量存储、嵌入器和检索器统一封装为一个外观类。
sequenceDiagram
participant User
participant KB as KnowledgeBase
participant IP as IngestPipeline
participant RP as RAGPipeline
participant EMB as Embedder
participant Store as VectorStore
participant RET as HybridRetriever
User->>KB: add(text, metadata)
KB->>EMB: embed_query(text)
EMB-->>KB: embedding
KB->>Store: add([doc], [embedding])
Store-->>KB: doc_id
User->>KB: ingest_file(path)
KB->>IP: ingest_file(path)
IP->>IP: load → normalize → chunk
IP->>EMB: embed_documents(chunks)
EMB-->>IP: embeddings
IP->>Store: add(docs, embeddings)
User->>KB: search(query, k, mode)
KB->>RP: search(query, options)
RP->>RET: search(query, k, options)
RET->>RET: BM25 + Dense → 归一化 → 融合
RET-->>RP: results
RP-->>KB: results
KB-->>User: List[SearchResult]
User->>KB: as_tool()
KB-->>User: KnowledgeSearchTool
add() 方法提供了最简单的单文档摄入:创建 Document → 调用 embed_query 生成向量 → 缓存向量到 Document → 通过 store.add 持久化。
as_tool() 方法将 _rag 管道包装为 KnowledgeSearchTool,使其可以直接注入到 Agent 的工具列表中。这意味着只需 3 行代码即可将知识库作为 Agent 的可搜索外部记忆使用:
KnowledgeSearchTool 是一个标准的 Tool 子类,定义了 knowledge_search 的输入 schema(query、k、mode),run() 方法内部调用 RAGPipeline.search() 并将结果格式化为编号列表返回。
配置体系:RagEmbeddingConfig 与 LimitsConfig¶
RAG 模块的配置分布在两处:
RagEmbeddingConfig(位于 jianmu/config/project.py,属于 jianmu.yaml 的 rag.embedding 节)控制嵌入器行为:
| 字段 | 类型 | 默认值 | 说明 |
|---|---|---|---|
provider |
str |
"openai" |
首选嵌入提供商 |
model |
Optional[str] |
None |
嵌入模型名(None 时使用各 Provider 内置默认) |
base_url |
Optional[str] |
None |
自定义 API 端点 |
LimitsConfig 中三个字段控制摄取和检索行为:
| 字段 | 类型 | 默认值 | 说明 |
|---|---|---|---|
chunk_size |
int |
500 |
文本分块大小(字符数) |
chunk_overlap |
int |
50 |
相邻分块重叠字符数 |
rerank_alpha |
float |
0.6 |
混合检索中稠密分数的权重 |
配置值通过 _default_chunker() 工厂和 _get_default_alpha() 工厂注入到 IngestPipeline 和 HybridRetriever 的默认构造中。在 jianmu.yaml 中配置示例:
rag:
embedding:
provider: openai
model: text-embedding-3-large
base_url: https://api.openai.com/v1
limits:
chunk_size: 800
chunk_overlap: 100
rerank_alpha: 0.7
示例用法:从文档摄取到 Agent 问答¶
完整的 RAG 使用流程在 examples/rag/rag_demo.py 中展示。其核心步骤:
-
构建知识库:使用
SQLiteStore作为持久化后端,配置自定义的SimpleTextChunker(chunk_size=300, overlap=50) -
文档摄取:通过
--path参数指定文件或目录,支持递归遍历和扩展名过滤(默认.md,.txt,.pdf,.docx,.doc),每个文件由DefaultDocumentLoader加载后调用kb.ingest_text()摄入 -
Agent 集成:调用
kb.as_tool()获取KnowledgeSearchTool,将其注入 ReAct Agent 的工具列表 -
问答执行:Agent 在推理过程中可以调用
knowledge_search工具,按需检索知识库中的相关内容
这种设计使得 RAG 能力完全融入 Jianmu 的 Agent 运行时:Agent 通过工具调用访问知识库,检索结果会自动进入消息历史,供 LLM 在下一次推理时引用。
扩展指南:接入自定义组件¶
得益于协议驱动的设计,替换或扩展任意组件只需实现对应的 Protocol。例如,接入 Cohere 的重排序服务:
from jianmu.rag.base import RerankerProtocol
from jianmu.rag.types import SearchResult
class CohereReranker(RerankerProtocol):
def rerank(self, query: str, results: List[SearchResult]) -> List[SearchResult]:
# 调用 Cohere Rerank API,按 relevance_score 重排
...
return reranked
kb = KnowledgeBase(reranker=CohereReranker())
类似地,可以通过实现 EmbedderProtocol 接入任何嵌入服务,通过实现 VectorStoreProtocol 接入任何向量数据库。所有管线类都接受通过构造函数注入的自定义实现,无需修改框架代码。
阅读建议:RAG 流水线向上承接 上下文构建器:消息过滤、Token 预算控制与多源 Prompt 装配 的消息整合能力,向下为 Agent 的 ReAct 节点工厂 提供外部知识检索工具。在 Swarm 多 Agent 场景中,不同 Agent 可以共享同一个 KnowledgeBase 实例或各自持有独立的知识库,由 Swarm 运行时 统一编排。