langchain/docs/extras/modules/data_connection/retrievers/MultiQueryRetriever.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8cc82b48",
   "metadata": {},
   "source": [
    "# MultiQueryRetriever\n",
    "\n",
    "Distance-based vector database retrieval embeds (represents) queries in high-dimensional space and finds similar embedded documents based on \"distance\". But, retrieval may produce difference results with subtle changes in query wording or if the embeddings do not capture the semantics of the data well. Prompt engineering / tuning is sometimes done to manually address these problems, but can be tedious.\n",
    "\n",
    "The `MultiQueryRetriever` automates the process of prompt tuning by using an LLM to generate multiple queries from different perspectives for a given user input query. For each query, it retrieves a set of relevant documents and takes the unique union across all queries to get a larger set of potentially relevant documents. By generating multiple perspectives on the same question, the `MultiQueryRetriever` might be able to overcome some of the limitations of the distance-based retrieval and get a richer set of results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "994d6c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a sample vectorDB\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.document_loaders import WebBaseLoader\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "# Load blog post\n",
    "loader = WebBaseLoader(\"https://lilianweng.github.io/posts/2023-06-23-agent/\")\n",
    "data = loader.load()\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)\n",
    "splits = text_splitter.split_documents(data)\n",
    "\n",
    "# VectorDB\n",
    "embedding = OpenAIEmbeddings()\n",
    "vectordb = Chroma.from_documents(documents=splits, embedding=embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cca8f56c",
   "metadata": {},
   "source": [
    "`Simple usage`\n",
    "\n",
    "Specify the LLM to use for query generation, and the retriver will do the rest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "edbca101",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "\n",
    "question = \"What are the approaches to Task Decomposition?\"\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "retriever_from_llm = MultiQueryRetriever.from_llm(\n",
    "    retriever=vectordb.as_retriever(), llm=llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9e6d3b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set logging for the queries\n",
    "import logging\n",
    "\n",
    "logging.basicConfig()\n",
    "logging.getLogger(\"langchain.retrievers.multi_query\").setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e5203612",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:langchain.retrievers.multi_query:Generated queries: ['1. How can Task Decomposition be approached?', '2. What are the different methods for Task Decomposition?', '3. What are the various approaches to decomposing tasks?']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unique_docs = retriever_from_llm.get_relevant_documents(query=question)\n",
    "len(unique_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c54a282f",
   "metadata": {},
   "source": [
    "`Supplying your own prompt`\n",
    "\n",
    "You can also supply a prompt along with an output parser to split the results into a list of queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d9afb0ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from langchain import LLMChain\n",
    "from pydantic import BaseModel, Field\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.output_parsers import PydanticOutputParser\n",
    "\n",
    "\n",
    "# Output parser will split the LLM result into a list of queries\n",
    "class LineList(BaseModel):\n",
    "    # \"lines\" is the key (attribute name) of the parsed output\n",
    "    lines: List[str] = Field(description=\"Lines of text\")\n",
    "\n",
    "\n",
    "class LineListOutputParser(PydanticOutputParser):\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__(pydantic_object=LineList)\n",
    "\n",
    "    def parse(self, text: str) -> LineList:\n",
    "        lines = text.strip().split(\"\\n\")\n",
    "        return LineList(lines=lines)\n",
    "\n",
    "\n",
    "output_parser = LineListOutputParser()\n",
    "\n",
    "QUERY_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"question\"],\n",
    "    template=\"\"\"You are an AI language model assistant. Your task is to generate five \n",
    "    different versions of the given user question to retrieve relevant documents from a vector \n",
    "    database. By generating multiple perspectives on the user question, your goal is to help\n",
    "    the user overcome some of the limitations of the distance-based similarity search. \n",
    "    Provide these alternative questions seperated by newlines.\n",
    "    Original question: {question}\"\"\",\n",
    ")\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "\n",
    "# Chain\n",
    "llm_chain = LLMChain(llm=llm, prompt=QUERY_PROMPT, output_parser=output_parser)\n",
    "\n",
    "# Other inputs\n",
    "question = \"What are the approaches to Task Decomposition?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6660d7ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:langchain.retrievers.multi_query:Generated queries: [\"1. What is the course's perspective on regression?\", '2. Can you provide information on regression as discussed in the course?', '3. How does the course cover the topic of regression?', \"4. What are the course's teachings on regression?\", '5. In relation to the course, what is mentioned about regression?']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "11"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run\n",
    "retriever = MultiQueryRetriever(\n",
    "    retriever=vectordb.as_retriever(), llm_chain=llm_chain, parser_key=\"lines\"\n",
    ")  # \"lines\" is the key (attribute name) of the parsed output\n",
    "\n",
    "# Results\n",
    "unique_docs = retriever.get_relevant_documents(\n",
    "    query=\"What does the course say about regression?\"\n",
    ")\n",
    "len(unique_docs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}