75 LlamaIndex 查询管道入门指南

LlamaIndex 查询管道入门指南

概述

LlamaIndex 提供了一种声明性查询 API，允许你将不同的模块串联在一起，以 orchestrate 简单到高级的工作流程。

这主要围绕我们的 QueryPipeline 抽象展开。加载各种模块（从 LLM 到提示到检索器再到其他管道），将它们全部连接成一个顺序链或 DAG，并端到端运行。

注意：你可以不使用声明性管道抽象（通过使用模块命令式并编写自己的函数）来 orchestrate 所有这些工作流程。那么 QueryPipeline 的优势是什么？

• 用更少的代码/样板文件表达常见的工作流程
• 更高的可读性
• 与常见的低代码/无代码解决方案（例如 LangFlow）更好的集成点
• [未来] 声明性接口允许轻松序列化管道组件，提供管道的可移植性/更容易部署到不同的系统

食谱

在这个食谱中，我们将向你介绍我们的 QueryPipeline 接口，并展示你可以处理的一些基本工作流程。

• 将提示和 LLM 串联在一起
• 将查询重写（提示 + LLM）与检索串联在一起
• 将完整的 RAG 查询管道（查询重写、检索、重排序、响应合成）串联在一起
• 设置自定义查询组件
• 逐步执行管道

设置

这里我们设置了一些数据 + 索引（来自 PG 的散文），我们将在食谱的其余部分使用。

%pip install llama-index-embeddings-openai
%pip install llama-index-postprocessor-cohere-rerank
%pip install llama-index-llms-openai

# 设置 Arize Phoenix 进行日志记录/可观察性
import phoenix as px
px.launch_app()

import llama_index.core
llama_index.core.set_global_handler("arize_phoenix")

import os
os.environ["OPENAI_API_KEY"] = "sk-..."

from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

from llama_index.core import SimpleDirectoryReader
reader = SimpleDirectoryReader("../data/paul_graham")
docs = reader.load_data()

import os
from llama_index.core import (
    StorageContext,
    VectorStoreIndex,
    load_index_from_storage,
)

if not os.path.exists("storage"):
    index = VectorStoreIndex.from_documents(docs)
    index.set_index_id("vector_index")
    index.storage_context.persist("./storage")
else:
    storage_context = StorageContext.from_defaults(persist_dir="storage")
    index = load_index_from_storage(storage_context, index_id="vector_index")

1. 将提示和 LLM 串联在一起

在这一部分，我们展示了一个非常简单的工作流程，将提示与 LLM 串联在一起。

我们只需在初始化时定义链。这是查询管道的一个特例，其中组件纯粹是顺序的，我们自动将输出转换为下一个输入的正确格式。

from llama_index.core.query_pipeline import QueryPipeline
from llama_index.core import PromptTemplate

prompt_str = "请生成与 {movie_name} 相关的电影"
prompt_tmpl = PromptTemplate(prompt_str)
llm = OpenAI(model="gpt-3.5-turbo")

p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)
output = p.run(movie_name="The Departed")
print(str(output))

2. 将查询重写工作流程（提示 + LLM）与检索串联在一起

这里我们尝试一个稍微复杂一些的工作流程，在开始检索之前，通过两个提示发送输入。

1. 生成关于给定主题的问题。
2. 幻觉给定问题的答案，以更好地进行检索。

由于每个提示只接受一个输入，QueryPipeline 会自动将 LLM 输出链接到提示，然后再链接到 LLM。

from llama_index.postprocessor.cohere_rerank import CohereRerank

prompt_str1 = "请生成一个关于 Paul Graham 生活关于以下主题 {topic} 的简明问题"
prompt_tmpl1 = PromptTemplate(prompt_str1)
prompt_str2 = (
    "请写一段回答问题的文字\n"
    "尽量包含尽可能多的关键细节。\n"
    "\n"
    "\n"
    "{query_str}\n"
    "\n"
    "\n"
    '段落:"""\n'
)
prompt_tmpl2 = PromptTemplate(prompt_str2)

llm = OpenAI(model="gpt-3.5-turbo")
retriever = index.as_retriever(similarity_top_k=5)
p = QueryPipeline(
    chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever], verbose=True
)
nodes = p.run(topic="college")
len(nodes)

3. 将完整的 RAG 管道作为 DAG

这里我们将一个完整的 RAG 管道串联在一起，包括查询重写、检索、重排序和响应合成。

由于某些模块依赖于多个输入（例如，响应合成期望检索到的节点和原始问题），我们不能使用链式语法。相反，我们将通过 add_modules 和 add_link 显式构建一个 DAG。

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")

response = p.run(topic="YC")
print(str(response))

4. 定义自定义组件

你可以轻松定义一个自定义组件。只需子类化一个 QueryComponent，实现验证/运行函数 + 一些辅助函数，并将其插入。

让我们将第一个示例中的相关电影生成提示+LLM 链包装成一个自定义组件。

from llama_index.core.query_pipeline import (
    CustomQueryComponent,
    InputKeys,
    OutputKeys,
)
from typing import Dict, Any
from llama_index.core.llms.llm import LLM
from pydantic import Field

class RelatedMovieComponent(CustomQueryComponent):
    """Related movie component."""

    llm: LLM = Field(..., description="OpenAI LLM")

    def _validate_component_inputs(
        self, input: Dict[str, Any]
    ) -> Dict[str, Any]:
        return input

    @property
    def _input_keys(self) -> set:
        return {"movie"}

    @property
    def _output_keys(self) -> set:
        return {"output"}

    def _run_component(self, **kwargs) -> Dict[str, Any]:
        prompt_str = "请生成与 {movie_name} 相关的电影"
        prompt_tmpl = PromptTemplate(prompt_str)
        p = QueryPipeline(chain=[prompt_tmpl, llm])
        return {"output": p.run(movie_name=kwargs["movie"])}

让我们尝试一下自定义组件！我们还将添加一个步骤，将输出转换为莎士比亚风格。

llm = OpenAI(model="gpt-3.5-turbo")
component = RelatedMovieComponent(llm=llm)

prompt_str = """\
这里是一些文字：

{text}

你能用莎士比亚的风格重写这个吗？
"""
prompt_tmpl = PromptTemplate(prompt_str)

p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)
output = p.run(movie="Love Actually")
print(str(output))

5. 逐步执行管道

如果你想更好地调试执行顺序、在每一步记录数据、向用户提供反馈等，逐步执行管道是一个好主意。

要执行管道，你必须创建一个 run_state，然后循环执行。一个基本示例如下。

from llama_index.core.query_pipeline import QueryPipeline
from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI

prompt_str = "请生成与 {movie_name} 相关的电影"
prompt_tmpl = PromptTemplate(prompt_str)
llm = OpenAI(model="gpt-3.5-turbo")

p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)
run_state = p.get_run_state(movie_name="The Departed")

next_module_keys = p.get_next_module_keys(run_state)

while True:
    for module_key in next_module_keys:
        module = run_state.module_dict[module_key]
        module_input = run_state.all_module_inputs[module_key]

        output_dict = module.run_component(**module_input)

        p.process_component_output(
            output_dict,
            module_key,
            run_state,
        )

    next_module_keys = p.get_next_module_keys(
        run_state,
    )

    if not next_module_keys:
        run_state.result_outputs[module_key] = output_dict
        break

print(run_state.result_outputs[module_key]["output"].message.content)

通过这个指南，你应该能够轻松构建和运行各种复杂的查询管道，提升你的数据处理和分析效率。更多详细信息，请查看我们的深入查询转换指南。