[GraphRAG] [Query]

448 阅读25分钟

[GraphRAG] [Query]

Local Search

Local Search 将来自 知识图谱 中的 结构化数据 与来自 输入文档非结构化数据 结合起来,在查询时用相关实体信息增强 LLM 上下文。

这种方法特别适合于回答那些 要求理解输入文档中提及的具体实体问题(比如“洋甘菊有什么治愈性质?”。

flowchart LR
    AA[User Query] --> |Entity Description Embedding|DD[Extracted Entities]
    CH[Conversation History] --> |Entity Description Embedding|DD[Extracted Entities]
    
    subgraph ConnectedEntities[Connected Entities]
        direction TB
        B[Candidate Text Units]
         E[Candidate Community Reports]
       H[Candidate Entities]
        K[Candidate Relationships]
        N[Candidate Covariates] 
    end
    
    subgraph PrioritizedEntities[Prioritized Entities]
        direction TB
        C[Prioritized Text Units]
        F[Prioritized Community Reports]
        I[Prioritized Entities]
        L[Prioritized Relationships]
        O[Prioritized Covariates]
    end
    
    DD -->|Entity-Text Unit Mapping|B[Candidate Text Units]
    DD -->|Entity-Report Mapping| E[Candidate Community Reports]
    DD -->|Entity-Entity Relationships|H[Candidate Entities]
    DD -->|Entity-Entity Relationships|K[Candidate Relationships]
    DD --> |Entity-Covariate Mappings| N[Candidate Covariates] 
   
   CH  -->  Q[Conversation History]
   
B[Candidate Text Units]-->|Ranking + Filtering|C[Prioritized Text Units]
E[Candidate Community Reports]-->|Ranking + Filtering|F[Prioritized Community Reports]
H[Candidate Entities]-->|Ranking + Filtering|I[Prioritized Entities]
K[Candidate Relationships]-->|Ranking + Filtering|L[Prioritized Relationships]
N[Candidate Covariates]-->|Ranking + Filtering|O[Prioritized Covariates]
 
PrioritizedEntities --> S[Response]
Q[Conversation History]--> S[Response]

给定用户的 查询 以及(如果有的话)会话历史记录

  • local search 会在知识图谱上找到一组与用户输入相关的 实体

    • 抽取更多该实体的相关信息:关联的实体关系实体协变量社群报告

    • 识别出的实体有关联的原始输入文件中抽取出相关的 文本片段

  • 对这些候选项数据来源进行 优先级排序和过滤,并且使其适应预设的单个上下文窗口。

数据准备

resolve_parquet_file

读取 parquet_list 得到:

  • final_nodes:create_final_nodes.parquet

    {
      "id": "node_id(uuid)",
      "human_readable_id": "node_id",
      "title": "node_name(非唯一)",
      "community": "node所属社群",
      "level": "node所属社群的层级",
      "degree": "与该node相连的node个数"
    }
    
  • final_entities:create_final_entities.parquet

    {
      "id": "node_name(uuid)",
      "human_readable_id": "node_id",
      "title": "node_name(唯一)",
      "type": "node_type(person, event, geo,...)",
      "description": "node描述",
      "text_unit_ids": "List[chucked_text]"
    }
    
  • final_communities:create_final_communities.parquet

    {
      "id": "community_id(uuid)",
      "human_readable_id": "community_id",
      "community": "community_name",
      "level": "社群层级",
      "title": "Community+community_name",
      "entity_ids": "社群中包含的node_id(对应final_entities中的id)",
      "relationship_ids": null,
      "text_unit_ids":  "社群中包含的node_id所属的 chucked_text",
      "period": "更新日期",
      "size": "社群node个数"
    }
    
  • final_community_reports:create_final_community_reports.parquet

    'create_final_relationships'{
      "id": "community_report_id(uuid)",
      "human_readable_id": "community_report_id",
      "community": "community_name",
      "level":  "社群层级",
      "title": "LLM 总结的community_name",
      "summary": "LLM 总结的summary",
      "full_content": "full_content",
      "rank": "影响力评分",
      "rank_explanation": "影响力评分的原因",
      "findings": "findings",
      "full_content_json": "full_content_json",
      "period": "更新日期",
      "size": "社群node个数"
    }
    
  • final_relationships:create_final_relationships.parquet

    'create_final_relationships'{
      "id": "relationship_id(uuid)",
      "human_readable_id": "relationship_id",
      "source": "source_node_name",
      "target": "target_node_name",
      "description": "relationship 描述",
      "weight": "relationship 紧密程度""combined_degree""source_degree+target_degree"
      "text_unit_ids""社群中包含的node_id所属的 chucked_text",
    }
    
read_indexer_entities

final_nodesfinal_entities中读取实体信息,并根据社群层级进行筛选和处理

  • 根据 community_level 筛选社群报告和节点数据

  • 处理final_nodes数据

    • 提取 final_nodes 中的 iddegreecommunity
    • 将所有节点的 community 列缺失值填充为 -1,并转换为整数类型,将 degree 列转换为整数类型。
    • 将所有节点按 iddegree 分组,聚合 community 列,去重,集合(set)转换为字符串列表。
  • 合并final_nodesfinal_entities数据为final_df

    • 使用 id 列将处理后的节点数据 nodes_df 和实体数据 entities_df 合并,得到一个包含实体信息的 DataFrame,去重。
  • 最终数据处理

    final_df的每一行转化为Entity 对象得到一个列表 _entities

    Entity(
        id = id,
        short_id = human_readable_id,
        title = title,
        type = type,
        description = description, 
        description_embedding = description_embedding,
        name_embedding = name_embedding, 
        community_ids = community_ids,
        text_unit_ids = text_unit_ids,
        rank = rank, 
        attributes = None
    )
    
read_indexer_reports

final_community_reports中读取社群报告,并根据社群层级和配置选项进行筛选、嵌入、合并等处理,最终返回处理后的社群报告列表。

  • 根据 community_level 筛选社群报告和节点数据

  • 社群报告的合并(基于最大社群层级)

    • 如果 dynamic_community_selectionFalse(即不启用动态社群选择),函数会根据每个节点的社群层级,确定该节点所在的最大社群层级。
    • nodes_df 中的每个节点分配其所属的最大社群,并筛选出这些社群 filtered_community_df
      • 将所有节点 community列缺失值填充为 -1,并转换为整数类型
      • 根据每个节点的 title 聚合,并选择每个 title 的最大社群层级。
    • 根据filtered_community_df(经过层级筛选后的社群列表)筛选出reports_df 中的社群报告
  • 报告内容嵌入

    • 如果配置对象 config 存在且报告数据中未包含内容嵌入列,或者该列存在但包含缺失值,则使用 get_text_embedder 获取一个嵌入模型(例如,OpenAI 的嵌入模型),并为报告内容生成嵌入向量,存储在 content_embedding_col 列中
  • 最终数据处理

    reports_df的每一行转化为CommunityReport 对象得到一个列表 reports

    CommunityReport(
        id = id,
        short_id = community,
        title = title,
        community_id = community,
        summary = summary,
        full_content = full_content,
        rank = rank,
        summary_embedding = summary_embedding,
        full_content_embedding = full_content_embedding,
        attributes = None,
        size = size,
        period = period
    )
    
read_indexer_text_units
TextUnit(
    id = id,
    short_id = index,
    text = text
    text_embedding = text_embedding,
    entity_ids = entity_ids,
    relationship_ids = relationship_ids,
    covariate_ids = covariate_ids,
    n_tokens = 1200,
    document_ids = document_ids,
    attributes = None
)
read_indexer_relationships
Relationship(
    id ='ce84c7567aef4778b493efa54fbda3bc',
    short_id ='0',
    source ='PROJECT GUTENBERG',
    target ='SUZANNE SHELL',
    weight = 7.0,
    description = "Suzanne Shell produced the Project Gutenberg eBook of 'A Christmas Carol'",
    description_embedding = None,
    text_unit_ids = ['d6583840046247f428a9f02738842a7c'],
    rank = 7,
    attributes = None
)
向量存储库配置
  • 设置向量存储类型:
    • 从配置中获取向量存储类型(vector_store_type),并根据存储类型(如 LanceDB)配置相应的参数。特别是对于 LanceDB 类型,更新数据库路径。
  • 连接embedding 存储
  • 获取本地搜索引擎

向量数据库:

  • description_embedding_store: entity.description 实体数据库

构建上下文

map_query_to_entities

根据query匹配 entities

  • 匹配实体

    • 如果query不为空:

      • 将query转为embedding

      • text_embedding_vectorstore 中进行基于embedding的相似度搜索,返回与查询最相关的embedding。

        k * oversample_scaler 表示放大返回的搜索结果数量。

      • 遍历搜索结果,根据 embedding_vectorstore_key 提取实体。根据嵌入向量的 ID 匹配相应的实体。

    • 如果query为空,则按实体的 rank 字段对所有实体进行排序,选择前 k 个实体。

  • 过滤与排序

    • 排除实体:根据 exclude_entity_names 中指定的名称过滤掉不需要的实体。
    • 包含实体:根据 include_entity_names 中的名称强制包括某些实体。
  • 最终返回 included_entities + matched_entities的列表。

build_community_context

如果 selected_entitiesself.community_reports 为空,直接返回空上下文文本和空的 DataFrame。

社群匹配
  • 对每个选中的实体,其所属的社群 match 数量 +1。
  • 在社群报告筛选中出这些社群,并按照 match 数量排序,如果有多个社群匹配的实体数相同,还会按社群 rank 进行二次排序,得到selected_communities
构建社群上下文
  1. 计算社群权重

    如果 entities 参数不为空,并且需要计算社群权重(即 include_community_weight=True),则为每个社群报告计算权重。计算的权重是与该社群相关联的文本单元的数量。

    • 建立社群和文本单元的关联

      • community_text_unitsDict{community_id: List[chuncked_doc_id, ]}存储每个社群 ID 和与之相关联的 text_unit_id 列表。

        对于每个实体,遍历其所属的社群 ID,并将该实体的文本单元 ID 添加到对应社群的列表中。

    • 为社群报告计算权重

      • 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的text_unit_id数量。
      • 将该数量存储为社群报告的 attributes: 中 ({"occurrence weight": int})。
    • 权重归一化(可选):

      如果 normalize==True,则会对所有社群报告的权重进行归一化。

      归一化的方式是将每个社群的权重除以所有社群中最大权重的值,使得最大权重为 11,其他社群的权重相对最大权重进行调整

  2. 选择报告并打乱报告顺序

    • 选择社群层级大于等于 min_community_rank 的报告。

    • 如果 shuffle_data==True,则会打乱报告的顺序,以避免模型偏向于某些社群。

  3. 生成报告上下文文本

    • 将社群报告按批次处理,每个批次的 token 数量不能超过 max_tokens

      • 当一个批次的 tokens 数量没有超过限制时,直接拼接
      • 当一个批次的 tokens 数量超过限制时,会将当前批次的文本转化为 DataFrame 并保存,然后开始一个新的批次。
    • 生成表头

      对每个报告,提取 header = ['id', 'title', 'occurrence weight', 'content', 'rank']

      batch_text = (
          f "-----{context_name}-----" + "\n" + column_delimiter.join(header) + "\n"
      )
      
      -----Reports-----
      id|title|occurrence weight|content|rank
      
    • 上下文数据

      • 对每个报告:
        • 提取 short_idtitleattributes.occurrence_weightsummary(如果没有则使用full_content)、rank组成 listnew_context
        • column_delimiter|)拼接new_context得到 new_context_text
        • 计算 new_context_texttokens 数量
  4. 返回结果

    • community_context: List[str]: nn 个 batch 的上下文文本

    • community_context_data: Dict{"reports": pd.Dataframe},Dataframe 里记录每个报告的详细信息

      idtitleoccurence weightcontentrank

  • 如果 return_candidate_contextTrue,还会调用 get_candidate_communities 来获取候选社群的数据,并在上下文中标注哪些社群已经包含在当前上下文中。
