数据截至 (上游 commit c4cb19d8fa20)
第 3 章 · 知识图谱检索:从一条关系出发,逐跳长出子图
这是全书技术含量最高的一章。它要解决的小问题:给一个问题,如何在「实体—关系」图里捞出一小撮最相关、又有结构关联的实体和关系? AutoFlow 的答案不是简单的向量 top-k,而是一套带权重、分距离段、多跳扩散的检索。
3.1 图谱长什么样
先建立心智模型。知识图谱里两类东西:
- 实体 Entity:一个概念/对象,带
name、description、description_vec(描述的向量)。 - 关系 Relationship:一条有向边
source_entity → description → target_entity,带weight(权重)、description_vec(关系描述的向量)。
关键设计:关系本身也是向量化的(get_relationship_description_embedding,graph_store/helpers.py:100,把「源实体(描述)→ 关系 → 目标实体(描述)」拼成一句话再 embed)。所以检索的主体是「关系」,不是「实体」——这是和很多 GraphRAG 的不同之处。
3.2 三层结构(和 chunk 检索对称)
KnowledgeGraphFusionRetriever (融合层:多 KB 子图合并)
└── KnowledgeGraphSimpleRetriever (执行层:单 KB 一次图检索)
└── TiDBGraphStore.retrieve_with_weight() (灵魂算法)
执行层 KnowledgeGraphSimpleRetriever._retrieve(retrievers/knowledge_graph/simple_retriever.py:55)非常薄,核心就一句:
# 真实源码节选,knowledge_graph/simple_retriever.py:60
entities, relationships = self._kg_store.retrieve_with_weight(
query_bundle.query_str,
embedding=[],
depth=self.config.depth, # 默认 2 跳
include_meta=self.config.include_meta,
with_degree=self.config.with_degree, # 是否用度数加分
relationship_meta_filters=metadata_filters,
)
所有难点都在 retrieve_with_weight 里。
3.3 retrieve_with_weight 的全景
这个函数(graph_store/tidb_graph_store.py:478)做的事,一张图说清:
query → embedding
│
▼
第 1 跳:search_relationships_weight(全距离范围)
找出与问题最相关的一批关系 + 它们两端的实体
│ 记下 visited(已访问的关系/实体 id)
▼
第 2..depth 跳:对每一跳,按「距离分段」逐段找新邻居
┌─ 距离段 (0.00,0.25) 比例 1
├─ 距离段 (0.25,0.35) 比例 0.7
├─ 距离段 (0.35,0.45) 比例 0.2
└─ 距离段 (0.45,0.55) 比例 0.1
每段从「已访问实体」出发,找未访问过的新关系,带权打分取 top
│
▼
补充:fetch_similar_entities(synopsis 概要实体 top 2)
│
▼
汇总所有 entities + relationships → 返回
怎么读: 第 1 跳是「锚点」——直接对问题向量找最近的关系。之后每一跳都从已访问实体往外长,而且越近的距离段配额越高(比例 1 > 0.7 > 0.2 > 0.1),保证优先吃掉高相关的邻居。
3.4 第 1 跳:从问题直接找最相关的关系
# 真实源码节选,tidb_graph_store.py:489
if not embedding:
embedding = get_query_embedding(query, self._embed_model)
relationships, entities = self.search_relationships_weight(
embedding, [], [], # 还没 visited 任何东西
with_degree=with_degree,
relationship_meta_filters=relationship_meta_filters,
)
all_relationships = set(relationships)
all_entities = set(entities)
visited_entities = set(e.id for e in entities)
visited_relationships = set(r.id for r in relationships)
search_relationships_weight(下节)负责「按问题向量找一批关系并打分」。第一跳不限制起点(visited 为空),所以它纯按与问题的语义相似度找种子关系。
3.5 多跳扩散:按「距离分段 + 配额」往外长
这是最精巧的部分(tidb_graph_store.py:507)。它不是无脑找所有邻居,而是把「每跳要找 10 个」这个配额,按距离远近分摊给不同的距离段:
# 真实源码节选,tidb_graph_store.py:511 (循环 depth-1 次)
search_number_each_depth = 10
for search_config in DEFAULT_RANGE_SEARCH_CONFIG: # 4 个距离段
search_ratio = search_config[1] # 该段配额比例
search_distance_range = search_config[0] # 该段距离范围
...
new_relationships, new_entities = self.search_relationships_weight(
embedding,
visited_relationships, # 排除已访问的关系
visited_entities, # 只从已访问实体往外扩
search_distance_range, # 限定这一段距离
rank_n=expected_number, # 这段要取几个
with_degree=with_degree, ...)
all_relationships.update(new_relationships)
visited_entities.update(e.id for e in new_entities)
visited_relationships.update(r.id for r in new_relationships)
配额表 DEFAULT_RANGE_SEARCH_CONFIG(graph_store/helpers.py:20):
| 距离范围(余弦距离,越小越像) | 配额比例 |
|---|---|
| (0.00, 0.25) | 1 |
| (0.25, 0.35) | 0.7 |
| (0.35, 0.45) | 0.2 |
| (0.45, 0.55) | 0.1 |
这套设计妙在哪:
- 从 visited 实体往外扩(
search_relationships_weight里query.where(source_entity_id.in_(visited_entities)),tidb_graph_store.py:694)——保证扩出来的是和已命中节点结构相连的,而不是图里随便一个相似关系。这才叫「图」检索。 - 近的距离段拿大配额:越接近问题的关系越值钱,优先填满;远的只给一点点,防止跑偏。距离 > 0.55 直接不要。
- 排除 visited:不重复访问,避免兜圈子。
3.6 关系打分公式:相似度 + 权重 + 度数
每一跳里,候选关系不止按距离排,而是按一个综合分排(search_relationships_weight → calculate_relationship_score)。公式在 graph_store/helpers.py:51:
# 真实源码节选,helpers.py:51
def calculate_relationship_score(embedding_distance, weight, in_degree, out_degree,
alpha, weight_coefficient_config, degree_coefficient,
with_degree=False):
weighted_score = get_weight_score(weight, weight_coefficient_config)
degree_score = 0
if with_degree:
degree_score = get_degree_score(in_degree, out_degree, degree_coefficient)
return alpha * (1 / embedding_distance) + weighted_score + degree_score
三个加项,各管一件事:
| 加项 | 含义 | 直觉 |
|---|---|---|
alpha * (1 / embedding_distance) | 语义相似度 | 离问题越近,分越高(距离倒数) |
weighted_score | 关系的累积权重 | 被反复印证的关系更可信 |
degree_score(可选) | 入度 − 出度 | 偏好「被指向多」的枢纽实体 |
权重分是分段累进的(get_weight_score,helpers.py:30),系数随权重变大而递减(DEFAULT_WEIGHT_COEFFICIENT_CONFIG,helpers.py:9):前 100 的权重每点算 0.01,100~1000 每点 0.001…… 这样高权重边不会因为权重爆表而碾压一切——边际收益递减,防止单条关系霸榜。
度数分(get_degree_score,helpers.py:47)= (in_degree - out_degree) * 系数:入度高、出度低的实体(被很多东西指向、自己很少指出去)更像一个核心概念,加分。
3.7 SQL 层:为什么先取 limit×10 再排
search_relationships_weight(tidb_graph_store.py:632)的查询结构值得看一眼:它先用 TiDB 向量索引取一批候选,再在 Python 里用综合分重排:
# 真实源码节选,tidb_graph_store.py:650
subquery = (select(self._relationship_model,
self._relationship_model.description_vec.cosine_distance(embedding).label("embedding_distance"))
.order_by(asc("embedding_distance"))
.limit(limit * 10) # 先用向量索引多取 10 倍候选
).subquery()
# ... 加上 visited 过滤、距离段过滤、起点实体过滤 ...
query = query.order_by(asc("embedding_distance")).limit(limit)
和第 2 章 chunk 的 oversampling 是同一个思路:ANN 先粗筛多取,再精筛。粗筛只能按单一的 embedding_distance 排(数据库会的),综合打分(加 weight、degree)在应用层做(tidb_graph_store.py:721-746)。
还有个小优化:候选数 ≤ rank_n 时直接返回,跳过打分(tidb_graph_store.py:702)——没必要为几条候选算分排序。
3.8 补充:synopsis 概要实体
扩散结束后,额外捞 2 个 synopsis(概要)实体(tidb_graph_store.py:551)。synopsis 是建库时对一组实体生成的「摘要节点」(EntityType.synopsis),提供更宏观的上下文,补足逐跳扩散偏「局部」的不足。
3.9 融合层:多 KB 子图怎么合
跨多个知识库时,KnowledgeGraphFusionRetriever._knowledge_graph_fusion(knowledge_graph/fusion_retriever.py:97)把各 KB 的子图合成一个:
# 真实源码节选,knowledge_graph/fusion_retriever.py:113
merged_entities.update(node.entities) # 实体并集
for r in node.relationships:
key = r.rag_description # "A -> 关系 -> B" 作 key
if key not in merged_relationships:
merged_relationships[key] = r
else:
merged_relationships[key].weight += r.weight # 同名关系权重相加
合并策略: 实体取并集;关系按 rag_description(人类可读的「源→关系→目标」)去重,同一条关系在多个库都出现就把权重相加——跨库印证的关系更可信。和 chunk 融合的「max 取 分」不同,这里是「权重累加」,因为关系权重本身就是「被印证次数」的语义。
3.10 小结
- 检索主体是关系(关系被整体向量化),实体跟着关系两端被带出来。
- 核心算法
retrieve_with_weight:第 1 跳锚定 → 按距离分段配额多跳扩散 → 综合分(相似度+权重+度数)排序。 - 处处是「ANN 粗筛多取 + 应用层精排」的两段式,和 chunk 检索同构。
- 多 KB 融合按关系
rag_description去重、权重累加。
第 4 章收尾:把这些巧妙之处提炼成可借鉴清单,并给出边界、横向对比和代码地图。