[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_nodes和final_entities中读取实体信息,并根据社群层级进行筛选和处理
-
根据
community_level筛选社群报告和节点数据 -
处理
final_nodes数据- 提取
final_nodes中的id、degree和community列 - 将所有节点的
community列缺失值填充为-1,并转换为整数类型,将degree列转换为整数类型。 - 将所有节点按
id和degree分组,聚合community列,去重,集合(set)转换为字符串列表。
- 提取
-
合并
final_nodes和final_entities数据为final_df- 使用
id列将处理后的节点数据nodes_df和实体数据entities_df合并,得到一个包含实体信息的 DataFrame,去重。
- 使用
-
最终数据处理
将
final_df的每一行转化为Entity对象得到一个列表_entitiesEntity( 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_selection为False(即不启用动态社群选择),函数会根据每个节点的社群层级,确定该节点所在的最大社群层级。 - 将
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对象得到一个列表reportsCommunityReport( 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_entities 或 self.community_reports 为空,直接返回空上下文文本和空的 DataFrame。
社群匹配
- 对每个选中的实体,其所属的社群 match 数量 +1。
- 在社群报告筛选中出这些社群,并按照 match 数量排序,如果有多个社群匹配的实体数相同,还会按社群
rank进行二次排序,得到selected_communities
构建社群上下文
-
计算社群权重:
如果
entities参数不为空,并且需要计算社群权重(即include_community_weight=True),则为每个社群报告计算权重。计算的权重是与该社群相关联的文本单元的数量。-
建立社群和文本单元的关联:
-
community_text_units:Dict{community_id: List[chuncked_doc_id, ]}存储每个社群 ID 和与之相关联的text_unit_id列表。对于每个实体,遍历其所属的社群 ID,并将该实体的文本单元 ID 添加到对应社群的列表中。
-
-
为社群报告计算权重:
- 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的
text_unit_id数量。 - 将该数量存储为社群报告的
attributes:中 ({"occurrence weight": int})。
- 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的
-
权重归一化(可选):
如果
normalize==True,则会对所有社群报告的权重进行归一化。归一化的方式是将每个社群的权重除以所有社群中最大权重的值,使得最大权重为 ,其他社群的权重相对最大权重进行调整
-
-
选择报告并打乱报告顺序:
-
选择社群层级大于等于
min_community_rank的报告。 -
如果
shuffle_data==True,则会打乱报告的顺序,以避免模型偏向于某些社群。
-
-
生成报告上下文文本:
-
将社群报告按批次处理,每个批次的 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_id、title、attributes.occurrence_weight、summary(如果没有则使用full_content)、rank组成list:new_context - 用
column_delimiter(|)拼接new_context得到new_context_text - 计算
new_context_texttokens 数量
- 提取
- 对每个报告:
-
-
返回结果
-
community_context:List[str]: 个 batch 的上下文文本 -
community_context_data:Dict{"reports": pd.Dataframe},Dataframe 里记录每个报告的详细信息id title occurence weight content rank
-
- 如果
return_candidate_context为True,还会调用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
| id | entity | description | number of relationships |
|---|---|---|---|
-
current_context_text:List[str]: 个 batch 的上下文文本 -
record_df:Dict{"reports": pd.Dataframe},Dataframe 里记录每个实体的详细信息
构建relationships上下文:
-
筛选关系:
筛选出与所选实体相关的relationship,并限制返回最多
relationship_budget个关系。-
筛选
in_network_relationships:- source 和 target 均在
selected_entity中
- source 和 target 均在
-
筛选
out_network_relationships:- source 和 target 任一在
selected_entity中 - 按照
ranking_attribute(rank) 排序
- source 和 target 任一在
-
计算
out_network_entities, 表示out_network_relationships中,不在selected_entity的部分。 -
为每个
out_network_entity计算链接数:(即它们作为source或target在out_network_relationships中出现的次数)。对out_network_relationships按照链接数排序(链接数相同的情况下,再按rank或weight排序)。 -
计算
relationship_budgetrelationship_budget = top_k_relationships * len(selected_entities) -
返回
in_network_relationships+out_network_relationships的前relationship_budget个关系。
-
-
构建relationship上下文 同构建社群上下文
返回结果
-----Relationships----- id|source|target|description|weight|linksid source target description weight links
- 协变量上下文:对于每个协变量,函数会使用
build_covariates_context来生成上下文。
构建协变量上下文
build_text_unit_context
如果没有选中的实体或没有文本单元数据,则直接返回空的上下文文本和空的 DataFrame。
-
筛选文本单元:
-
遍历每个选定entity,获取与该 entity 匹配的 relationship(该实体属于 relationship的
source或target)作为entity_relationships。 -
每个entity 对应一系列
text_unit_ids,对每个text_unit,取该 entity 的entity_relationships和 text_unit 中的 entities 交集,计算基数num_relationships -
根据 entity 的顺序(
index)和关系数(num_relationships)排列文本单元。
-
-
**构建文本单元上下文:**同构建社群上下文
返回结果
-----Sources----- id|textid text
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
---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生成社群层级结构-
分组聚合:
函数首先使用
groupby对final_nodesDataFrame 按社群 (community_column) 和层级 (level_column) 进行分组,并将每组内的社群名称 (name_column) 聚合成列表。得到一个每个社群及其层级的 DataFrame:community_dfcommunity level title int int list[str] -
构建社群层级结构
- 创建一个字典
community_levels,以层级作为键(level),每个层级对应一个社群的字典(community)和节点名称列表(name)。 - 遍历
community_df中的每一行,并填充community_levels字典,将每个层级的社群与其对应的节点名称关联起来。
{ "level":{ "community": [node_name,...],... },... } - 创建一个字典
-
生成社群层级关系
-
对层级关系从高到低进行排序,遍历每个层级来构建父子社群关系。
-
对于每个父层级(当前层级),查看下一个层级的所有社群,检查是否当前层级的社群是下一个层级社群的父社群。通过比较节点的集合来判断子社群是否属于父社群。
-
如果子社群的节点集合完全包含在父社群的节点集合中,就确定这个子社群为父社群的子社群。
-
在每个层级之间,如果发现子社群属于某个父社群,则将父社群、子社群以及子社群的大小(
SUB_COMMUNITY_SIZE)信息添加到community_hierarchy列表中 -
将
community_hierarchy转成 DataFramecommunity level sub_community sub_community_size int int int int -
将
community_hierarchy按community分组,并将每组内的子社群 (sub_community) 聚合成列表,重排改名community sub_community_ids int int
-
-
-
合并层级信息到社群数据:
将重建的社群层级信息(子社群 ID)
community_hierarchy合并回final_communitiesDataFrame。对于没有子社群的社群,sub_community_ids列会被设置为空列表 -
最终数据处理
将
final_communities的每一行转化为Community对象得到一个列表_communitiesCommunity( 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
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]."
构建上下文
处理对话历史
构建社群上下文
-
计算社群权重:
如果
entities参数不为空,并且需要计算社群权重(即include_community_weight=True),则为每个社群报告计算权重。计算的权重是与该社群相关联的文本单元的数量。
-
建立社群和文本单元的关联:
-
community_text_units:Dict{community_id: List[chuncked_doc_id, ]}存储每个社群 ID 和与之相关联的text_unit_id列表。对于每个实体,遍历其所属的社群 ID,并将该实体的文本单元 ID 添加到对应社群的列表中。
-
-
为社群报告计算权重:
- 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的
text_unit_id数量。 - 将该数量存储为社群报告的
attributes:中 ({"occurrence weight": int})。
- 遍历每个社群报告,检查该报告的社群 ID,并计算其对应的
-
权重归一化(可选):
如果
normalize==True,则会对所有社群报告的权重进行归一化。归一化的方式是将每个社群的权重除以所有社群中最大权重的值,使得最大权重为 ,其他社群的权重相对最大权重进行调整
-
选择报告并打乱报告顺序:
-
选择社群层级大于等于
min_community_rank的报告。 -
如果
shuffle_data==True,则会打乱报告的顺序,以避免模型偏向于某些社群。
-
-
生成报告上下文文本:
-
将社群报告按批次处理,每个批次的 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_id、title、attributes.occurrence_weight、summary(如果没有则使用full_content)、rank组成list:new_context - 用
column_delimiter(|)拼接new_context得到new_context_text - 计算
new_context_texttokens 数量
- 提取
- 对每个报告:
-
-
返回结果
-
community_context:List[str]: 个 batch 的上下文文本 -
community_context_data:Dict{"reports": pd.Dataframe},Dataframe 里记录每个报告的详细信息id title occurence weight content rank
-
拼接最终上下文
- 对话历史
context_prefix = f "{conversation_history_context}\n\n"
- 将
context_prefix与community_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
---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,从每个响应中提取answer和score(答案和分数)。 -
过滤掉分数为 的回答,意味着这些回答被认为没有提供有价值的信息。
-
然后按照
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
---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会对比用户输入的查询与系统中最相关的前 篇社群报道,并以此为基础给出初步的答案以及一系列后续问题用于更深入地挖掘相关信息。
-
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
使用随机的社群报告模板展开查询。
- 选择社群报告模板:从
self.reports(社群报告列表)中随机选择一个报告,并提取其完整内容(full_content)。
template = secrets.choice(self.reports).full_content
- 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.
-
返回扩展后的文本(
text)以及 token 使用的详细信息(token_ct字典)。{ "llm_calls": 1, "prompt_tokens": prompt_tokens, "output_tokens": output_tokens, } -
将 text转为 embedding
计算相似度
使用余弦相似度公式计算expand_query 的embedding 和每个report之间的相似度
根据相似度排序并选取前 k 个最相关的报告
初始化搜索 primer search
中间结果生成
- 拆分报告, 将报告拆分为多个子集,以便并行处理。
- 异步执行查询分解:对每个报告子集生成多个并发任务,异步执行所有任务并收集结果。
response:
- score:(1-100) intermediate_answer 对查询的处理效果如何。
- intermediate_answer:此答案应与社群报告中的详细程度和长度相匹配,Markdown格式,并且必须有一个标题作为开头,解释接下来的文字是如何与查询关联起来的
- 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_followups个actions,执行后续的搜索操作(asearch_step)
asearch_step
-
创建任务列表,并发执行所有
actions。 -
action:DriftAction:对每一个 follow_up_query, 执行 Local Searchprompt:(跟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_dataself.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_context为True,则构建一个context_data字典,存储每个节点的查询和上下文数据。上下文文本
context_text是上下文数据的字符串表示。
response
response 为根节点问题的回答。即 node_list[0]的回答