build_local_context

final_context_text

-----Entities-----
id|entity|description|number of relationships

-----Relationships-----
id|source|target|description|weight|links

-----Claims-----
构建entities上下文

构建社群上下文

返回结果

-----Entities-----
id|entity|description|number of relationships
identitydescriptionnumber of relationships
  • current_context_text: List[str]: nn 个 batch 的上下文文本

  • record_df: Dict{"reports": pd.Dataframe},Dataframe 里记录每个实体的详细信息

构建relationships上下文:
  • 筛选关系

    筛选出与所选实体相关的relationship,并限制返回最多 relationship_budget 个关系。

    • 筛选in_network_relationships

      • source 和 target 均在 selected_entity
    • 筛选 out_network_relationships

      • source 和 target 任一selected_entity
      • 按照 ranking_attribute(rank) 排序
    • 计算 out_network_entities, 表示 out_network_relationships 中,不在selected_entity的部分。

    • 为每个out_network_entity计算链接数:(即它们作为 sourcetargetout_network_relationships中出现的次数)。对 out_network_relationships 按照链接数排序(链接数相同的情况下,再按 rankweight 排序)。

    • 计算 relationship_budget

      relationship_budget = top_k_relationships * len(selected_entities)
      
    • 返回in_network_relationships + out_network_relationships 的前relationship_budget 个关系。

  • 构建relationship上下文构建社群上下文

    返回结果

    -----Relationships-----
    id|source|target|description|weight|links
    
    idsourcetargetdescriptionweightlinks

  • 协变量上下文:对于每个协变量,函数会使用 build_covariates_context 来生成上下文。
