语义检索与树扩展:一次查询如何取回带上下文的片段

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:37`
`NodeStoreComponent`	按 id 取回节点(扩展时要重新拉出周围节点)	`components/node_store/node_store_component.py`
`Chunk`	结果模型:命中正文 + 前后上下文 + 元数据	`components/chunk/models.py:18`

主线走一遍(高层): 查询文本进来 → 嵌入成向量、在指定 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:145),不完全依赖向量库返回的顺序;score 为空时兜底成 0.0。
token 预算在这里定下来:final_token_limit = token_limit or llm.metadata.context_window // 8(semantic_search_service.py:148)。调用方不指定,就取 LLM 上下文窗口的 1/8 作为「每个命中最多扩展多少 token」。这个数会一路传给树扩展算法(§3.3)。
expand 开关:关掉就退化成朴素检索(只返回命中块,不补上下文);打开才走 _expand_nodes(semantic_search_service.py:151-154)。

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:80)把上面的声明翻译成向量库能懂的 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:87-89):命中任一 artifact 都算。
逻辑多租户下还会再叠一层 collection 过滤,把上面的条件和租户条件用 AND 组合(vector_store_component.py:98-107),确保跨租户不串数据。

get_retriever(vector_store_component.py:111)最后把过滤器、similarity_top_k(默认 2,但主流程传的是 limit)、score_threshold 一起塞进 LlamaIndex 的 VectorIndexRetriever。注意 score_threshold 是有值才塞(vector_store_component.py:122-124)——传 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:437)。
跑扩展算法。 把定位到的节点和 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:88)把命中节点当起点,一次性朝三个方向递归扩展(_expand_recursively,document_expander.py:222):

                        父节点(向上 ↑)
                     ┌──────┴──────┐
     左兄弟(← 左右)  [命中节点]   右兄弟(左右 →)
                     ┌──┴──┐
                  子节点(向下 ↓)

三个方向各有分工(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:473)——读取时不用现算,直接查表,快。

两阶段的巧思。 向下用 STOP(预算不够时,已经收进来的半棵子树留着——多点细节无妨);左右和向上用 ROLLBACK(预算不够时,把这一整棵候选子树整个丢掉——宁可不要,也不要半截的兄弟/父级把上下文搞得支离破碎)。一句话:正文可以只补一半,但结构相邻的整块要么全要要么不要。

一段示意代码:命中 → 按预算向四周扩展 → 拼回连续上下文

下面用极简的树把整个思路演一遍(真实实现见上面各 file:line):

# 示意,非源码:命中一个叶节点,按 token 预算沿树补齐四周
def expand_around_hit(hit_node, token_budget):
    picked = {hit_node.id}                 # 结果集:先放命中节点自己
    remaining = token_budget - hit_node.tokens

    # 一个通用的“想收一个节点就先问预算”的动作
    def try_take(node):
        nonlocal remaining
        if node.id in picked or node.tokens > remaining:
            return False                   # 已收过 / 预算不够 → 放弃
        picked.add(node.id)
        remaining -= node.tokens
        return True

    # ① 向下:命中块自己的子树(更细的正文)优先补全
    for child in hit_node.children:
        try_take(child)

    # ② 左右:在左右兄弟间“乒乓”交替,离得近的先收
    left = hit_node.left_siblings()[::-1]   # 越近越先
    right = hit_node.right_siblings()
    for l, r in zip_longest(left, right):
        if r: try_take(r)
        if l: try_take(l)

    # ③ 向上:爬到父节点(章节),补进“这块属于哪一节”
    if hit_node.parent:
        try_take(hit_node.parent)

    # 收尾:把选中的节点按文档原始顺序排好,以命中为界劈成前/后文
    ordered = sorted(picked, key=lambda nid: abs_index_of(nid))  # 按 abs_idx
    i = ordered.index(hit_node.id)
    return ordered[:i], ordered[i + 1:]     # (previous_texts, next_texts)

重点看: 每收一个节点都先过 try_take 的预算关(对应真实代码 document_expander.py:248-254);最后按文档顺序排好、以命中为轴劈成前后两段(对应 semantic_search_service.py:76-92)。这就是「比固定 prev/next 窗口更聪明」的地方——补进来的是结构上真正相邻的内容,而不是机械的前一块后一块。

3.4 结果模型:`Chunk.from_node`

它要解决的小问题: 检索内部用的是 LlamaIndex 的 NodeWithScore,但对外得给一个稳定、带上下文字段的结果对象。

Chunk(chunk/models.py:18)就是这个对外模型,关键字段:

字段	含义	定义
`text`	命中块的正文	`chunk/models.py:39`
`score`	相关度(0~1,越高越相关)	`chunk/models.py:29`
`previous_texts`	树扩展补出的前文列表	`chunk/models.py:52`
`next_texts`	树扩展补出的后文列表	`chunk/models.py:61`
`document`	命中所属文档(含 `artifact` 等元数据)	`chunk/models.py:36`

Chunk.from_node(chunk/models.py:137)负责从 NodeWithScore 造出 Chunk:抽取正文(树节点用 get_content(TreeMetadataMode.USER),即面向用户展示的模式)、把 score 兜底为非负、从元数据里挑出 abs_idx/rel_idx 等定位信息。注意 previous_texts/next_texts 在 from_node 里默认取自节点元数据(chunk/models.py:170-171),而树扩展的结果是在 retrieve_semantic_relevant 里回填覆盖上去的(semantic_search_service.py:152-153)——也就是先造壳、再塞扩展好的上下文。

最后,SourceType = Chunk | Website(chunk/models.py:271):检索结果和网页搜索结果共用一个「来源」类型,方便下游(03/05 章)统一处理。

4. 巧妙之处(可借鉴的技术)

用文档树替代固定窗口做上下文回填。 命中一块后补的是「结构上相邻」的内容(子/兄弟/父),而非机械的前一块后一块——补进来的更可能真相关。见 document_expander.py:37 + semantic_search_service.py:46。
两阶段失败策略(STOP vs ROLLBACK)。 向下允许保留半棵子树(细节多点无害),左右/向上要么整块收要么整块丢(避免半截结构污染上下文)。见 fill_window 的注释与 document_expander.py:88-113。
token 数在摄取时算好、读取时查表。 扩展要反复问「加这个节点会不会超预算」,靠 TreeNode.token_count/get_sum_token_count(tree_node.py:82/473)读现成的数,不在检索热路径上现算。
预算 = 上下文窗口的 1/8。 context_window // 8(semantic_search_service.py:148)是个克制的默认值:给每个命中留够上下文,又不至于让若干命中把整个 prompt 撑爆(多个 Chunk 会一起进 prompt,见 05 章)。
过滤声明与执行分离。 ContextFilter(数据)只说「在哪搜」,get_filters(逻辑)负责翻译成向量库过滤、并按需叠加多租户条件——调用方不必懂向量库的过滤语法。

5. 边界与局限(诚实)

扩展依赖树结构存在。 _expand_nodes 一上来就判断命中是不是 TreeNode,不是就直接返回空的前后文(semantic_search_service.py:52-54)——非树来源(如某些网页片段)拿不到树扩展的好处。
要重新回库取节点。 扩展阶段用 get_sorted_nodes 按 id 重新拉回周围节点(semantic_search_service.py:72),命中越多、扩展越广,回库读取的开销越大;expand=False 可关掉换取速度。
artifact 过滤会做就绪校验。 若传了 artifacts 且 validate=True,会逐个检查对应索引是否已就绪,没就绪直接报错(semantic_search_service.py:111-122)——好处是不会静默漏搜,代价是对未摄取完的文档不宽容。
score_threshold 语义依赖底层向量库。 阈值只是透传给 VectorIndexRetriever 的 vector_store_kwargs(vector_store_component.py:122-124),不同向量库对「分数」的定义/量纲未必一致,阈值需按后端调。

6. 代码地图(导航索引)

主题	文件	符号
读取侧总入口 / 主流程	`private_gpt/server/primitives/semantic_search_service.py`	`SemanticSearchService.retrieve_semantic_relevant`
扩展调度(定位命中→跑算法→劈前后文)	`private_gpt/server/primitives/semantic_search_service.py`	`SemanticSearchService._expand_nodes`
检索器装配	`private_gpt/components/vector_store/vector_store_component.py`	`VectorStoreComponent.get_retriever`
过滤翻译(artifact / 多租户)	`private_gpt/components/vector_store/vector_store_component.py`	`VectorStoreComponent.get_filters`
检索范围声明	`private_gpt/chat/extensions/context_filter.py`	`ContextFilter`
树扩展算法	`private_gpt/components/postprocessor/tree_expansion/document_expander.py`	`DocumentTreeExpander.fill_window`
三方向扩展	`private_gpt/components/postprocessor/tree_expansion/document_expander.py`	`_expand_downwards` / `_expand_horizontally` / `_expand_upwards`
按 id 取回节点	`private_gpt/components/node_store/node_store_component.py`	`NodeStoreComponent.get_sorted_nodes`
节点 token 数	`private_gpt/components/readers/nodes/tree_node.py`	`TreeNode.token_count` / `get_sum_token_count`
结果模型	`private_gpt/components/chunk/models.py`	`Chunk.from_node` / `SourceType`

相关章节: 片段从哪来 → 01-ingestion-tree-nodes;片段怎么锚回原文 → 03-citations;检索怎么被封成工具挂进 agent → 04-agentic-rag-chat-loop;片段怎么进 prompt → 05-context-stack-assembly;全局总览 → index。

1. 这是什么(零基础也能懂)​

2. 顶层全景(一次查询怎么转)​

3. 核心原理(逐个机制,由浅入深)​

3.1 主流程:retrieve_semantic_relevant​

3.2 检索器与过滤:把检索限定到某些 artifact​

3.3 关键机制:树扩展 / 窗口回填(本章重点)​

为什么不用「固定 prev/next 窗口」​

服务侧:_expand_nodes 怎么调度​

算法侧:DocumentTreeExpander.fill_window 的三个方向​

一段示意代码:命中 → 按预算向四周扩展 → 拼回连续上下文​

3.4 结果模型:Chunk.from_node​

4. 巧妙之处(可借鉴的技术)​

5. 边界与局限(诚实)​

6. 代码地图(导航索引)​