wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
364 lines
9.7 KiB
Plaintext
364 lines
9.7 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "c91d7b0a",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/custom_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "c2598b32-dfc1-48bf-8ccc-719873aeb05f",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Defining a Custom Query Engine\n",
|
|
"\n",
|
|
"You can (and should) define your custom query engines in order to plug into your downstream LlamaIndex workflows, whether you're building RAG, agents, or other applications.\n",
|
|
"\n",
|
|
"We provide a `CustomQueryEngine` that makes it easy to define your own queries."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f47027d5-19b4-4850-bc14-3f054899f472",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup\n",
|
|
"\n",
|
|
"We first load some sample data and index it."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "0193f7c3",
|
|
"metadata": {},
|
|
"source": [
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "76a072e4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-llms-openai"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c136286f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install llama-index"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1b6e547e-af15-4506-9a4f-ec28c4bf3ce8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import VectorStoreIndex, SimpleDirectoryReader"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "29e58b49",
|
|
"metadata": {},
|
|
"source": [
|
|
"Download Data"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "8cdaf0f6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!mkdir -p 'data/paul_graham/'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4a3b2e77-fbea-451a-b714-696f2c9eb396",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# load documents\n",
|
|
"documents = SimpleDirectoryReader(\"./data//paul_graham/\").load_data()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b94287a5-4ae8-4369-a245-4a8c53790b17",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"index = VectorStoreIndex.from_documents(documents)\n",
|
|
"retriever = index.as_retriever()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0697f50b-5238-4b96-a925-68dfc70138fb",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Building a Custom Query Engine\n",
|
|
"\n",
|
|
"We build a custom query engine that simulates a RAG pipeline. First perform retrieval, and then synthesis.\n",
|
|
"\n",
|
|
"To define a `CustomQueryEngine`, you just have to define some initialization parameters as attributes and implement the `custom_query` function.\n",
|
|
"\n",
|
|
"By default, the `custom_query` can return a `Response` object (which the response synthesizer returns), but it can also just return a string. These are options 1 and 2 respectively."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d3d9c2af-bf04-4ce6-b2ed-b6dd12c7176c",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.query_engine import CustomQueryEngine\n",
|
|
"from llama_index.core.retrievers import BaseRetriever\n",
|
|
"from llama_index.core import get_response_synthesizer\n",
|
|
"from llama_index.core.response_synthesizers import BaseSynthesizer"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "480d6dde-84af-4185-a505-78c95dbf884d",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Option 1 (`RAGQueryEngine`)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "8fd6c9cb-2240-40da-8e69-1bb3d60913dd",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"class RAGQueryEngine(CustomQueryEngine):\n",
|
|
" \"\"\"RAG Query Engine.\"\"\"\n",
|
|
"\n",
|
|
" retriever: BaseRetriever\n",
|
|
" response_synthesizer: BaseSynthesizer\n",
|
|
"\n",
|
|
" def custom_query(self, query_str: str):\n",
|
|
" nodes = self.retriever.retrieve(query_str)\n",
|
|
" response_obj = self.response_synthesizer.synthesize(query_str, nodes)\n",
|
|
" return response_obj"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "de9b50c1-8c0f-4ca1-9ca0-aa984d7270ae",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Option 2 (`RAGStringQueryEngine`)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ef7b34bb-f14a-4e6e-a677-8213cb53f4b8",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Option 2: return a string (we use a raw LLM call for illustration)\n",
|
|
"\n",
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"from llama_index.core import PromptTemplate\n",
|
|
"\n",
|
|
"qa_prompt = PromptTemplate(\n",
|
|
" \"Context information is below.\\n\"\n",
|
|
" \"---------------------\\n\"\n",
|
|
" \"{context_str}\\n\"\n",
|
|
" \"---------------------\\n\"\n",
|
|
" \"Given the context information and not prior knowledge, \"\n",
|
|
" \"answer the query.\\n\"\n",
|
|
" \"Query: {query_str}\\n\"\n",
|
|
" \"Answer: \"\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"class RAGStringQueryEngine(CustomQueryEngine):\n",
|
|
" \"\"\"RAG String Query Engine.\"\"\"\n",
|
|
"\n",
|
|
" retriever: BaseRetriever\n",
|
|
" response_synthesizer: BaseSynthesizer\n",
|
|
" llm: OpenAI\n",
|
|
" qa_prompt: PromptTemplate\n",
|
|
"\n",
|
|
" def custom_query(self, query_str: str):\n",
|
|
" nodes = self.retriever.retrieve(query_str)\n",
|
|
"\n",
|
|
" context_str = \"\\n\\n\".join([n.node.get_content() for n in nodes])\n",
|
|
" response = self.llm.complete(\n",
|
|
" qa_prompt.format(context_str=context_str, query_str=query_str)\n",
|
|
" )\n",
|
|
"\n",
|
|
" return str(response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cb13986c-0431-4f29-9bc3-924424832373",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Trying it out\n",
|
|
"\n",
|
|
"We now try it out on our sample data."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "022d0ae1-74e8-4cae-a460-8f895ed3b293",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Trying Option 1 (`RAGQueryEngine`)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e2b22e9b-fd1d-40ff-ad66-1718384cb5ed",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"synthesizer = get_response_synthesizer(response_mode=\"compact\")\n",
|
|
"query_engine = RAGQueryEngine(\n",
|
|
" retriever=retriever, response_synthesizer=synthesizer\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "26fd76a0-451c-4796-97f8-4524af499c3a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\"What did the author do growing up?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "71d885de-10fc-4c85-b782-108a7c33805e",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"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.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d4da52b0-77a6-43dc-9860-d1a52cee907d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(response.source_nodes[0].get_content())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d0c7c860-06c4-445c-95a7-858240053648",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Trying Option 2 (`RAGStringQueryEngine`)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "09bdd1c9-7bd2-4de0-a807-3eb20f6ede46",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
|
|
"\n",
|
|
"query_engine = RAGStringQueryEngine(\n",
|
|
" retriever=retriever,\n",
|
|
" response_synthesizer=synthesizer,\n",
|
|
" llm=llm,\n",
|
|
" qa_prompt=qa_prompt,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "370a3d91-e16f-4424-b839-5a57b28f21a6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\"What did the author do growing up?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "307831fa-2c8d-4ab2-9d95-78e6c0c61a97",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"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.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "llama_index_v2",
|
|
"language": "python",
|
|
"name": "llama_index_v2"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|