RAG: Multi-Query Retriever with Langchain Notebook

notebooks/langchain/multi-query-retriever-examples/chatbot-with-multi-query-retriever.ipynb

@@ -0,0 +1,289 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Question Answering with Langchain, OpenAI, and MultiQuery Retriever\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/notebooks/langchain/notebooks/langchain/multi-query-retriever-examples/chatbot-with-multi-query-retriever.ipynb)\n",
+    "\n",
+    "This interactive workbook demonstrates example of Elasticsearch's [MultiQuery Retriever](https://api.python.langchain.com/en/latest/retrievers/langchain.retrievers.multi_query.MultiQueryRetriever.html) to generate similar queries for a given user input and apply all queries to retrieve a larger set of relevant documents from a vectorstore.\n",
+    "\n",
+    "Before we begin, we first split the fictional workplace documents into passages with `langchain` and uses OpenAI to transform these passages into embeddings and then store these into Elasticsearch.\n",
+    "\n",
+    "We will then ask a question, generate similar questions using langchain and OpenAI, retrieve relevant passages from the vector store, and use langchain and OpenAI again to provide a summary for the questions."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install packages and import modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.2\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49m/Library/Developer/CommandLineTools/usr/bin/python3 -m pip install --upgrade pip\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python3 -m pip install -qU jq lark elasticsearch langchain_openai tiktoken\n",
+    "\n",
+    "from langchain_openai.embeddings import OpenAIEmbeddings\n",
+    "from langchain.vectorstores import ElasticsearchStore\n",
+    "from langchain_openai.llms import OpenAI\n",
+    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
+    "from getpass import getpass"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Connect to Elasticsearch\n",
+    "\n",
+    "ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook. If you don't have an Elastic Cloud deployment, sign up [here](https://cloud.elastic.co/registration?utm_source=github&utm_content=elasticsearch-labs-notebook) for a free trial. \n",
+    "\n",
+    "We'll use the **Cloud ID** to identify our deployment, because we are using Elastic Cloud deployment. To find the Cloud ID for your deployment, go to https://cloud.elastic.co/deployments and select your deployment.\n",
+    "\n",
+    "We will use [ElasticsearchStore](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html) to connect to our elastic cloud deployment, This would help create and index data easily.  We would also send list of documents that we created in the previous step"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n",
+    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n",
+    "\n",
+    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n",
+    "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n",
+    "\n",
+    "# https://platform.openai.com/api-keys\n",
+    "OPENAI_API_KEY = getpass(\"OpenAI API key: \")\n",
+    "\n",
+    "embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)\n",
+    "\n",
+    "vectorstore = ElasticsearchStore(\n",
+    "    es_cloud_id=ELASTIC_CLOUD_ID, \n",
+    "    es_api_key=ELASTIC_API_KEY,\n",
+    "    index_name=\"chatbot-multi-query-demo\",\n",
+    "    embedding= embeddings,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Indexing Data into Elasticsearch\n",
+    "Let's download the sample dataset and deserialize the document."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from urllib.request import urlopen\n",
+    "import json\n",
+    "\n",
+    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/example-apps/chatbot-rag-app/data/data.json\"\n",
+    "\n",
+    "response = urlopen(url)\n",
+    "data = json.load(response)\n",
+    "\n",
+    "with open('temp.json', 'w') as json_file:\n",
+    "    json.dump(data, json_file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Split Documents into Passages\n",
+    "\n",
+    "We’ll chunk documents into passages in order to improve the retrieval specificity and to ensure that we can provide multiple passages within the context window of the final question answering prompt.\n",
+    "\n",
+    "Here we are chunking documents into 800 token passages with an overlap of 400 tokens.\n",
+    "\n",
+    "Here we are using a simple splitter but Langchain offers more advanced splitters to reduce the chance of context being lost."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.document_loaders import JSONLoader \n",
+    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
+    "\n",
+    "def metadata_func(record: dict, metadata: dict) -> dict:\n",
+    "    metadata[\"name\"] = record.get(\"name\")\n",
+    "    metadata[\"summary\"] = record.get(\"summary\")\n",
+    "    metadata[\"url\"] = record.get(\"url\")\n",
+    "    metadata[\"category\"] = record.get(\"category\")\n",
+    "    metadata[\"updated_at\"] = record.get(\"updated_at\")\n",
+    "\n",
+    "    return metadata\n",
+    "\n",
+    "# For more loaders https://python.langchain.com/docs/modules/data_connection/document_loaders/\n",
+    "# And 3rd party loaders https://python.langchain.com/docs/modules/data_connection/document_loaders/#third-party-loaders\n",
+    "loader = JSONLoader(\n",
+    "    file_path=\"temp.json\",\n",
+    "    jq_schema=\".[]\",\n",
+    "    content_key=\"content\",\n",
+    "    metadata_func=metadata_func,\n",
+    ")\n",
+    "\n",
+    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=512, chunk_overlap=256)\n",
+    "docs = loader.load_and_split(text_splitter=text_splitter)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Bulk Import Passages\n",
+    "\n",
+    "Now that we have split each document into the chunk size of 800, we will now index data to elasticsearch using [ElasticsearchStore.from_documents](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html#langchain.vectorstores.elasticsearch.ElasticsearchStore.from_documents).\n",
+    "\n",
+    "We will use Cloud ID, Password and Index name values set in the `Create cloud deployment` step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "documents = vectorstore.from_documents(\n",
+    "    docs, \n",
+    "    embeddings,  \n",
+    "    index_name=\"chatbot-multi-query-demo\",\n",
+    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
+    "    es_api_key=ELASTIC_API_KEY\n",
+    ")\n",
+    "\n",
+    "llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)\n",
+    "\n",
+    "retriever = MultiQueryRetriever.from_llm(vectorstore.as_retriever(), llm)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Question Answering with MultiQuery Retriever\n",
+    "\n",
+    "Now that we have the passages stored in Elasticsearch, we can now ask a question to get the relevant passages."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:langchain.retrievers.multi_query:Generated queries: ['1. Can you provide information on the sales team at NASA?', '2. How does the sales team operate within NASA?', '3. What are the responsibilities of the NASA sales team?']\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "---- Answer ----\n",
+      "The NASA sales team is a part of the Americas region in the sales organization of the company. It is led by two Area Vice-Presidents, Laura Martinez for North America and Gary Johnson for South America. The team is responsible for promoting and selling the company's products and services in the North and South American markets. They work closely with other departments, such as marketing, product development, and customer support, to ensure the success of the company's sales objectives in the region.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from langchain.schema.runnable import RunnableParallel, RunnablePassthrough\n",
+    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
+    "from langchain.schema import format_document\n",
+    "\n",
+    "import logging\n",
+    "\n",
+    "logging.basicConfig()\n",
+    "logging.getLogger(\"langchain.retrievers.multi_query\").setLevel(logging.INFO)\n",
+    "\n",
+    "LLM_CONTEXT_PROMPT = ChatPromptTemplate.from_template(\n",
+    "    \"\"\"You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as verbose and educational in your response as possible. \n",
+    "    \n",
+    "    context: {context}\n",
+    "    Question: \"{question}\"\n",
+    "    Answer:\n",
+    "    \"\"\"\n",
+    ")\n",
+    "\n",
+    "LLM_DOCUMENT_PROMPT = PromptTemplate.from_template(\"\"\"\n",
+    "---\n",
+    "SOURCE: {name}\n",
+    "{page_content}\n",
+    "---\n",
+    "\"\"\")\n",
+    "\n",
+    "def _combine_documents(\n",
+    "    docs, document_prompt=LLM_DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
+    "):\n",
+    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
+    "    return document_separator.join(doc_strings)\n",
+    "\n",
+    "_context = RunnableParallel(\n",
+    "    context=retriever | _combine_documents,\n",
+    "    question=RunnablePassthrough(),\n",
+    ")\n",
+    "\n",
+    "chain = (_context | LLM_CONTEXT_PROMPT | llm)\n",
+    "\n",
+    "ans = chain.invoke(\"what is the nasa sales team?\")\n",
+    "\n",
+    "print(\"---- Answer ----\")\n",
+    "print(ans)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}