local-search – 码农练功房

引言

GraphRAG 的 Local Search 在查询时需要从知识图谱关联的原始文本片段（Text Units）中选择最相关的内容填入 LLM 上下文窗口。这个选择策略看似简单——按实体相似度排序、逐条填入——但在实际场景中会暴露出一个显著的局限性：热门实体可能霸占整个 Text Unit 预算，导致其他实体的关键文本被截断。

本文将深入分析这个问题的成因、它所解决的核心问题，以及可能的改进方向。

当前策略是什么

Local Search 的 Text Unit 选择分四步：

遍历选中实体（按向量相似度排名），收集每个实体关联的 text_unit_ids
去重：同一个 TU 只归属于最先遍历到的实体
排序：按 (entity_index, -num_relationships) —— 实体顺序优先，同一实体内按关系密度降序
逐条填入上下文，直到达到 token 上限（默认总预算的 50%，约 6000 tokens）

核心代码：

for index, entity in enumerate(selected_entities):
    entity_relationships = [rel for rel in relationships if rel.source == entity.title or rel.target == entity.title]
    for text_id in entity.text_unit_ids or []:
        if text_id not in text_unit_ids_set and text_id in self.text_units:
            num_relationships = count_relationships(entity_relationships, self.text_units[text_id])
            text_unit_ids_set.add(text_id)
            unit_info_list.append((self.text_units[text_id], index, num_relationships))

unit_info_list.sort(key=lambda x: (x[1], -x[2]))

问题场景：热门实体霸占预算

具体例子

假设用户问：”甘菊蓝的抗炎机制是什么？”

向量搜索返回的实体排名：

排名	实体	关联 TU 数量	说明
0	洋甘菊	50	高频实体，几乎所有草药文档都提到它
1	甘菊蓝	4	洋甘菊的活性成分，专业文献较少
2	NF-κB 通路	2	具体的抗炎分子机制

遍历去重后的 TU 归属：

index 0 "洋甘菊": TU1, TU2, TU3, ..., TU50  (50条)
index 1 "甘菊蓝": TU51, TU52              (TU1, TU5 已被洋甘菊占有)
index 2 "NF-κB":  TU53                    (仅1条未被占有)

排序结果：

TU1(index=0, rel=5) → TU2(index=0, rel=4) → ... → TU50(index=0, rel=0)
→ TU51(index=1, rel=2) → TU52(index=1, rel=1)
→ TU53(index=2, rel=1)

假设 token 预算为 6000 tokens，每个 TU 平均 300 tokens，那么只能容纳约 20 个 TU。

结果：前 20 个位置全部被”洋甘菊”的 TU 占据。用户真正关心的”甘菊蓝的抗炎机制”相关文本（TU51、TU52、TU53）全部被截断。LLM 拿到的上下文中充斥着”洋甘菊”的泛泛介绍，却没有任何关于甘菊蓝具体分子机制的原文支撑。

为什么要这样设计：它解决了什么问题

这个策略并非随意设计，它解决的是一个更基础的问题：确保语义最相关的实体获得最充分的原文支撑。

它解决的场景

假设用户问：”洋甘菊在欧洲传统医学中的地位如何？”

向量搜索返回：

排名	实体	关联 TU 数量
0	洋甘菊	50
1	欧洲草药学	8
2	薰衣草	30

在这个场景中，”洋甘菊”确实是最核心的实体，用户就是在问它。如果采用 round-robin 策略（每个实体轮流取 1 个 TU），那么”薰衣草”的 30 个 TU 会与”洋甘菊”平分预算——但用户根本没问薰衣草。

当前策略的优势在于：
– 尊重语义排名：向量相似度最高的实体获得最多原文支撑，这在大多数情况下是正确的
– 关系密度排序保证质量：同一实体的多个 TU 中，信息最密集的排在前面
– 去重避免冗余：同一个 TU 不会因为被多个实体关联而重复出现

核心权衡

这是一个经典的 相关性深度 vs. 覆盖广度 的权衡：

