网硕互联帮助中心GraphRAG实现原理-17:四种查询方式之二：Local Search介绍了Local Search查询模式的涉及和实现原理。RIFT Search在Local Search之上增加了一层社区级别的全局引导。普通Local Search直接从查询映射到实体，起点完全依赖实体相似度。DRIFT则先用社区报告做一轮宽泛的语义铺开，由社区报告教会系统应该追问什么，再将追问下放到实体层面执行。这使得DRIFT能触及比单次Local Search更广泛的事实集合。
1. RIFT Search针对Local Search的改进
DRIFT Search（Dynamic Reasoning and Inference with Flexible Traversal）的核心设计思路是先全局铺开、再局部深挖：以社区报告作为宽泛的起点生成初步答案和后续查询，再通过多轮局部搜索逐层细化，最终汇总为综合回答。整个流程总体划分为如下三个阶段：
- Primer: 从索引表community_reports中提取top k个社区报告，并基于报告的内容调用LLM生成一个初始答案和置信度评分，和一批为了完成最终的任务需要继续追问的查询文本；
- Follow-Up: 采用并发的方式采用Local Search模式对每个查询文本实时搜索并生成新的中间答案和追问查询，然后以相同当时迭代处理这些追问查询，直至满足退出的条件；
- Reduce: 将生成的中间答案提交给LLM进行汇总并生成最终的答案。
2. 配置
DRIFT Search利用DRIFTSearchConfig提供了丰富的配置，GraphRagConfig的drift_search字段返回的正是这个对象。
class GraphRagConfig(BaseModel):
drift_search: DRIFTSearchConfig = Field(
description="The drift search configuration.", default=DRIFTSearchConfig()
)
class DRIFTSearchConfig(BaseModel):
prompt: str | None = Field(
description="The DRIFT Search prompt to use.",
default=graphrag_config_defaults.drift_search.prompt,
)
reduce_prompt: str | None = Field(
description="The DRIFT Search reduce prompt to use.",
default=graphrag_config_defaults.drift_search.reduce_prompt,
)
completion_model_id: str = Field(
description="The model ID to use for drift search.",
default=graphrag_config_defaults.drift_search.completion_model_id,
)
embedding_model_id: str = Field(
description="The model ID to use for drift search.",
default=graphrag_config_defaults.drift_search.embedding_model_id,
)
data_max_tokens: int = Field(
description="The data llm maximum tokens.",
default=graphrag_config_defaults.drift_search.data_max_tokens,
)
reduce_max_tokens: int | None = Field(
description="The reduce llm maximum tokens response to produce.",
default=graphrag_config_defaults.drift_search.reduce_max_tokens,
)
reduce_temperature: float = Field(
description="The temperature to use for token generation in reduce.",
default=graphrag_config_defaults.drift_search.reduce_temperature,
)
reduce_max_completion_tokens: int | None = Field(
description="The reduce llm maximum tokens response to produce.",
default=graphrag_config_defaults.drift_search.reduce_max_completion_tokens,
)
concurrency: int = Field(
description="The number of concurrent requests.",
default=graphrag_config_defaults.drift_search.concurrency,
)
drift_k_followups: int = Field(
description="The number of top global results to retrieve.",
default=graphrag_config_defaults.drift_search.drift_k_followups,
)
primer_folds: int = Field(
description="The number of folds for search priming.",
default=graphrag_config_defaults.drift_search.primer_folds,
)
primer_llm_max_tokens: int = Field(
description="The maximum number of tokens for the LLM in primer.",
default=graphrag_config_defaults.drift_search.primer_llm_max_tokens,
)
n_depth: int = Field(
description="The number of DRIFT Searchsteps to take.",
default=graphrag_config_defaults.drift_search.n_depth,
)
local_search_text_unit_prop: float = Field(
description="The proportion of search dedicated to text units.",
default=graphrag_config_defaults.drift_search.local_search_text_unit_prop,
)
local_search_community_prop: float = Field(
description="The proportion of search dedicated to community properties.",
default=graphrag_config_defaults.drift_search.local_search_community_prop,
)
local_search_top_k_mapped_entities: int = Field(
description="The number of top K entities to map during local search.",
default=graphrag_config_defaults.drift_search.local_search_top_k_mapped_entities,
)
local_search_top_k_relationships: int = Field(
description="The number of top K relationships to map during local search.",
default=graphrag_config_defaults.drift_search.local_search_top_k_relationships,
)
local_search_max_data_tokens: int = Field(
description="The maximum context size in tokens for local search.",
default=graphrag_config_defaults.drift_search.local_search_max_data_tokens,
)
local_search_temperature: float = Field(
description="The temperature to use for token generation in local search.",
default=graphrag_config_defaults.drift_search.local_search_temperature,
)
local_search_top_p: float = Field(
description="The top-p value to use for token generation in local search.",
default=graphrag_config_defaults.drift_search.local_search_top_p,
)
local_search_n: int = Field(
description="The number of completions to generate in local search.",
default=graphrag_config_defaults.drift_search.local_search_n,
)
local_search_llm_max_gen_tokens: int | None = Field(
description="The maximum number of generated tokens for the LLM in local search.",
default=graphrag_config_defaults.drift_search.local_search_llm_max_gen_tokens,
)
local_search_llm_max_gen_completion_tokens: int | None = Field(
description="The maximum number of generated tokens for the LLM in local search.",
default=graphrag_config_defaults.drift_search.local_search_llm_max_gen_completion_tokens,
)
个配置选项说明如下：
- prompt：Primer阶段使用的提示词模板，用于生成假设答案、中间答案与后续查询。
- reduce_prompt：Reduce阶段的提示词模板，用于把各分支中间答案融合为最终响应。
- completion_model_id：Completion模型ID。
- embedding_model_id：嵌入模型ID。
- data_max_tokens：数据侧LLM的最大token数，控制送入LLM的上下文数据量。
- reduce_max_tokens/reduce_max_completion_tokens：Reduce阶段LLM响应可产生的最大token数。
- reduce_temperature：Reduce阶段调用LLM的温度参数，控制输出随机性。
- concurrency：并发请求数量，控制Follow-up阶段进行搜索采用的最大并发度。
- drift_k_followups：每轮生成的后续追问（follow-up）数量上限，即从全局结果中取top-k作为追问。
- primer_folds：对提取的社区表格进行分组的数量。
- primer_llm_max_tokens：Primer阶段LLM的最大token数。
- n_depth：Follow-up阶段进行搜索迭代的深度（步数）。
Follow-up阶段采用Local search模式，DRIFTSearchConfig提供了如下的配置选项：
- local_search_text_unit_prop：分配给文本单元的上下文预算比例。
- local_search_community_prop：分配给社区报告的上下文预算比例。
- local_search_top_k_mapped_entities：在实体映射阶段召回的实体数量上限。
- local_search_top_k_relationships：保留的关系数量上限。
- local_search_max_data_tokens：上下文数据的最大token数。
- local_search_temperature：的token生成温度。
- local_search_top_p：的top-p采样参数。
- local_search_n：生成的补全条数。
- local_search_llm_max_gen_tokens/local_search_llm_max_gen_completion_tokens：的LLM最大生成token数。
3. 利用QueryState维护迭代状态
QueryState状态的核心是一个通过networkx.MultiDiGraph对象的维护的多重有向图，即两个节点之间可以包含多条边。MultiDiGraph在QueryState中的节点绑定为一个DriftAction对象，该对象可以视为一个最小推理单元。DriftAction是一个有状态的对象，它维护了如下的状态：
- query：推理的查询文本。
- answer：推理生成的中间答案，None表示尚未完成。如果此答案存在，表示推理操作是否以完成的is_complete字段返回True；
- score：置信度评分，None表示未评分；-inf表示评不出分/失败，用于排序时沉底。
- follow_ups：后续追问查询文本列表。
- metadata：元数据字典
- llm_calls：推理触发的LLM调用次数
- prompt_tokens：推理消耗的提示token数
- output_tokens：推理消耗的输出token数
- context_data: 本次搜索返回的上下文数据
class DriftAction:
def __init__(
self,
query: str,
answer: str | None = None,
follow_ups: list["DriftAction"] | None = None,
):
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,
}
@property
def is_complete(self) –> bool:
return self.answer is not None
async def search(
self,
search_engine: Any,
global_query: str,
k_followups: int,
scorer: Any = None,
):
if self.is_complete:
return self
search_result = await search_engine.search(
query=self.query,
drift_query=global_query,
k_followups=k_followups,
)
_, response = try_parse_json_object(search_result.response, verbose=False)
self.answer = response.pop("response", None)
self.score = float(response.pop("score", "-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", [])
if scorer:
self.compute_score(scorer)
return self
def compute_score(self, scorer: Any):
score = scorer.compute_score(self.query, self.answer)
self.score = (
score if score is not None else float("-inf")
)
def try_parse_json_object(input: str, verbose: bool = True) –> tuple[str, dict]
DriftAction承载的推理体现在它的search方法上，它接受如下的三个参数：
- search_engine：搜索引擎实例，负责实际执行当前动作的查询并返回答案、评分、上下文及追问。由于Follow-up阶段采用Local search，所以这里是要给LocalSearch对象。
- global_query：原始查询文本。
- k_followups：期望生成的追问查询上限。
- scorer：评分器对象，其compute_score方法用来评分。
DriftAction对象的search方法按照如下的流程执行搜索：
- 将当前查询文本、原始查询文本和k_followups作为参数调用search_engine的search方法，并得到表示搜索结果的search_result对象；
- 调用try_parse_json_object一个JSON形式将格式化响应内容解析成response表示的字段。
- 根据search_result和response修改当前状态：
- answer: response.pop(“response”, None);
- score: float(response.pop(“score”, “-inf”));
- follow_ups: response.pop(“follow_up_queries”, [])
- metadata:
- context_data：search_result.context_data
- llm_calls: +1;
- prompt_tokens: += search_result.prompt_tokens
- output_tokens: += search_result.output_tokens
DriftAction还提供了如下这个serialize方法将它的状态（query、answer、score、metadata和follow_ups）序列化成字典。
class DriftAction:
def serialize(self, include_follow_ups: bool = True) –> dict[str, Any]:
data = {
"query": self.query,
"query": self.answer,
"score": self.score,
"metadata": self.metadata,
}
if include_follow_ups:
data["follow_ups"] = [action.serialize() for action in self.follow_ups]
return data
QueryState本质上就是对一个MultiDiGraph对象的封装，并提供了add_action和relate_actions方法来添加作为节点的DriftAction对象和父子节点之间的有向边。当一轮推理完成后，LLM会为当前DriftAction生成一组继续追问的查询文本，后者将被转换对应的子DriftAction通过调用add_all_follow_ups方法添加到图中。add_all_follow_ups方法最终会调用add_action和relate_actions方法完成节点和边的添加。
class QueryState:
def __init__(self):
self.graph = nx.MultiDiGraph()
def add_action(self, action: DriftAction, metadata: dict[str, Any] | None = None):
self.graph.add_node(action, **(metadata or {}))
def relate_actions(self, parent: DriftAction, child: DriftAction, weight: float = 1.0):
self.graph.add_edge(parent, child, weight=weight)
def add_all_follow_ups(
self,
action: DriftAction,
follow_ups: list[DriftAction] | list[str],
weight: float = 1.0,
):
for follow_up in follow_ups:
if isinstance(follow_up, str):
follow_up = DriftAction(query=follow_up)
self.add_action(follow_up)
self.relate_actions(action, follow_up, weight)
QueryState还定义其他的方法：
- find_incomplete_actions：找出尚未完成的（没有生成答案）的DriftAction节点。
- rank_incomplete_actions：根据指定的scorer对未完成的DriftAction节点进行排序（从高到低）。如果没有指定scorer，则调用random.shuffle方法重新洗牌。
- action_token_ct：合计所有DriftAction节点的LLM调用和token统计。
- serialize：对状态实施序列化，最终的输出取决于include_context参数
- True： 返回一个三元组，分别是由节点和边组成的字典（{"nodes": nodes, "edges": edges}），每个节点上名为context_data的元数据元数据字典（以节点的查询文本作为Key）和该字点的字符串表示；
- False：只返回由节点和边组成的字典（{"nodes": nodes, "edges": edges}）
class QueryState:
def find_incomplete_actions(self) –> list[DriftAction]:
return [node for node in self.graph.nodes if not node.is_complete]
def rank_incomplete_actions(
self, scorer: Callable[[DriftAction], float] | None = None
) –> list[DriftAction]:
unanswered = self.find_incomplete_actions()
if scorer:
for node in unanswered:
node.compute_score(scorer)
return sorted(
unanswered,
key=lambda node: (
node.score if node.score is not None else float("-inf")
),
reverse=True,
)
random.shuffle(unanswered)
return list(unanswered)
def action_token_ct(self) –> dict[str, int]:
llm_calls, prompt_tokens, output_tokens = 0, 0, 0
for action in self.graph.nodes:
llm_calls += action.metadata["llm_calls"]
prompt_tokens += action.metadata["prompt_tokens"]
output_tokens += action.metadata["output_tokens"]
return {
"llm_calls": llm_calls,
"prompt_tokens": prompt_tokens,
"output_tokens": output_tokens,
}
def serialize(self, include_context: bool = True
) –> dict[str, Any] | tuple[dict[str, Any], dict[str, Any], str]:
nodes: list[dict[str, Any]] = [
{
**node.serialize(include_follow_ups=False),
"id": node_to_id[node],
**self.graph.nodes[node],
}
for node in self.graph.nodes()
]
edges: list[dict[str, Any]] = [
{
"source": node_to_id[u],
"target": node_to_id[v],
"weight": edge_data.get("weight", 1.0),
}
for u, v, edge_data in self.graph.edges(data=True)
]
if include_context:
context_data = {
node["query"]: node["metadata"]["context_data"]
for node in nodes
if node["metadata"].get("context_data") and node.get("query")
}
context_text = str(context_data)
return {"nodes": nodes, "edges": edges}, context_data, context_text
return {"nodes": nodes, "edges": edges}
4. Primer阶段
作为搜索引擎的DRIFTSearch会在Primer阶段从索引表community_reports中提取与查询文本相似度最高的top k（通过GraphRagConfig的drift_k_followups字段配置）个社区报告，并基于报告的内容调用LLM生成一个初始答案和置信度评分，和一批为了完成最终的任务需要继续追问的查询文本。
4.1 采用HyDE生成更利于社区报告检索的查询文本
为了更加精准地找到与原始查询匹配的社区报告，DRIFTSearch采用了典型的**假设文档嵌入（HyDE：Hypothetical Document Embeddings）**的方式。HyDE核心思路是：如果向量数据库存储的嵌入向量对应的文本具有某种固有的模式（比如它们都是LLM针对类似的提示词生成），那么比直接使用原始查询文本进行检索更好的方式是：让LLM生成一份与原始查询文本具有相同语义又符合此模式的文本来实时检索，因为新的查询文本与向量存储的嵌入空间对得更齐。DRIFTSearch会利用一个类型为PrimerQueryProcessor的可执行对象采用HyDE的方式为原始的查询文本生成更契合community_reports嵌入空间的新查询文本，并将其转换成嵌入向量。
class PrimerQueryProcessor:
def __init__(
self,
chat_model: "LLMCompletion",
text_embedder: "LLMEmbedding",
reports: list[CommunityReport],
tokenizer: Tokenizer | None = None,
):
self.chat_model = chat_model
self.text_embedder = text_embedder
self.tokenizer = tokenizer or chat_model.tokenizer
self.reports = reports
async def expand_query(self, query: str) –> tuple[str, dict[str, int]]:
template = secrets.choice(self.reports).full_content # nosec S311
prompt = (
textwrap.dedent(f"""\\
Create a hypothetical answer to the following query: {query}
Format it to follow the structure of the template below:
""")
+ template
+ textwrap.dedent("""
Ensure that the hypothetical answer does not reference new named entities that are not present in the original query.""")
)
model_response: LLMCompletionResponse = await self.chat_model.completion_async(
messages=prompt
) # type: ignore
text = model_response.content
prompt_tokens = len(self.tokenizer.encode(prompt))
output_tokens = len(self.tokenizer.encode(text))
token_ct = {
"llm_calls": 1,
"prompt_tokens": prompt_tokens,
"output_tokens": output_tokens,
}
if text == "":
return query, token_ct
return text, token_ct
async def __call__(self, query: str) –> tuple[list[float], dict[str, int]]:
hyde_query, token_ct = await self.expand_query(query)
return self.text_embedder.embedding(
input=[hyde_query]
).first_embedding, token_ct
当传入初始查询文本调用__call__方法时，内部会调用expand_query方法，后者会执行如下的操作：
- 从列表中随机提取一个CommunityReport对象，然后根据原始的查询文本定义一条提示词，该提示词旨在指导LLM生成一条与该社区报告内容描述（对应CommunityReport对象的full_content字段）模式匹配的新查询文本;
- 执行LLMCompletion的completion_async方法调用LLM并得到响应内容和token统计；
- 返回[str, dict[str, int]]二元组：
- 更利于社区报告检索的查询文本
- LLM调用次数和token统计字典
__call__方法进一步调用LLMEmbedding的embedding方法将生成的新查询文本转换成嵌入向量，并与LLM调用次数和token统计字典组成返回的二元组。
4.2 检索社区报告和生成搜索上下文
PrimerQueryProcessor采用HyDE生成的查询文本对应的向量最终被DRIFTSearchContextBuilder用来检索社区报告，并进一步生成调用LLM的上下文。DRIFTSearchContextBuilder就是DRIFTSearch专属的ContextBuilder，如下这个DRIFTContextBuilder是它的基类。抽象方法返回一个[pd.DataFrame, dict[str, int]]二元组，前者表示用来承载检索社区报告的DataFrame，后者就是上述的LLM调用次数和token统计字典。
class DRIFTContextBuilder(ABC):
@abstractmethod
async def build_context(
self,
query: str,
**kwargs,
) –> tuple[pd.DataFrame, dict[str, int]]
DRIFTSearchContextBuilder不仅仅实现了为了Primer搜索提供上下文的build_context方法，还定义了init_local_context_builder方法为Follow-up阶段的Local search创建LocalSearchMixedContext。构造函数的很多参数都用在了这个方法上。
class DRIFTSearchContextBuilder(DRIFTContextBuilder):
def __init__(
self,
model: "LLMCompletion",
config: DRIFTSearchConfig,
text_embedder: "LLMEmbedding",
entities: list[Entity],
entity_text_embeddings: VectorStore,
text_units: list[TextUnit] | None = None,
reports: list[CommunityReport] | None = None,
relationships: list[Relationship] | None = None,
covariates: dict[str, list[Covariate]] | None = None,
tokenizer: Tokenizer | None = None,
embedding_vectorstore_key: str = EntityVectorStoreKey.ID,
local_system_prompt: str | None = None,
local_mixed_context: LocalSearchMixedContext | None = None,
reduce_system_prompt: str | None = None,
response_type: str | None = None,
)
def init_local_context_builder(self) –> LocalSearchMixedContext:
return LocalSearchMixedContext(
community_reports=self.reports,
text_units=self.text_units,
entities=self.entities,
relationships=self.relationships,
covariates=self.covariates,
entity_text_embeddings=self.entity_text_embeddings,
embedding_vectorstore_key=self.embedding_vectorstore_key,
text_embedder=self.text_embedder,
tokenizer=self.tokenizer,
)
DRIFTSearchContextBuilder的build_context方法采用如下的方式为Primer搜索创建上下文：
- 调用创建的PrimerQueryProcessor对象针对原始查询文本生成服务HyDE策略的新查询文本，并返回其嵌入向量和LLM调用和token统计；
- 对查询向量进行归一化后，采用计算点积的方法提取出相似度最高的k个（由DRIFTSearchConfig配置的drift_k_followups字段配置）社区报告；
- 将提取出来的社区报告转换成DataFrame（保留short_id, community_id和full_content三个字段）；
- 返回上述DataFrame和LLM调用和token统计字典构成的二元组。
class DRIFTSearchContextBuilder(DRIFTContextBuilder):
async def build_context(
self, query: str, **kwargs
) –> tuple[pd.DataFrame, dict[str, int]]:
query_processor = PrimerQueryProcessor(
chat_model=self.model,
text_embedder=self.text_embedder,
tokenizer=self.tokenizer,
reports=self.reports,
)
query_embedding, token_ct = await query_processor(query)
report_df = self.convert_reports_to_df(self.reports)
query_norm = np.linalg.norm(query_embedding)
document_norms = np.linalg.norm(report_df["full_content_embedding"].to_list(), axis=1)
dot_products = np.dot(np.vstack(report_df["full_content_embedding"].to_list()), query_embedding)
report_df["similarity"] = dot_products / (document_norms * query_norm)
top_k = report_df.nlargest(self.config.drift_k_followups, "similarity")
return top_k.loc[:, ["short_id", "community_id", "full_content"]], token_ct
4.3 首轮搜索
DRIFTSearchContextBuilder利用build_context方法创建的上下文最终会被DRIFTPrimer使用，后者完成了Primer阶段的搜索工作。DRIFTPrimer在初始时需要指定作为配置的DRIFTSearchConfig对象、作为Completion模型的LLMCompletion对象和用来计算token数的Tokenizer。
class DRIFTPrimer:
def __init__(
self,
config: DRIFTSearchConfig,
chat_model: "LLMCompletion",
tokenizer: Tokenizer | None = None,
):
self.chat_model = chat_model
self.config = config
self.tokenizer = tokenizer or chat_model.tokenizer
async def decompose_query(
self, query: str, reports: pd.DataFrame
) –> tuple[dict, dict[str, int]]:
community_reports = "\\n\\n".join(reports["full_content"].tolist())
prompt = DRIFT_PRIMER_PROMPT.format(
query=query, community_reports=community_reports
)
model_response: LLMCompletionResponse[
PrimerResponse
] = await self.chat_model.completion_async(
messages=prompt, response_format=PrimerResponse
) # type: ignore
parsed_response = model_response.formatted_response.model_dump() # type: ignore
token_ct = {
"llm_calls": 1,
"prompt_tokens": len(self.tokenizer.encode(prompt)),
"output_tokens": len(self.tokenizer.encode(model_response.content)),
}
return parsed_response, token_ct
async def search(
self,
query: str,
top_k_reports: pd.DataFrame,
) –> SearchResult:
start_time = time.perf_counter()
report_folds = self.split_reports(top_k_reports)
tasks = [self.decompose_query(query, fold) for fold in report_folds]
results_with_tokens = await tqdm_asyncio.gather(*tasks, leave=False)
completion_time = time.perf_counter() – start_time
return SearchResult(
response=[response for response, _ in results_with_tokens],
context_data={"top_k_reports": top_k_reports},
context_text=top_k_reports.to_json() or "",
completion_time=completion_time,
llm_calls=len(results_with_tokens),
prompt_tokens=sum(ct["prompt_tokens"] for _, ct in results_with_tokens),
output_tokens=sum(ct["output_tokens"] for _, ct in results_with_tokens),
)
def split_reports(self, reports: pd.DataFrame) –> list[pd.DataFrame]:
primer_folds = self.config.primer_folds or 1
if primer_folds == 1:
return [reports]
return [
reports.iloc[indices]
for indices in np.array_split(np.arange(len(reports)), primer_folds)
]
由DRIFTPrimer完成搜索工作实现在它的search方法中，方法的top_k_reports参数对应的DataFrame就是检索得到的社区报告。整个执行流程如下所示：
- 根据GraphRagConfig配置的primer_folds字段对社区报告列表包括进行分割，并将它们分批存入不同的DataFrame中；
- 并行处理每一批社区报告，具体的处理流程如下：
- 提取默认的提示词模板；
- 使用query参数表示的原始查询文本替换模板的{query}占位符
- 将每一份社区报告的full_content字段提取出来，并使用\\n\\n分隔符合并生成的字符串来替换模板的{community_reports}占位符；
- 执行LLMCompletion对象的completion_async方法调用LLM，除了传入完整的提示词外，还基于响应类型PrimerResponse实现结构化响应；
- 从返回的LLMCompletionResponse对象的formatted_response字段直接提取结构化的响应内容，并返回它与LLM调用/token统计组成的二元组。
- 汇总所有任务的执行结果，并返回生成的SearchResult对象。
作为响应类型的PrimerResponse定义如下，它为LLM提供描述输出结构的Schema，所以针对Primer搜索的LLM调用结果将包括：
- intermediate_answer：中间答案；
- score：置信度得分；
- follow_up_queries：后面继续追问的查询文本列表。
class PrimerResponse(BaseModel):
intermediate_answer: str = Field(
description="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.",
)
score: int = Field(
description="A score on 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."
)
follow_up_queries: list[str] = Field(
description="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."
)
如下使调用LLM使用的提示词模板。可以看出此提示词模板最后给出的JSON格式其实是有问题的，'QueryState': List[str]应该替换成'follow_up_queries': List[str]。
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,
'QueryState': List[str]}}
Begin:
5. 完整的流程
DRIFT Search对应的搜索引擎类型为DRIFTSearch，接下来我们就来看看实现在它上面的完整流程。如代码所示，DRIFTSearch派生自BaseSearch[DRIFTSearchContextBuilder]，并将上面介绍的DRIFTSearchContextBuilder作为上下文构建器类型。
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
def __init__(
self,
model: "LLMCompletion",
context_builder: DRIFTSearchContextBuilder,
tokenizer: Tokenizer | None = None,
query_state: QueryState | None = None,
callbacks: list[QueryCallbacks] | None = None,
):
super().__init__(model, context_builder, tokenizer)
self.context_builder = context_builder
self.tokenizer = tokenizer or model.tokenizer
self.query_state = query_state or QueryState()
self.primer = DRIFTPrimer(
config=self.context_builder.config,
chat_model=model,
tokenizer=self.tokenizer,
)
self.callbacks = callbacks or []
self.local_search = self.init_local_search()
def init_local_search(self) –> LocalSearch:
local_context_params = {
"text_unit_prop": self.context_builder.config.local_search_text_unit_prop,
"community_prop": self.context_builder.config.local_search_community_prop,
"top_k_mapped_entities": self.context_builder.config.local_search_top_k_mapped_entities,
"top_k_relationships": self.context_builder.config.local_search_top_k_relationships,
"include_entity_rank": True,
"include_relationship_weight": True,
"include_community_rank": False,
"return_candidate_context": False,
"embedding_vectorstore_key": EntityVectorStoreKey.ID,
"max_context_tokens": self.context_builder.config.local_search_max_data_tokens,
}
model_params = {
"temperature": self.context_builder.config.local_search_temperature,
"n": self.context_builder.config.local_search_n,
"top_p": self.context_builder.config.local_search_top_p,
"max_completion_tokens": self.context_builder.config.local_search_llm_max_gen_completion_tokens,
"response_format_json_object": True,
}
return LocalSearch(
model=self.model,
system_prompt=self.context_builder.local_system_prompt,
context_builder=self.context_builder.local_mixed_context,
tokenizer=self.tokenizer,
model_params=model_params,
context_builder_params=local_context_params,
response_type="multiple paragraphs",
callbacks=self.callbacks,
)
我们在初始化DRIFTSearch对象时需要指定作为Completion模型的LLMCompletion对象和DRIFTSearchContextBuilder对象，后者提供完整的DRIFTSearchConfig配置。如果没有指定tokenizer参数，将会使用LLMCompletion对象提供的Tokenizer。如果没有指定QueryState，会创建一个空的QueryState对象作为驱动整个流程执行的状态。除此之外，初始化时还会创建DRIFTPrimer和LocalSearch。上述的三个阶段的执行流程就体现在它的DRIFTSearch的search方法中。
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
async def search(
self,
query: str,
conversation_history: Any = None,
reduce: bool = True,
**kwargs,
) –> SearchResult:
llm_calls, prompt_tokens, output_tokens = {}, {}, {}
start_time = time.perf_counter()
# Primer
if not self.query_state.graph:
primer_context, token_ct = await self.context_builder.build_context(query)
llm_calls["build_context"] = token_ct["llm_calls"]
prompt_tokens["build_context"] = token_ct["prompt_tokens"]
output_tokens["build_context"] = token_ct["output_tokens"]
primer_response = await self.primer.search(
query=query, top_k_reports=primer_context
)
llm_calls["primer"] = primer_response.llm_calls
prompt_tokens["primer"] = primer_response.prompt_tokens
output_tokens["primer"] = primer_response.output_tokens
init_action = self._process_primer_results(query, primer_response)
self.query_state.add_action(init_action)
self.query_state.add_all_follow_ups(init_action, init_action.follow_ups)
# Follow-up
epochs = 0
while epochs < self.context_builder.config.n_depth:
actions = self.query_state.rank_incomplete_actions()
if len(actions) == 0:
logger.debug("No more actions to take. Exiting DRIFT loop.")
break
actions = actions[: self.context_builder.config.drift_k_followups]
results = await self._search_step(
global_query=query,
k_followups=self.context_builder.config.drift_k_followups,
search_engine=self.local_search,
actions=actions,
)
for action in results:
self.query_state.add_action(action)
self.query_state.add_all_follow_ups(action, action.follow_ups)
epochs += 1
t_elapsed = time.perf_counter() – start_time
token_ct = self.query_state.action_token_ct()
llm_calls["action"] = token_ct["llm_calls"]
prompt_tokens["action"] = token_ct["prompt_tokens"]
output_tokens["action"] = token_ct["output_tokens"]
response_state, context_data, context_text = self.query_state.serialize(
include_context=True
)
reduced_response = response_state
# Reduce
if reduce:
for callback in self.callbacks:
callback.on_reduce_response_start(response_state)
model_params = {
"temperature": self.context_builder.config.reduce_temperature,
"max_completion_tokens": self.context_builder.config.reduce_max_completion_tokens,
}
reduced_response = await self._reduce_response(
responses=response_state,
query=query,
llm_calls=llm_calls,
prompt_tokens=prompt_tokens,
output_tokens=output_tokens,
model_params=model_params,
)
for callback in self.callbacks:
callback.on_reduce_response_end(reduced_response)
return SearchResult(
response=reduced_response,
context_data=context_data,
context_text=context_text,
completion_time=t_elapsed,
llm_calls=sum(llm_calls.values()),
prompt_tokens=sum(prompt_tokens.values()),
output_tokens=sum(output_tokens.values()),
llm_calls_categories=llm_calls,
prompt_tokens_categories=prompt_tokens,
output_tokens_categories=output_tokens,
)
我们根据search方法定义简单说说三个阶段的执行流程：
5.1 Primer阶段
- 调用DRIFTSearchContextBuilder的build_context方法根据查询文本构建上下文，并从中得到top k个与查询文本具有最高相似度的社区报告（被封装在一个DataFrame中）和LLM调用和token统计；
- 将查询文本和承载社区报告列表的DataFrame作为参数调用DRIFTPrimer对象的search方法发起首轮搜索，得到的结构化响应包含初始的中间答案、后续追问查询文本列表和置信度得分；
- 基于响应结果调用_process_primer_results函数创建一个DriftAction，并调用QueryState对象的add_action方法将它作为初始节点添加到用来描述整个推理流程的图中；
5.2 Follow-up阶段
- 开启一个循环，并将GraphRagConfig.n_depth的设置作为迭代次数的上限。每轮迭代执行流程如下：
- 调用QueryState对象的rank_incomplete_actions方法提取得分最高的未完成DriftAction列表，如果列表为空，退出迭代；
- 将GraphRagConfig.drift_k_followups作为限额对DriftAction列表作进一步过滤；
- 将构造函数中创建的LocalSearch对象作为参数调用_search_step方法，后者会并发调用每个DriftAction的search方法;
- 调用QueryState对象的add_action方法添加``DriftAction，这里出现了重复添加，不过add_action`方法的幂等性确保不会出错；
- 调用QueryState对象的add_all_follow_ups方法为生成的追问查询创建对应的DriftAction，并作为子DriftAction的方式添加到推理图中。
- 调用QueryState对象的action_token_ct方法，并对结果提供的LLM调用、输入和输出Token进行统计。
- 调用QueryState对象的serialize方法状态进行序列化（include_context=True），如果不进行后续的Reduce，那么作为序列化结果的第一部分（包含推理图节点和边）将最为返回的SearchResult的响应内容。
上面提到的在每个跌断中并发执行Local search的_search_step方法定义如下
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
async def _search_step(
self,
global_query: str,
k_followups: int,
search_engine: LocalSearch,
actions: list[DriftAction],
) –> list[DriftAction]:
tasks = [
action.search(
search_engine=search_engine,
global_query=global_query,
k_followups=k_followups,
)
for action in actions
]
return await tqdm_asyncio.gather(*tasks, leave=False)
5.3 Reduce阶段
最终的Reduct节点的执行流程主要体现在如下这个_reduce_response方法上，具体流程如下：
- 从QueryState对象的序列化结果中提取每个节点（DriftAction对象）携带的中间答案；
- 从DRIFTSearchContextBuilder的reduce_system_prompt字段提取配置或者默认的系统提示词模板，并将上面获取的中间答案列表和response_type字段返回的响应类型（由命令行参数指定）作为参数进行格式化生成完整的提示词；
- 利用CompletionMessagesBuilder将系统提示词和查询文本转换成角色分别为system和user的两条消息；
- 将消息列表作为参数调用LLMCompletion对象的completion_async方法对LLM发起调用，并调用gather_completion_response_async函数解析响应结果，该结果将作为最终生成的SearchResult对象的响应结果。
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
async def _reduce_response(
self,
responses: str | dict[str, Any],
query: str,
llm_calls: dict[str, int],
prompt_tokens: dict[str, int],
output_tokens: dict[str, int],
**llm_kwargs,
) –> str:
reduce_responses = []
if isinstance(responses, str):
reduce_responses = [responses]
else:
reduce_responses = [
response["answer"]
for response in responses.get("nodes", [])
if response.get("answer")
]
search_prompt = self.context_builder.reduce_system_prompt.format(
context_data=reduce_responses,
response_type=self.context_builder.response_type,
)
messages_builder = (
CompletionMessagesBuilder()
.add_system_message(search_prompt)
.add_user_message(query)
)
model_response = await self.model.completion_async(
messages=messages_builder.build(),
**llm_kwargs.get("model_params", {}),
)
reduced_response = await gather_completion_response_async(model_response)
llm_calls["reduce"] = 1
prompt_tokens["reduce"] = len(self.tokenizer.encode(search_prompt)) + len(
self.tokenizer.encode(query)
)
output_tokens["reduce"] = len(self.tokenizer.encode(reduced_response))
return reduced_response
如下所示的是Reduce阶段整合中间答案调用LLM采用的默认系统提示词模板，占位符{context_data}将被中间答案列表替换。
—Role—
You are a helpful assistant responding to questions about data in the reports provided.
—Goal—
Generate a response of the target length and format that responds to the user's question, summarizing all information in the input reports appropriate for the response length and format, and incorporating any relevant general knowledge while being as specific, accurate and concise as possible.
If you don't know the answer, just say so. Do not make anything up.
Points supported by data should list their data references as follows:
"This is an example sentence supported by multiple data references [Data: <dataset name> (record ids); <dataset name> (record ids)]."
Do not list more than 5 record ids in a single reference. Instead, list the top 5 most relevant record ids and add "+more" to indicate that there are more.
For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Sources (1, 5, 15)]."
Do not include information where the supporting evidence for it is not provided.
If you decide to use general knowledge, you should add a delimiter stating that the information is not supported by the data tables. For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing. [Data: General Knowledge (href)]"
—Data Reports—
{context_data}
—Target response length and format—
{response_type}
—Goal—
Generate a response of the target length and format that responds to the user's question, summarizing all information in the input reports appropriate for the response length and format, and incorporating any relevant general knowledge while being as specific, accurate and concise as possible.
If you don't know the answer, just say so. Do not make anything up.
Points supported by data should list their data references as follows:
"This is an example sentence supported by multiple data references [Data: <dataset name> (record ids); <dataset name> (record ids)]."
Do not list more than 5 record ids in a single reference. Instead, list the top 5 most relevant record ids and add "+more" to indicate that there are more.
For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Sources (1, 5, 15)]."
Do not include information where the supporting evidence for it is not provided.
If you decide to use general knowledge, you should add a delimiter stating that the information is not supported by the data tables. For example:
"Person X is the owner of Company Y and subject to many allegations of wrongdoing. [Data: General Knowledge (href)]".
Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown. Now answer the following query using the data above:
6. 流式调用
DRIFTSearch定义search这个阻塞式调用的方法返回最终的结果之前，它还定义了如下这个基于流式调用的stream_search方法。该方法的实现很清晰：
- 将指定的查询文本和对话历史作为参数调用search方法，由于reduce=False，所以得到的结果仅仅是Follow-up阶段结束后的状态。说得具体一点，得到的结果是QueryState序列化的结果，包含推理图的所有结果和边的状态（{"nodes": nodes, "edges": edges}）。
- 调用_reduce_response_streaming方法以流形式完成Reduce阶段的工作，并实时输出生成的结果。
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
async def stream_search(
self, query: str, conversation_history: ConversationHistory | None = None
) –> AsyncGenerator[str, None]:
result = await self.search(query=query, conversation_history=conversation_history, reduce=False)
if isinstance(result.response, list):
result.response = result.response[0]
for callback in self.callbacks:
callback.on_reduce_response_start(result.response)
model_params = {
"temperature": self.context_builder.config.reduce_temperature,
"max_completion_tokens": self.context_builder.config.reduce_max_completion_tokens,
}
full_response = ""
async for resp in self._reduce_response_streaming(
responses=result.response,
query=query,
model_params=model_params,
):
full_response += resp
yield resp
for callback in self.callbacks:
callback.on_reduce_response_end(full_response)
除了在调用LLMCompletion的completion_async方法时将stream参数设置为True之外，整个流程与实现在search中的Reduce逻辑完全一样。
class DRIFTSearch(BaseSearch[DRIFTSearchContextBuilder]):
async def _reduce_response_streaming(
self,
responses: str | dict[str, Any],
query: str,
model_params: dict[str, Any],
) –> AsyncGenerator[str, None]:
reduce_responses = []
if isinstance(responses, str):
reduce_responses = [responses]
else:
reduce_responses = [
response["answer"]
for response in responses.get("nodes", [])
if response.get("answer")
]
search_prompt = self.context_builder.reduce_system_prompt.format(
context_data=reduce_responses,
response_type=self.context_builder.response_type,
)
messages_builder = (
CompletionMessagesBuilder()
.add_system_message(search_prompt)
.add_user_message(query)
)
response_search: AsyncIterator[LLMCompletionChunk] = await self.model.completion_async(
messages=messages_builder.build(), stream=True, **model_params
)
async for chunk in response_search:
response_text = chunk.choices[0].delta.content or ""
for callback in self.callbacks:
callback.on_llm_new_token(response_text)
yield response_text
评论前必须登录！
注册