构建协变量上下文
build_text_unit_context

如果没有选中的实体或没有文本单元数据,则直接返回空的上下文文本和空的 DataFrame。

  • 筛选文本单元

    • 遍历每个选定entity,获取与该 entity 匹配的 relationship(该实体属于 relationship的 sourcetarget)作为 entity_relationships

    • 每个entity 对应一系列 text_unit_ids,对每个text_unit,取该 entity 的 entity_relationships 和 text_unit 中的 entities 交集,计算基数 num_relationships

    • 根据 entity 的顺序(index)和关系数(num_relationships)排列文本单元。

  • **构建文本单元上下文:**同构建社群上下文

    返回结果

    -----Sources-----
    id|text
    
    idtext
ContextBuilderResult

context_result

class ContextBuilderResult:
    "" "A class to hold the results of the build_context." ""

    context_chunks: str | list [str]
    context_records: dict [str, pd.DataFrame]
    llm_calls: int = 0
    prompt_tokens: int = 0
    output_tokens: int = 0

context_chunks:(context_data)

-----Reports-----
id|title|content

-----Entities-----
id|entity|description|number of relationships

-----Relationships-----
id|source|target|description|weight|links

-----Sources-----
id|text

context_records:

{
    "claims": pd.DataFrame,
    "entities": pd.DataFrame,
    "relationships": pd.DataFrame,
    "reports": pd.DataFrame,
    "sources": pd.DataFrame,
}

