定义自定义查询引擎¶

你可以（也应该）定义自己的自定义查询引擎，以便将其接入下游的 LlamaIndex 工作流，无论你是在构建 RAG、代理还是其他应用。

我们提供一个 CustomQueryEngine，可以轻松定义你自己的查询。

设置¶

我们首先加载一些示例数据并为其建立索引。

如果你在 colab 上打开此 Notebook，你可能需要安装 LlamaIndex 🦙。

输入 [ ]

%pip install llama-index-llms-openai

%pip install llama-index-llms-openai

输入 [ ]

!pip install llama-index

!pip install llama-index

输入 [ ]

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

下载数据

输入 [ ]

!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'

!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'

输入 [ ]

# load documents
documents = SimpleDirectoryReader("./data//paul_graham/").load_data()

# load documents documents = SimpleDirectoryReader("./data//paul_graham/").load_data()

输入 [ ]

index = VectorStoreIndex.from_documents(documents)
retriever = index.as_retriever()

index = VectorStoreIndex.from_documents(documents) retriever = index.as_retriever()

构建自定义查询引擎¶

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

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

默认情况下，custom_query 可以返回一个 Response 对象（响应合成器返回的），但它也可以只返回一个字符串。这分别是选项 1 和选项 2。

输入 [ ]

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

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 (`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

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

选项 2 (`RAGStringQueryEngine`)¶

输入 [ ]

# Option 2: return a string (we use a raw LLM call for illustration)

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)

# 选项 2：返回字符串（我们使用原始 LLM 调用作为示例） from llama_index.llms.openai import OpenAI from llama_index.core import PromptTemplate qa_prompt = PromptTemplate( "上下文信息如下。\n" "---------------------\n" "{context_str}\n" "---------------------\n" "根据上下文信息，而非先验知识，" "回答查询。\n" "查询：{query_str}\n" "回答：" ) 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)

试用¶

我们现在在示例数据上试用它。

试用选项 1 (`RAGQueryEngine`)¶

输入 [ ]

synthesizer = get_response_synthesizer(response_mode="compact")
query_engine = RAGQueryEngine(
    retriever=retriever, response_synthesizer=synthesizer
)

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

response = query_engine.query("What did the author do growing up?")

输入 [ ]

print(str(response))

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

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

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

response = query_engine.query("What did the author do growing up?")

输入 [ ]

print(str(response))

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.