数据截至 (上游 commit 4a030776a31a)
语义检索与树扩展:一次查询如何取回带上下文的片段
30 秒导读: 上一章 01-ingestion-tree-nodes 把文件切成了一棵棵「文档树」的节点存进向量库。本章讲读取侧:用户一句话进来,系统怎么把最相关的几个片段捞出来——而且不是捞出孤零零的一句,而是顺着树把它前后、上级的上下文一起补齐,拼成一段读得懂的连续文本。这是 PrivateGPT 检索的核心算法。
本章只讲「取回带上下文的片段」这一步。至于这些片段之后怎么进 prompt、怎么触发检索,见 04-agentic-rag-chat-loop 与 05-context-stack-assembly;怎么把答案锚回原文(引用格式)见 03-citations。
1. 这是什么(零基础也能懂)
一句话定义: 给一段查询文本,返回一组按相关度排序、且各自带着前后上下文的文本片段(Chunk)。
它解决什么问题。 朴素的向量检索有个老毛病:你把文档切成小块存进库,查询命中的往往只是其中一小块——比如命中了「营收增长了 15%」,但「相比哪个季度、为什么增长」在相邻的块里,没被捞出来。模型拿到一句没头没尾的话,容易答歪。
PrivateGPT 的解法: 命中一块之后,顺着摄取时建好的文档树,把这块的前一段、后一段、乃至它所属章节的上下文一起补回来,直到攒够一个 token 预算为止。于是模型看到的不是一句孤立的话,而是一小段连贯的原文。
一句话直觉/类比: 你在一本书里用关键词搜到某一行,好的阅读器不会只给你那一行,而是把那一段、必要时那一节都显示出来让你读懂上下文。PrivateGPT 做的就是这件事,只不过「段落/章节」的边界来自它自己建的树。
用起来什么样。 调用方给一段文本 + 一个 ContextFilter(限定在哪个知识库/哪些文件里搜),拿回一列 Chunk:
# 示意,非源码:一次检索长什么样
chunks = semantic_search.retrieve_semantic_relevant(
text="第三季度营收为什么增长?",
context_filter=ContextFilter(collection="finance", artifacts=["q3_report"]),
limit=10, # 最多取回 10 个命中
score_threshold=0.3, # 相关度低于 0.3 的丢掉
expand=True, # 打开树扩展
)
# 每个 chunk:命中的正文 + 前文列表 + 后文列表
chunks[0].text # "营收增长了 15%..."(命中的那块)
chunks[0].previous_texts # ["第三季度财务摘要", "本报告覆盖..."](向前补的上下文)
chunks[0].next_texts # ["科技板块贡献了 60%..."](向后补的上下文)
本节不碰代码细节。记住一句:命中一块 → 沿树补齐四周 → 拼回连续上下文。
2. 顶层全景(一次查询怎么转)
整个读取侧就一个入口:retrieve_semantic_relevant(private_gpt/server/primitives/semantic_search_service.py:94)。它把一次查询串成一条流水线。
怎么读下面这张图: 从上到下是一次查询的先后步骤;右侧标注了每一步落在哪个文件。虚线框「树扩展」是本章重点,单独在 §3.3 展开。
用户查询文本 "第三季度营收为什么增长?"
│
▼
┌────────────────────────────────────────────────┐
│ ① 建检索索引 + 嵌入查询 │ semantic_search_service.py:124-135
│ VectorStoreIndex.from_vector_store(...) │
└────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────────┐
│ ② 装配检索器(带过滤器) │ vector_store_component.py:111 get_retriever
│ similarity_top_k / score_threshold / filters │ vector_store_component.py:80 get_filters
│ 过滤条件来自 ContextFilter │ context_filter.py:6
└────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────────┐
│ ③ 检索 + 按 score 降序排 │ semantic_search_service.py:144-145
│ retriever.retrieve(text) → nodes.sort(...) │
└────────────────────────────────────────────────┘
│
▼
┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐
④ 树扩展 / 窗口回填(逐个命中做) │ semantic_search_service.py:46 _expand_nodes
│ 沿树向下/左右/向上补上下文,受 token 预算约束 │ document_expander.py:37 DocumentTreeExpander
└ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
│
▼
┌────────────────────────────────────────────────┐
│ ⑤ 转成结果模型 Chunk │ chunk/models.py:137 Chunk.from_node
│ 正文 + previous_texts + next_texts │
└────────────────────────────────────────────────┘
│
▼
list[Chunk] ← 返回给调用方
部件一句话职责:
| 部件 | 干什么 | 在哪 |
|---|---|---|
SemanticSearchService | 读取侧总入口,把整条流水线串起来 | server/primitives/semantic_search_service.py:27 |
VectorStoreComponent | 造检索器、把 ContextFilter 翻译成向量库的过滤条件 | components/vector_store/vector_store_component.py:28 |
ContextFilter | 声明「在哪搜」——限定 collection / artifacts / 元数据 | chat/extensions/context_filter.py:6 |
DocumentTreeExpander | 命中后沿树扩展上下文的算法,受 token 预算约束 | components/postprocessor/tree_expansion/document_expander.py:36 |
NodeStoreComponent | 按 id 取回节点(扩展时要重新拉出周围节点) | components/node_store/node_store_component.py |
Chunk | 结果模型:命中正文 + 前后上下文 + 元数据 | components/chunk/models.py:21 |
主线走一遍(高层): 查询文本进来 → 嵌入成向量、在指定 collection 上建一个只读检索索引(①)→ 用 ContextFilter 生成过滤器、装配检索器(②)→ 检索出 top-k 命中并按相关度排序(③)→ 对每个命中沿树扩展上下文(④) → 打包成 Chunk 返回(⑤)。
3. 核心原理(逐个机制,由浅入深)
3.1 主流程:retrieve_semantic_relevant
它要解决的小问题: 把「一段文本」变成「一组相关片段」,这中间有几步固定动作,得有个地方按顺序做完。
思路。 用 LlamaIndex 的现成件搭骨架:向量库 + 索引 + 检索器,PrivateGPT 只在过滤和扩展这两处塞进自己的逻辑。
真实实现(骨架,已省略中间行):
# semantic_search_service.py:124-157(节选)
index = VectorStoreIndex.from_vector_store( # 在该 collection 上建只读检索索引
self.vector_store_component.vector_store(collection),
embed_model=self.embedding_component.get_embed(embed_model_id), # 用它把查询嵌成向量
...
)
vector_index_retriever = self.vector_store_component.get_retriever( # 装配检索器 + 过滤器
index=index,
artifacts=context_filter.artifacts,
filter_dicts=context_filter.metadata_filter,
similarity_top_k=limit,
score_threshold=score_threshold,
)
nodes = vector_index_retriever.retrieve(text) # 真正检索
nodes.sort(key=lambda n: n.score or 0.0, reverse=True) # 按相关度降序
几个要点:
- 嵌入模型来自
EmbeddingComponent,查询和文档必须用同一个模型嵌入,向量才可比。 - 排序是显式的(
semantic_search_service.py:144),不完全依赖向量库返回的顺序;score为空时兜底成0.0。 - token 预算在这里定下来:
final_token_limit = token_limit or llm.metadata.context_window // 8(semantic_search_service.py:147)。调用方不指定,就取 LLM 上下文窗口的 1/8 作为「每个命中最多扩展多少 token」。这个数会一路传给树扩展算法(§3.3)。 expand开关:关掉就退化成朴素检索(只返回命中块,不补上下文);打开才走_expand_nodes(semantic_search_service.py:150-153)。
3.2 检索器与过滤:把检索限定到某些 artifact
它要解决的小问题: 用户往往不想在整个知识库里搜,而是「只在这份 Q3 报告里搜」或「只搜带某个标签的内容」。检索必须能被限定范围。
声明侧:ContextFilter。 这是个纯数据类(context_filter.py:6),调用方用它说清「在哪搜」,三层由粗到细:
| 字段 | 含义 | 定义 |
|---|---|---|
collection | 在哪个知识库里搜(最粗的范围) | context_filter.py:22 |
artifacts | 只搜这些 artifact(具体某些文件/文档) | context_filter.py:27 |
metadata_filter | 按元数据键值再筛(如 file_id == ...) | context_filter.py:32 |
类的 docstring 点明了组合语义:若同时给了 artifacts 和 metadata_filter,结果是两者的交集(context_filter.py:18-20)。
执行侧:get_filters → get_retriever。 VectorStoreComponent.get_filters(vector_store_component.py:83)把上面的声明翻译成向量库能懂的 MetadataFilters:
# vector_store_component.py:92-96(节选)——把每个 artifact 变成一条元数据过滤
if artifacts:
for artifact in artifacts:
filters.filters.append(
MetadataFilter(key=MetadataKeys.ARTIFACT_ID.value, value=artifact)
)
- artifact 过滤 = 按
ARTIFACT_ID这个元数据键匹配。摄取时(见 01 章)每个节点都盖了artifact_id,这里就靠它把检索圈到指定文档。 - 多个 artifact 用
OR连接(filters初始化时condition=FilterCondition.OR,vector_store_component.py:90-92):命中任一 artifact 都算。 - 逻辑多租户下还会再叠一层 collection 过滤,把上面的条件和租户条件用
AND组合(vector_store_component.py:101-110),确保跨租户不串数据。
get_retriever(vector_store_component.py:114)最后把过滤器、similarity_top_k(默认 2,但主流程传的是 limit)、score_threshold 一起塞进 LlamaIndex 的 VectorIndexRetriever。注意 score_threshold 是有值才塞(vector_store_component.py:125-127)——传 None 就不设阈值,不会误把所有结果卡成 0 分。
3.3 关键机制:树扩展 / 窗口回填(本章重点)
这是 PrivateGPT 检索最有含金量的一段。前面检索只给出一个个孤立的命中节点;这一步负责把每个命中回填成一段连续、够用的上下文。
为什么不用「固定 prev/next 窗口」
最朴素的做法是:命中第 5 块,就无脑带上第 4 块和第 6 块(固定窗口 ±1)。问题是文档是有结构的——第 4 块可能属于上一个章节、和命中内容毫无关系,而命中块真正需要的上下文其实是它的父节点(章节标题/引言)。固定窗口对结构一无所知,补进来的经常是噪声。
PrivateGPT 改成沿着摄取时建好的文档树扩展。命中一个叶节点后,它按「先向下、再左右、最后向上」的顺序,在 token 预算内尽量把树上相邻的、结构上真正相关的节点收进来。
服务侧:_expand_nodes 怎么调度
_expand_nodes(semantic_search_service.py:46)是服务层的胶水,干四件事:
- 找到命中节点在整棵树里的位置。 命中节点可能只是个「局部节点」,得先按
root_id把整棵树捞回来,再find_self_or_child_by_id定位到它(semantic_search_service.py:56-63;find_self_or_child_by_id见tree_node.py:439)。 - 跑扩展算法。 把定位到的节点和 token 预算交给
DocumentTreeExpander,fill_window()返回一组该收进来的节点 id(semantic_search_service.py:67-71)。 - 按顺序取回这些节点。 用
get_sorted_nodes把这些 id 对应的节点从库里拉回来(semantic_search_service.py:72-74)。 - 以命中节点为界,劈成前文/后文两半。 找到命中节点在列表里的下标,前面的进
previous_texts,后面的进next_texts(semantic_search_service.py:76-92)。
# semantic_search_service.py:76-82(节选)——以命中节点为轴,劈成前后两段
hit_node_idx = next(
(i for i, node in enumerate(sorted_nodes) if node.id_ == hit_node.id_), -1
)
prev_nodes = sorted_nodes[:hit_node_idx] # 命中之前 → previous_texts
next_nodes = sorted_nodes[hit_node_idx + 1 :] # 命中之后 → next_texts
算法侧:DocumentTreeExpander.fill_window 的三个 方向
fill_window(document_expander.py:87)把命中节点当起点,一次性朝三个方向递归扩展(_expand_recursively,document_expander.py:221):
父节点(向上 ↑)
┌──────┴──────┐
左兄弟(← 左右) [命中节点] 右兄弟(左右 →)
┌──┴──┐
子节点(向下 ↓)
三个方向各有分工(document_expander.py):
| 方向 | 方法 | 直觉 | 失败策略 |
|---|---|---|---|
| 向下 | _expand_downwards:293 | 先把命中节点自己的子树(更细的正文)收全 | STOP——保留已收的部分子树 |
| 左右 | _expand_horizontally:337 | 用「乒乓」策略在左右兄弟间交替扩,离得近、token 少、同类型的优先 | ROLLBACK——整棵子树一起丢,保证完整 |
| 向上 | _expand_upwards:424 | 爬到父节点(章节),再从父节点向下/左右铺开,补进「这块属于哪一节」 | ROLLBACK |
token 预算是硬约束。 每要收一个节点,先算它(或它整棵子树)的 token 数,超预算就按该方向的失败策略处理:
# document_expander.py:248-254(节选)——加节点前先看预算够不够
potential_tokens = self._get_token_count(current_node, include_subtree=False)
if potential_tokens > self._remaining_tokens:
return self._current_failure_result # 超了 → STOP 或 ROLLBACK
self._remaining_tokens -= potential_tokens # 没超 → 扣预算,收下
节点的 token 数来自摄取时算好的元数据(TreeNode.token_count,tree_node.py:82;整棵子树用 get_sum_token_count,tree_node.py:475)——读取时不用现算,直接查表,快。
两阶段的巧思。 向下用 STOP(预算不够时,已经收进来的半棵子树留着——多点细节无妨);左右和向上用 ROLLBACK(预算不够时,把这一整棵候选子树整个丢掉——宁可不要,也不要半截的兄弟/父级把上下文搞得支离破碎)。一句话:正文可以只补一半,但结构相邻的整块要么全要要么不要。