response

system_prompt

---Role---

---Goal---

---Target response length and format---
{response_type}

---Data tables---
{context_data}

Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.

Global Search

Baseline RAG 很难处理那些需要在整个数据集上聚合信息来构成答案的问题。

Baseline RAG 对于诸如“数据中前5个主题是什么?”之类的问题表现得很差,因为其依赖于在数据集中通过语义相似性搜索找到相关的文本内容。但是,在这个问题中并没有提供任何指导它去查找正确信息的东西。

GraphRAG 可以解答这类问题,因为LLM生成知识图结构揭示了整个数据集的整体架构(及其主题)。这就使得私人数据集可以被组织成已被预处理过的有意义语义群组。利用Global Search,当面对用户的此类查询时,LLM会用到这些群组以概括出这些主题。

flowchart LR
    A [User Query]
    B [Conversation History]
    
    subgraph ShuffledCommunity [Shuffled Community Report]
        direction LR
        D [Batch 1]
        E [Batch 2] 
        q [...]
        F [Batch N]
    end
    
    subgraph RIR [RIR]
        direction LR
        C1 [Rated Intermediate Response 1]
        C2 [Rated Intermediate Response 2]
        C3 [...]
        CN [Rated Intermediate Response N]
    end
    
    A --> ShuffledCommunity
    B --> ShuffledCommunity
    
    ShuffledCommunity -->|Processed Batches| RIR
    RIR -->|Ranking + Filtering| G [Aggregated Intermediate Responses]
    G --> I [Response]

给定用户的查询以及(如果有的话)会话历史记录

  • Global Search 利用社群图谱层级结构上的某一特定层级中一系列由LLM生成的社群报告作为上下文数据,以map-reduce方式生成响应。
    • Map阶段,社群报告被分割成预定义大小的文本块。随后,每一个文本块都会用来生产出一个中间回复,其中包含一个要点列表,每个要点旁边都有一个数值评分,表示该要点的重要性。
    • Reduce阶段,对这些中间回复中最重要的一些点进行筛选,并把它们聚合起来形成最后的回应所依赖的上下文。

Global Search response的质量可能受社群层级的影响。较低的层级及其详细报告通常会提供更完整的答复,但也低层级的报告数量增更多会导致生成最后答案所需的时问和LLM资源增加。

数据准备

resolve_parquet_file

读取parquet_list得到:

  • final_nodes:create_final_nodes.parquet

  • final_entities:create_final_entities.parquet

  • final_communities:create_final_communities.parquet

  • final_community_reports:create_final_community_reports.parquet

read_indexer_communities

final_nodes 重构社群层级信息,在final_communities中增加子社群字段

  • 确保社群与社群报告的匹配

