跳到主要内容

第 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

这套设计妙在哪:

  1. 从 visited 实体往外扩(search_relationships_weight 里 query.where(source_entity_id.in_(visited_entities)),tidb_graph_store.py:694)——保证扩出来的是和已命中节点结构相连的,而不是图里随便一个相似关系。这才叫「图」检索。
  2. 近的距离段拿大配额:越接近问题的关系越值钱,优先填满;远的只给一点点,防止跑偏。距离 > 0.55 直接不要。
  3. 排除 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 章收尾:把这些巧妙之处提炼成可借鉴清单,并给出边界、横向对比和代码地图。