{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chatbot with LangChain conversational chain and OpenAI 🤖💬\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/notebooks/generative-ai/chatbot.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "In this notebook we'll build a chatbot that can respond to questions about custom data, such as policies of an employer.\n",
    "\n",
    "The chatbot uses LangChain's `ConversationalRetrievalChain` and has the following capabilities:\n",
    "\n",
    "- Answer questions asked in natural language\n",
    "- Run hybrid search in Elasticsearch to find documents that answer the question\n",
    "- Extract and summarize the answer using OpenAI LLM\n",
    "- Maintain conversational memory for follow-up questions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Requirements 🧰\n",
    "\n",
    "For this example, you will need:\n",
    "\n",
    "- An Elastic deployment\n",
    "  - We'll be using [Elastic Cloud](https://www.elastic.co/guide/en/cloud/current/ec-getting-started.html) for this example (available with a [free trial](https://cloud.elastic.co/registration?onboarding_token=vectorsearch&utm_source=github&utm_content=elasticsearch-labs-notebook))\n",
    "- OpenAI account\n",
    "\n",
    "### Use Elastic Cloud\n",
    "\n",
    "If you don't have an Elastic Cloud deployment, follow these steps to create one.\n",
    "\n",
    "1. Go to [Elastic cloud Registration](https://cloud.elastic.co/registration?onboarding_token=vectorsearch&utm_source=github&utm_content=elasticsearch-labs-notebook) and sign up for a free trial\n",
    "2. Select **Create Deployment** and follow the instructions\n",
    "\n",
    "### Locally install Elasticsearch\n",
    "\n",
    "If you prefer to run Elasticsearch locally, the easiest way is to use Docker. See [Install Elasticsearch with Docker](https://www.elastic.co/search-labs/tutorials/install-elasticsearch/docker) for instructions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages 📦\n",
    "\n",
    "First we `pip install` the packages we need for this example.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "langserve 0.0.21 requires pydantic<2,>=1, but you have pydantic 2.3.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\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.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\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;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -qU langchain openai langchain-elasticsearch tiktoken"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize clients 🔌\n",
    "\n",
    "Next we input credentials with `getpass`. `getpass` is part of the Python standard library and is used to securely prompt for credentials.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "# 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: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and process documents 📄\n",
    "\n",
    "Time to load some data! We'll be using the workplace search example data, which is a list of employee documents and policies.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully loaded 15 documents\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "from urllib.request import urlopen\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/notebooks/generative-ai/data/workplace-docs.json\"\n",
    "\n",
    "response = urlopen(url)\n",
    "\n",
    "workplace_docs = json.loads(response.read())\n",
    "\n",
    "print(f\"Successfully loaded {len(workplace_docs)} documents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chunk documents into passages 🪓\n",
    "\n",
    "As we're chatting with our bot, it will run semantic searches on the index to find the relevant documents. In order for this to be accurate, we need to split the full documents into small chunks (also called passages). This way the semantic search will find the passage within a document that most likely answers our question.\n",
    "\n",
    "We'll use LangChain's `RecursiveCharacterTextSplitter` and split the documents' text at 800 characters with some overlap between chunks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split 15 documents into 24 passages\n"
     ]
    }
   ],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "metadata = []\n",
    "content = []\n",
    "\n",
    "for doc in workplace_docs:\n",
    "    content.append(doc[\"content\"])\n",
    "    metadata.append({\"name\": doc[\"name\"], \"summary\": doc[\"summary\"]})\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=512,\n",
    "    chunk_overlap=256,\n",
    ")\n",
    "docs = text_splitter.create_documents(content, metadatas=metadata)\n",
    "\n",
    "print(f\"Split {len(workplace_docs)} documents into {len(docs)} passages\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's generate the embeddings and index the documents with them.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_elasticsearch import ElasticsearchStore\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "\n",
    "embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "vector_store = ElasticsearchStore.from_documents(\n",
    "    docs,\n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    "    index_name=\"workplace-docs\",\n",
    "    embedding=embeddings,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat with the chatbot 💬\n",
    "\n",
    "Let's initialize our chatbot. We'll define Elasticsearch as a store for retrieving documents and for storing the chat session history, OpenAI as the LLM to interpret questions and summarize answers, then we'll pass these to the conversational chain.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain.memory import ElasticsearchChatMessageHistory\n",
    "from uuid import uuid4\n",
    "\n",
    "\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "llm = OpenAI(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "chat = ConversationalRetrievalChain.from_llm(\n",
    "    llm=llm, retriever=retriever, return_source_documents=True\n",
    ")\n",
    "\n",
    "session_id = str(uuid4())\n",
    "chat_history = ElasticsearchChatMessageHistory(\n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    "    session_id=session_id,\n",
    "    index=\"workplace-docs-chat-history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can ask questions from our chatbot!\n",
    "\n",
    "See how the chat history is passed as context for each question.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CHAT SESSION ID] f1ebce0a-edc1-46d6-a65b-f86f8ab15d2f\n",
      "[QUESTION] What does NASA stand for?\n",
      "[ANSWER]   NASA stands for North America South America region.\n",
      "          [SUPPORTING DOCUMENTS] ['Sales Organization Overview', 'Intellectual Property Policy', 'April Work From Home Update', 'New Employee Onboarding Guide']\n",
      "[QUESTION] Which countries are part of it?\n",
      "[ANSWER]   The North America South America region includes the United States, Canada, Mexico, as well as Central and South America.\n",
      "          [SUPPORTING DOCUMENTS] ['Sales Organization Overview', 'Wfh Policy Update May 2023', 'Fy2024 Company Sales Strategy', 'New Employee Onboarding Guide']\n",
      "[QUESTION] Who are the team's leads?\n",
      "[ANSWER]   Laura Martinez is the Area Vice-President of North America, and Gary Johnson is the Area Vice-President of South America.\n",
      "          [SUPPORTING DOCUMENTS] ['Sales Organization Overview', 'Fy2024 Company Sales Strategy', 'Wfh Policy Update May 2023', 'Fy2024 Company Sales Strategy']\n"
     ]
    }
   ],
   "source": [
    "# Define a convenience function for Q&A\n",
    "def ask(question, chat_history):\n",
    "    result = chat({\"question\": question, \"chat_history\": chat_history.messages})\n",
    "    print(\n",
    "        f\"\"\"[QUESTION] {question}\n",
    "[ANSWER]  {result[\"answer\"]}\n",
    "          [SUPPORTING DOCUMENTS] {list(map(lambda d: d.metadata[\"name\"], list(result[\"source_documents\"])))}\"\"\"\n",
    "    )\n",
    "    chat_history.add_user_message(result[\"question\"])\n",
    "    chat_history.add_ai_message(result[\"answer\"])\n",
    "\n",
    "\n",
    "# Chat away!\n",
    "print(f\"[CHAT SESSION ID] {session_id}\")\n",
    "ask(\"What does NASA stand for?\", chat_history)\n",
    "ask(\"Which countries are part of it?\", chat_history)\n",
    "ask(\"Who are the team's leads?\", chat_history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "💡 _Try experimenting with other questions or after clearing the workplace data, and observe how the responses change._\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Optional) Clean up 🧹\n",
    "\n",
    "Once we're done, we can clean up the chat history for this session...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_history.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... or delete the indices.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ObjectApiResponse({'acknowledged': True})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vector_store.client.indices.delete(index=\"workplace-docs\")\n",
    "vector_store.client.indices.delete(index=\"workplace-docs-chat-history\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.4 64-bit",
   "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.12.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}