​ 如果某些社群在报告中缺失,函数会记录警告并删除这些缺失报告的社群记录。确保了社群数据与报告数据的一致性。

  • 使用 final_nodes 生成社群层级结构

    • 分组聚合

      函数首先使用 groupbyfinal_nodes DataFrame 按社群 (community_column) 和层级 (level_column) 进行分组,并将每组内的社群名称 (name_column) 聚合成列表。得到一个每个社群及其层级的 DataFrame: community_df

      communityleveltitle
      intintlist[str]
    • 构建社群层级结构

      • 创建一个字典 community_levels,以层级作为键(level),每个层级对应一个社群的字典(community)和节点名称列表(name)。
      • 遍历 community_df 中的每一行,并填充 community_levels 字典,将每个层级的社群与其对应的节点名称关联起来。
      {
          "level":{
              "community": [node_name,...],...        
          },...
      }
      
    • 生成社群层级关系

      • 对层级关系从高到低进行排序,遍历每个层级来构建父子社群关系。

      • 对于每个父层级(当前层级),查看下一个层级的所有社群,检查是否当前层级的社群是下一个层级社群的父社群。通过比较节点的集合来判断子社群是否属于父社群。

      • 如果子社群的节点集合完全包含在父社群的节点集合中,就确定这个子社群为父社群的子社群。

      • 在每个层级之间,如果发现子社群属于某个父社群,则将父社群、子社群以及子社群的大小(SUB_COMMUNITY_SIZE)信息添加到 community_hierarchy 列表中

      • community_hierarchy 转成 DataFrame

        communitylevelsub_communitysub_community_size
        intintintint
      • community_hierarchycommunity 分组,并将每组内的子社群 (sub_community) 聚合成列表,重排改名

        communitysub_community_ids
        intint
  • 合并层级信息到社群数据

    将重建的社群层级信息(子社群 ID)community_hierarchy 合并回 final_communities DataFrame。对于没有子社群的社群,sub_community_ids 列会被设置为空列表

  • 最终数据处理

    final_communities的每一行转化为Community 对象得到一个列表 _communities

    Community(
        id = id,
        short_id = human_readable_id,
        title = title,
        level = level,
        entity_ids = entity_ids,
        relationship_ids = relationship_ids,
        covariate_ids = None,
        sub_community_ids = sub_community_ids,
        attributes = None,
        size = size,
        period = period
    )
    
read_indexer_reports
read_indexer_entities
knowledge_prompt

knowledge_system_prompt

The response may also include relevant real-world knowledge outside the dataset, but it must be explicitly annotated with a verification tag [LLM: verify]. For example:
"This is an example sentence supported by real-world knowledge [LLM: verify]."

构建上下文

处理对话历史
构建社群上下文
  1. 计算社群权重

    如果 entities 参数不为空,并且需要计算社群权重(即 include_community_weight=True),则为每个社群报告计算权重。计算的权重是与该社群相关联的文本单元的数量。

  • 建立社群和文本单元的关联

    • community_text_unitsDict{community_id: List[chuncked_doc_id, ]}存储每个社群 ID 和与之相关联的 text_unit_id 列表。

      对于每个实体,遍历其所属的社群 ID,并将该实体的文本单元 ID 添加到对应社群的列表中。

  • 为社群报告计算权重

    • 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的text_unit_id数量。
    • 将该数量存储为社群报告的 attributes: 中 ({"occurrence weight": int})。
  • 权重归一化(可选):

    如果 normalize==True,则会对所有社群报告的权重进行归一化。

    归一化的方式是将每个社群的权重除以所有社群中最大权重的值,使得最大权重为 11,其他社群的权重相对最大权重进行调整

  1. 选择报告并打乱报告顺序

    • 选择社群层级大于等于 min_community_rank 的报告。

    • 如果 shuffle_data==True,则会打乱报告的顺序,以避免模型偏向于某些社群。

  2. 生成报告上下文文本

    • 将社群报告按批次处理,每个批次的 token 数量不能超过 max_tokens

      • 当一个批次的 tokens 数量没有超过限制时,直接拼接
      • 当一个批次的 tokens 数量超过限制时,会将当前批次的文本转化为 DataFrame 并保存,然后开始一个新的批次。
    • 生成表头

      对每个报告,提取 header = ['id', 'title', 'occurrence weight', 'content', 'rank']

      batch_text = (
          f "-----{context_name}-----" + "\n" + column_delimiter.join(header) + "\n"
      )
      
      -----Reports-----
      id|title|occurrence weight|content|rank
      
    • 上下文数据

      • 对每个报告:
        • 提取 short_idtitleattributes.occurrence_weightsummary(如果没有则使用full_content)、rank组成 listnew_context
        • column_delimiter|)拼接new_context得到 new_context_text
        • 计算 new_context_texttokens 数量
  3. 返回结果

    • community_context: List[str]: nn 个 batch 的上下文文本

    • community_context_data: Dict{"reports": pd.Dataframe},Dataframe 里记录每个报告的详细信息

      idtitleoccurence weightcontentrank
拼接最终上下文
  • 对话历史
context_prefix = f "{conversation_history_context}\n\n"
  • context_prefixcommunity_context拼接得到 final_context:List[str]
