73 构建高效查询管道：QueryPipeline 的设置与使用指南

本文链接：https://blog.youkuaiyun.com/xycxycooo/article/details/141385920

构建高效查询管道：QueryPipeline 的设置与使用指南

在现代数据处理和分析中，构建高效的查询管道（Query Pipeline）是提升工作效率的关键。本文将深入探讨如何使用 QueryPipeline 来设置和运行各种复杂的查询流程，无论是简单的线性链还是复杂的DAG（有向无环图）结构。

设置查询管道

定义顺序链

对于一些简单的线性查询流程，QueryPipeline 提供了简化的链式语法。例如：

提示 -> LLM（大型语言模型） -> 输出解析
提示 -> LLM -> 提示 -> LLM
检索器 -> 响应合成器

以下是一个简单的示例代码：

from llama_index.core.query_pipeline import QueryPipeline

# 尝试链式基本提示
prompt_str = "请生成与 {movie_name} 相关的电影"
prompt_tmpl = PromptTemplate(prompt_str)
llm = OpenAI(model="gpt-3.5-turbo")

p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)

定义DAG

对于更复杂的查询流程，可能需要设置一个DAG。QueryPipeline 提供了低级API来添加模块及其键，并定义模块之间的链接。

from llama_index.postprocessor.cohere_rerank import CohereRerank
from llama_index.core.response_synthesizers import TreeSummarize

# 定义模块
prompt_str = "请生成一个关于 Paul Graham 生活的问题，主题是 {topic}"
prompt_tmpl = PromptTemplate(prompt_str)
llm = OpenAI(model="gpt-3.5-turbo")
retriever = index.as_retriever(similarity_top_k=3)
reranker = CohereRerank()
summarizer = TreeSummarize(llm=llm)

# 定义查询管道
p = QueryPipeline(verbose=True)
p.add_modules(
    {
        "llm": llm,
        "prompt_tmpl": prompt_tmpl,
        "retriever": retriever,
        "summarizer": summarizer,
        "reranker": reranker,
    }
)
p.add_link("prompt_tmpl", "llm")
p.add_link("llm", "retriever")
p.add_link("retriever", "reranker", dest_key="nodes")
p.add_link("llm", "reranker", dest_key="query_str")
p.add_link("reranker", "summarizer", dest_key="nodes")
p.add_link("llm", "summarizer", dest_key="query_str")

运行查询管道

单输入/单输出

输入是第一个组件的kwargs。如果最后一个组件的输出是单个对象，则直接返回该对象。

output = p.run(topic="YC")
# 输出类型是 Response
type(output)

多输入/多输出

如果你的DAG有多个根节点或输出节点，可以使用 run_multi。传入一个包含模块键 -> 输入字典的输入字典。输出是模块键 -> 输出字典的字典。

output_dict = p.run_multi({"llm": {"topic": "YC"}})
print(output_dict)

# 输出字典是 {"summarizer": {"output": response}}

定义部分输入

如果你想预填充某个模块的某些输入，可以使用 partial。然后DAG只需连接到未填充的输入。

summarizer = TreeSummarize(llm=llm)
summarizer_c = summarizer.as_query_component(partial={"nodes": nodes})
# 可以定义一个链，因为llm输出进入query_str，nodes是预填充的
p = QueryPipeline(chain=[prompt_tmpl, llm, summarizer_c])
# 运行管道
p.run(topic="YC")

批量输入

如果你想为几轮单/多输入运行管道，可以在函数调用中设置 batch=True。传入一个你希望运行的单/多输入列表。批量模式将按输入顺序返回响应列表。

output = p.run(topic=["YC", "RAG", "LlamaIndex"], batch=True)
# 输出是 [ResponseYC, ResponseRAG, ResponseLlamaIndex]
print(output)

中间输出

如果你想获取QueryPipeline中模块的中间输出，可以使用 run_with_intermediates 或 run_multi_with_intermediates。

output, intermediates = p.run_with_intermediates(topic="YC")
print(output)
print(intermediates)

# 输出是 (Response, {"module_key": ComponentIntermediates("inputs": {}, "outputs": {})})

定义自定义查询组件

你可以轻松定义自定义组件：将函数传递给 FnComponent 或子类化 CustomQueryComponent。

from llama_index.core.query_pipeline import FnComponent

def add(a: int, b: int) -> int:
    """Adds two numbers."""
    return a + b

add_component = FnComponent(fn=add, output_key="output")

# 输入键到add_component是"a"和"b"，输出键是'output'

确保输出兼容

通过在QueryPipeline中链接模块，一个模块的输出进入下一个模块的输入。通常你必须确保预期的输出和输入类型大致匹配。

我们说“大致”是因为我们对现有模块做了一些魔法，以确保“可字符串化”的输出可以传递到可以作为“字符串”查询的输入中。某些输出类型被视为Stringable - CompletionResponse, ChatResponse, Response, QueryBundle等。检索器/查询引擎会自动将字符串输入转换为QueryBundle对象。

这让你可以做一些原本需要样板字符串转换的工作流程，例如：