当前策略选择了深度：确保最相关实体有充分的原文证据
代价是广度：次要实体可能完全没有原文支撑

在大多数”关于特定实体的问题”中（Local Search 的设计目标），深度优先是合理的。问题出在查询涉及多个实体的交叉关系时。

问题的本质：单一排序维度无法表达多目标优化

Text Unit 选择本质上是一个多目标优化问题：

相关性：TU 与查询的语义相关度（通过实体排名间接表达）
信息密度：TU 中包含的关系数量
覆盖性：确保每个选中实体都有原文支撑
多样性：避免上下文中充斥同质化内容

当前策略用一个二元组 (entity_index, -num_relationships) 试图同时优化前两个目标，但完全忽略了后两个。

改进方向

方案 1：Per-Entity Cap（每实体上限）

最简单的改进——为每个实体设置 TU 贡献上限：

MAX_TU_PER_ENTITY = 5

for index, entity in enumerate(selected_entities):
    count = 0
    for text_id in entity.text_unit_ids or []:
        if count >= MAX_TU_PER_ENTITY:
            break
        if text_id not in text_unit_ids_set and text_id in self.text_units:
            # ... 添加逻辑不变
            count += 1

优点：实现简单，保证每个实体至少有机会贡献 TU
缺点：cap 值难以确定；如果某个实体确实需要大量原文支撑，会被人为限制

方案 2：Round-Robin（轮询）

每轮从每个实体取 1 个 TU（按关系密度选最优的），循环直到预算用完：

entity_queues = {i: sorted_tus_for_entity_i for i in range(len(selected_entities))}
result = []
while budget > 0 and any(entity_queues.values()):
    for i in range(len(selected_entities)):
        if entity_queues[i]:
            tu = entity_queues[i].pop(0)
            result.append(tu)
            budget -= token_count(tu)

优点：保证覆盖性，每个实体都有原文支撑
缺点：最相关实体的深度被稀释；排名靠后的不相关实体也获得了等量预算

方案 3：加权配额分配

按实体的向量相似度分数分配 TU 配额：

# 假设相似度分数: [0.95, 0.82, 0.71]
scores = [0.95, 0.82, 0.71]
total = sum(scores)
quotas = [int(max_tus * s / total) for s in scores]
# quotas ≈ [15, 13, 11] (假设 max_tus=39)

优点：兼顾深度和广度，相关性高的实体获得更多配额但不会独占
缺点：实现复杂度增加；需要从向量搜索结果中保留相似度分数（当前代码未保留）

方案 4：最小保证 + 剩余竞争

每个实体保证至少 N 个 TU（如 2 个），剩余预算按当前策略竞争：

# Phase 1: 每个实体保证 2 个最优 TU
for entity in selected_entities:
    guaranteed_tus = top_2_by_relationship_density(entity)
    result.extend(guaranteed_tus)

# Phase 2: 剩余预算按原策略排序填充
remaining = all_tus - guaranteed_tus
remaining.sort(key=lambda x: (x.entity_index, -x.num_relationships))
fill_until_budget(remaining)

优点：保证覆盖性的同时保留了原策略的深度优势
缺点：如果选中实体很多，保证阶段可能消耗大量预算

总结

维度	当前策略	问题
相关性深度	✅ 优秀	—
信息密度	✅ 优秀	—
覆盖广度	❌ 缺失	热门实体霸占预算
内容多样性	❌ 缺失	同质化风险

GraphRAG 当前的 Text Unit 选择策略是一个”深度优先”的设计，它在”关于单一实体的问题”场景下表现良好，但在涉及多实体交叉关系的查询中会暴露覆盖性不足的问题。

最务实的改进是方案 4（最小保证 + 剩余竞争）——它以最小的代码改动保证了每个选中实体至少有原文支撑，同时不破坏原策略在主流场景下的优势。

标签： local-search

GraphRAG Local Search 的 Text Unit 选择策略：设计权衡与改进方向

引言