final_context = (
    [f "{context_prefix}{context}" for context in community_context]
    if isinstance(community_context, list)
    else f "{context_prefix}{community_context}"
)
返回结果
class ContextBuilderResult:
    "" "A class to hold the results of the build_context." ""

    context_chunks: str | list [str] (final_context)
    context_records: dict [str, pd.DataFrame] (community_context_data)
    llm_calls: int = 0
    prompt_tokens: int = 0
    output_tokens: int = 0

map_responses

处理一批次内的社群报告数据,生成对用户查询的回答

class SearchResult:
    "" "A Structured Search Result." ""

    response: str | dict [str, Any] | list [dict[str, Any]]
    context_data: str | list [pd.DataFrame] | dict [str, pd.DataFrame]
    # actual text strings that are in the context window, built from context_data
    context_text: str | list [str] | dict [str, str]
    completion_time: float
    # total LLM calls and token usage
    llm_calls: int
    prompt_tokens: int
    output_tokens: int
    # breakdown of LLM calls and token usage
    llm_calls_categories: dict [str, int] | None = None
    prompt_tokens_categories: dict [str, int] | None = None
    output_tokens_categories: dict [str, int] | None = None
map_prompt

map_system_prompt

---Role---
{role}


---Goal---
{goal}

The response should be JSON formatted as follows:
{{
    "points": [
        {{"description": "Description of point 1 [Data: Reports (report ids)]", "score": score_value}},
        {{"description": "Description of point 2 [Data: Reports (report ids)]", "score": score_value}}
    ]
}}

For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Reports (2, 7, 64, 46, 34, +more)]. He is also CEO of company X [Data: Reports (1, 3)]"

where 1, 2, 3, 7, 34, 46, and 64 represent the id (not the index) of the relevant data report in the provided tables.

Do not include information where the supporting evidence for it is not provided.


---Data tables---

{context_data}
map_response
  {
    "answer": "{answer} [Data: Reports (19, 47, 12, 49, 13, +more)]",
    "score": 100
  }

reduce_response

  • 遍历每个 map_responses,从每个响应中提取 answerscore(答案和分数)。

  • 过滤掉分数为 00 的回答,意味着这些回答被认为没有提供有价值的信息。

  • 然后按照 score(分数)对剩余的答案进行降序排序,确保最相关的答案优先返回

  • 将每个关键点(包括分析员编号、分数和答案)格式化为字符串。

    formatted_response_data = []
    formatted_response_data.append(
        f'----Analyst {point ["analyst"] + 1}----'
    )
    formatted_response_data.append(
        f'Importance Score: {point ["score"]}'  # type: ignore
    )
    formatted_response_data.append(point ["answer"])  # type: ignore
    formatted_response_text = "\n".join(formatted_response_data)
    

    text data:

    ----Analyst {i}----
    Importance Score: {sore}
    {content1} [Data: Reports (1,2,3)]
    
  • 如果将格式化后的回答加入总数据后,token 数量超过了最大限制(max_data_tokens),则停止添加更多内容。

reduce_prompt

reduce_system_prompt

---Role---
{Role}

---Goal---
{goal}

For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Reports (2, 7, 34, 46, 64, +more)]. He is also CEO of company X [Data: Reports (1, 3)]"

where 1, 2, 3, 7, 34, 46, and 64 represent the id (not the index) of the relevant data record.

---Target response length and format---
{response_type}

---Analyst Reports---
{report_data}

---Target response length and format---
{response_type}

Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.

DRIFT Search

DRIFT Search (Dynamic Reasoning and Inference with Flexible Traversal) 引入了一种新的方式来进行本地搜索查询,即在搜索过程中加入社群信息。

DRIFT Search结合了全局和局部搜索的特性,以一种平衡计算成本和质量结果的方法生成详细的响应。

HyDE:Precise Zero-Shot Dense Retrieval without Relevance Labels

flowchart TD

    A --> B1[.]
    A --> B2[.]
    A --> B
    
    B1[.] --> C11[.]
    B1[.] --> C12[.]
    B1[.] --> C13[.]
    B1[.] --> C14[.]
    
    B2[.] --> C21[.]
    B2[.] --> C22[.]
    B2[.] --> C23[.]
    B2[.] --> C24[.]
    
    B --> C1[.]
    B --> C2[.]
    B --> C
    B --> C3[.]
    
    C13[.]--> D131[.]
    C13[.]--> D132[.]

	

