定义自定义查询引擎¶

您可以（也应该）定义自己的自定义查询引擎，以便接入下游的 LlamaIndex 工作流，无论您是在构建 RAG、智能体还是其他应用程序。

我们提供了 CustomQueryEngine，让您能够轻松定义自己的查询逻辑。

安装配置¶

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

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

In [ ]:

%pip install llama-index-llms-openai

%pip install llama-index-llms-openai

In [ ]:

!pip install llama-index

!pip install llama-index

In [ ]:

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

下载数据

In [ ]:

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

In [ ]:

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

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

In [ ]:

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。

In [ ]:

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

In [ ]:

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）¶

In [ ]:

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

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

试用体验¶

我们现在使用示例数据进行测试。

尝试方案一（RAGQueryEngine）¶

In [ ]:

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 )

In [ ]:

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

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

In [ ]:

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.

In [ ]:

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

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

尝试方案二（RAGStringQueryEngine）¶

In [ ]:

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

In [ ]:

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

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

In [ ]:

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.