llamaindex 自定义查询引擎的定义与实现

自定义查询引擎的定义与实现

1. 概述

在 LlamaIndex 中，您可以（并且应该）定义自定义查询引擎，以便将其集成到下游的 LlamaIndex 工作流中，无论是构建 RAG（检索增强生成）、代理还是其他应用程序。我们提供了一个 CustomQueryEngine，使定义您自己的查询变得容易。

2. 设置

首先，我们需要加载一些示例数据并对其进行索引。

安装依赖：

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

下载数据：

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

加载文档并创建索引：

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
index = VectorStoreIndex.from_documents(documents)
retriever = index.as_retriever()

3. 构建自定义查询引擎

我们构建一个自定义查询引擎，模拟 RAG 管道。首先执行检索，然后进行合成。

要定义一个 CustomQueryEngine，您只需定义一些初始化参数作为属性并实现 custom_query 函数。

导入必要的库：

from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.core import get_response_synthesizer
from llama_index.core.response_synthesizers import BaseSynthesizer

选项 1：返回一个 Response 对象（RAGQueryEngine）

class RAGQueryEngine(CustomQueryEngine):
    """RAG Query Engine."""

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        response_obj = self.response_synthesizer.synthesize(query_str, nodes)
        return response_obj

这段代码定义了一个名为 RAGQueryEngine 的自定义查询引擎类，继承自 CustomQueryEngine。这个类实现了一个检索增强生成（RAG）管道，首先执行检索操作，然后进行响应合成。下面是对代码的详细解释：

类定义

class RAGQueryEngine(CustomQueryEngine):
    """RAG Query Engine."""

• RAGQueryEngine 类继承自 CustomQueryEngine，表明它是一个自定义查询引擎。
• 类的文档字符串 """RAG Query Engine.""" 简要描述了这个类的用途。

类属性

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer

• retriever：这是一个 BaseRetriever 类型的属性，用于执行检索操作。检索器负责从索引中获取与查询相关的节点（文档片段）。
• response_synthesizer：这是一个 BaseSynthesizer 类型的属性，用于合成最终的响应。响应合成器将检索到的节点和查询字符串结合起来，生成一个连贯的响应。

自定义查询方法

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        response_obj = self.response_synthesizer.synthesize(query_str, nodes)
        return response_obj

• custom_query 方法定义了自定义查询的逻辑。
  • query_str: str：这个方法接受一个字符串参数 query_str，表示用户的查询。
  • nodes = self.retriever.retrieve(query_str)：调用 retriever 的 retrieve 方法，根据查询字符串 query_str 检索相关的节点（文档片段）。
  • response_obj = self.response_synthesizer.synthesize(query_str, nodes)：调用 response_synthesizer 的 synthesize 方法，将查询字符串和检索到的节点结合起来，生成最终的响应对象 response_obj。
  • return response_obj：返回生成的响应对象。

选项 2：返回一个字符串（使用原始 LLM 调用进行说明）

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

qa_prompt = PromptTemplate(
    "Context information is below.\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "Given the context information and not prior knowledge, "
    "answer the query.\n"
    "Query: {query_str}\n"
    "Answer: "
)

class RAGStringQueryEngine(CustomQueryEngine):
    """RAG String Query Engine."""

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer
    llm: OpenAI
    qa_prompt: PromptTemplate

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)

        context_str = "\n\n".join([n.node.get_content() for n in nodes])
        response = self.llm.complete(
            qa_prompt.format(context_str=context_str, query_str=query_str)
        )

        return str(response)

4. 尝试使用自定义查询引擎

尝试选项 1（RAGQueryEngine）

synthesizer = get_response_synthesizer(response_mode="compact")
query_engine = RAGQueryEngine(
    retriever=retriever, response_synthesizer=synthesizer
)
response = query_engine.query("What did the author do growing up?")
print(str(response))

输出：

The author worked on writing and programming outside of school before college. They wrote short stories and tried writing programs on an IBM 1401 computer using an early version of Fortran. They also mentioned getting a microcomputer, building it themselves, and writing simple games and programs on it.

查看源节点内容：

print(response.source_nodes[0].get_content())

尝试选项 2（RAGStringQueryEngine）

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

query_engine = RAGStringQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer,
    llm=llm,
    qa_prompt=qa_prompt,
)
response = query_engine.query("What did the author do growing up?")
print(str(response))

输出：

The author worked on writing and programming before college. They wrote short stories and started programming on the IBM 1401 computer in 9th grade. They later got a microcomputer and continued programming, writing simple games and a word processor.

5. 拓展内容

自定义查询引擎的高级功能：

1. 多语言支持： 自定义查询引擎可以支持多种语言，使得全球用户都能方便地使用。
2. 数据分析： 提供强大的数据分析功能，帮助用户更好地理解和利用数据。
3. 自动化工作流： 支持自动化工作流，减少手动操作，提高效率。
4. 安全性： 提供高级的安全措施，保护用户数据不被泄露。

示例：

# 多语言查询
response_zh = query_engine.query("作者小时候做了什么？")
print(str(response_zh))

# 数据分析查询
analysis_response = query_engine.query("分析最近一个季度的数据趋势")
print(analysis_response)

通过这些详细的讲解和示例，学生们可以更好地理解和掌握自定义查询引擎的定义与实现方法，从而在实际项目中高效地应用。