DRIFT Search 的三个核心阶段

  • A(引言部分):DRIFT会对比用户输入的查询与系统中最相关的前 kk 篇社群报道,并以此为基础给出初步的答案以及一系列后续问题用于更深入地挖掘相关信息。

  • B(跟进部分):通过利用本地搜索技术,DRIFT可以逐步优化查询内容,从而得到更多中间答案及后续问题,这些都将进一步提升检索结果的相关性,帮助搜索引擎更好地定位到包含丰富上下文信息的内容上。图中每一个节点旁边都会有一个符号用来表示该处算法对于是否需要继续拓展当前查询的信心程度如何。

  • **C(输出层级结构):**最后一步则是按照相关性高低对所有已获取的问题和答案进行整理归类形成一个树状结构,在这个过程中我们会尽量兼顾全局视角下的重要发现和细节层面的深度剖析,使得最终呈现的结果既具备良好的适应能力又不失全面性

数据准备

resolve_parquet_file

读取parquet_list得到:

  • final_nodes:create_final_nodes.parquet

  • final_entities:create_final_entities.parquet

  • final_communities:create_final_communities.parquet

  • final_community_reports:create_final_community_reports.parquet

  • final_relationships:create_final_relationships.parquet

向量存储库配置

向量数据库:

  • description_embedding_store: entity.description 实体数据库
  • full_content_embedding_store:community.full_content 社群数据库

构建 query_state

QueryState 类用于管理查询的状态,包括一个存储查询操作及其关系的图结构(使用 NetworkX 库的 MultiDiGraph),以及与查询相关的操作、上下文和追踪信息。

它有一系列方法用于添加操作、建立操作之间的关系、计算排名等功能

class QueryState:
    """Manage the state of the query, including a graph of actions."""

    def __init__(self):
        self.graph = nx.MultiDiGraph()

构建上下文

expand_query

使用随机的社群报告模板展开查询。

  1. 选择社群报告模板:从 self.reports(社群报告列表)中随机选择一个报告,并提取其完整内容(full_content)。
template = secrets.choice(self.reports).full_content
  1. prompt
Create a hypothetical answer to the following query: {query}

Format it to follow the structure of the template below:

{template}                
Ensure that the hypothetical answer does not reference new named entities that are not present in the original query.
  1. 返回扩展后的文本(text)以及 token 使用的详细信息(token_ct 字典)。

    {
        "llm_calls": 1,
        "prompt_tokens": prompt_tokens,
        "output_tokens": output_tokens,
    }
    
  2. 将 text转为 embedding

计算相似度

使用余弦相似度公式计算expand_query 的embedding 和每个report之间的相似度

根据相似度排序并选取前 k 个最相关的报告

初始化搜索 primer search

中间结果生成
  • 拆分报告, 将报告拆分为多个子集,以便并行处理。
  • 异步执行查询分解:对每个报告子集生成多个并发任务,异步执行所有任务并收集结果。

response

  1. score:(1-100) intermediate_answer 对查询的处理效果如何。
  2. intermediate_answer:此答案应与社群报告中的详细程度和长度相匹配,Markdown格式,并且必须有一个标题作为开头,解释接下来的文字是如何与查询关联起来的
  3. follow_up_queries: 可以要求进一步探索该主题的后续询问列表。

prompt

You are a helpful agent designed to reason over a knowledge graph in response to a user query.
This is a unique knowledge graph where edges are freeform text rather than verb operators. You will begin your reasoning looking at a summary of the content of the most relevant communites and will provide:

1. score: How well the intermediate answer addresses the query. A score of 0 indicates a poor, unfocused answer, while a score of 100 indicates a highly focused, relevant answer that addresses the query in its entirety.

2. intermediate_answer: This answer should match the level of detail and length found in the community summaries. The intermediate answer should be exactly 2000 characters long. This must be formatted in markdown and must begin with a header that explains how the following text is related to the query.

3. follow_up_queries: A list of follow-up queries that could be asked to further explore the topic. These should be formatted as a list of strings. Generate at least five good follow-up queries.

Use this information to help you decide whether or not you need more information about the entities mentioned in the report. You may also use your general knowledge to think of entities which may help enrich your answer.

You will also provide a full answer from the content you have available. Use the data provided to generate follow-up queries to help refine your search. Do not ask compound questions, for example: "What is the market cap of Apple and Microsoft?". Use your knowledge of the entity distribution to focus on entity types that will be useful for searching a broad area of the knowledge graph.

For the query:

{query}

The top-ranked community summaries:

{community_reports}

Provide the intermediate answer, and all scores in JSON format following:

{{'intermediate_answer': str,
'score': int,
'follow_up_queries': List[str]}}

Begin:

返回值:返回一个 SearchResult 对象,包含响应内容和相关的上下文信息。

class SearchResult:
    """A Structured Search Result."""

    response: str | dict[str, Any] | list[dict[str, Any]]
    context_data: str | list[pd.DataFrame] | dict[str, pd.DataFrame]
    # actual text strings that are in the context window, built from context_data
    context_text: str | list[str] | dict[str, str]
    completion_time: float
    # total LLM calls and token usage
    llm_calls: int
    prompt_tokens: int
    output_tokens: int
    # breakdown of LLM calls and token usage
    llm_calls_categories: dict[str, int] | None = None
    prompt_tokens_categories: dict[str, int] | None = None
    output_tokens_categories: dict[str, int] | None = None
中间结果处理

分别提取 intermediate_answer、follow_up_queries、score:

  • intermediate_answer: str 直接拼接

  • follow_up_queries:list 直接拼接

  • score:计算均值

返回一个 DriftAction 对象 (表示一个包含查询、答案、评分和后续操作的动作)

class DriftAction:
    """
    Represent an action containing a query, answer, score, and follow-up actions.

    This class encapsulates action strings produced by the LLM in a structured way.
    """

    def __init__(
        self,
        query: str,
        answer: str | None = None,
        follow_ups: list["DriftAction"] | None = None,
    ):
        """
        Initialize the DriftAction with a query, optional answer, and follow-up actions.

        Args:
            query (str): The query for the action.
            answer (Optional[str]): The answer to the query, if available.
            follow_ups (Optional[list[DriftAction]]): A list of follow-up actions.
        """
        self.query = query
        self.answer: str | None = answer  # Corresponds to an 'intermediate_answer'
        self.score: float | None = None
        self.follow_ups: list[DriftAction] = (
            follow_ups if follow_ups is not None else []
        )
        self.metadata: dict[str, Any] = {
            "llm_calls": 0,
            "prompt_tokens": 0,
            "output_tokens": 0,
        }
  • 将该 action 添加进 QueryState.graph
  • 将该 action的 每一个 follow_up_query 作为 query 构建新的 DriftAction 作为节点添加进QueryState.graph中。并为两者添加一条边。

主循环 Search Loop

循环执行搜索直到达到指定的迭代次数(self.config.n)或者没有更多可处理的 action。

  • 创建 actions 列表, 列表由未完成的 action (即 action.answer!=None)组成,并打乱顺序。
  • 选取前 self.config.drift_k_followupsactions,执行后续的搜索操作(asearch_step
asearch_step
  • 创建任务列表,并发执行所有 actions

  • actionDriftAction :对每一个 follow_up_query, 执行 Local Search

    prompt:(跟Local Search 相比多了最后一段)

    ---Role---
    
    ---Goal---
    
    ---Target response length and format---
    {response_type}
    
    ---Data tables---
    {context_data}
    
    Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
    
    Additionally provide a score between 0 and 100 representing how well the response addresses the overall research question: {global_query}. Based on your response, suggest up to five follow-up questions that could be asked to further explore the topic as it relates to the overall research question. Do not include scores or follow up questions in the 'response' field of the JSON, add them to the respective 'score' and 'follow_up_queries' keys of the JSON output. Format your response in JSON with the following keys and values:
    
    {{'response': str, Put your answer, formatted in markdown, here. Do not answer the global query in this section.
    'score': int,
    'follow_up_queries': List[str]}}
    
  • 更新 DriftAction 的字段并返回自身(action)

    context_data

    self.answer = response.pop("response", None)
    self.score = response.pop("score", float("-inf"))
    self.metadata.update({"context_data": search_result.context_data})
    
    self.metadata["llm_calls"] += 1
    self.metadata["prompt_tokens"] += search_result.prompt_tokens
    self.metadata["output_tokens"] += search_result.output_tokens
    
    self.follow_ups = response.pop("follow_up_queries", [])
    
  • 将返回的 actions 添加进 QueryState.graph

  • 将返回的 actions的 每一个 follow_up_query 作为 query 构建新的 DriftAction 作为节点添加进QueryState.graph中。并为两者添加一条边。

  • 迭代次数+1

序列化节点
  • QueryState.graph 中的 nodes 转为 List

    [
      {
        "query": "",
        "answer": "",
        "score": 95.0,
        "metadata": {
          "llm_calls": 0,
          "prompt_tokens": 0,
          "output_tokens": 0
        },
        "id": 0
      },
    ]
    
  • QueryState.graph 中的 nodes 转为 List

    [{"source": 0, "target": 1, "weight": 1.0},]
    
  • 上下文数据(如果需要)

    如果 include_contextTrue,则构建一个 context_data 字典,存储每个节点的查询和上下文数据。

    上下文文本 context_text 是上下文数据的字符串表示。

response

response 为根节点问题的回答。即 node_list[0]的回答