week5/community-contributions/day5 - generating answers with citations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ec3276da-2cc6-4558-beb3-00cf4dc1ac0a",
   "metadata": {},
   "source": [
    "# Generating answers with source citations\n",
    "### This Notebook contains a sample showing how to generate answers with inline & end of the answer citations pointing to the original source document used to answer the question\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "<h4><u>Prerequisite:</u> Please run the <a href=\"../day5.ipynb\" >Day 5 notebook</a> to create & populate the vector database before executing this notebook</h4>\n",
    "</div>\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59921cc4-ecb7-460a-a15a-1b4490f3cf25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import gradio as gr\n",
    "from openai import OpenAI\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_chroma import Chroma\n",
    "from IPython.display import Markdown, display, update_display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5d79d2-d7fb-473e-bd0f-79abbb7b69ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "MODEL = \"gpt-4o-mini\"\n",
    "db_name = \"vector_db\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b34c33fc-b7d8-4880-ba8f-0db51b24a8a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai = OpenAI()\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c92170d-f4e9-45d4-8b72-f3779103a551",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the existing vector database that you created from the Day5 notebook\n",
    "if os.path.exists(f\"..\\\\{db_name}\"):\n",
    "  vectorstore = Chroma(embedding_function=embeddings, persist_directory=f\"..\\\\{db_name}\")\n",
    "  print(f\"Vectorstore loaded with {vectorstore._collection.count()} documents\")\n",
    "else:\n",
    "  print(\"Vector store doesn't exist. Please run the Day 5 notebook first to create Chroma Vector DB & injest the data.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78809d3b-bea3-436d-bb79-6adc93757a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"You are an assistant for question-answering tasks. \n",
    "Use the following pieces of retrieved context to answer the question. \n",
    "If you don't know the answer, just say that you don't know.\n",
    "Use three sentences maximum and keep the answer concise.\n",
    "Use the following markdown format to answer the question along with the Source used to generate the answer, add inline citation for each sentence & add end of the answer citations:\n",
    "'CEO of Insurellm is Avery Lancaster [[1]](Source Link 1). Who is also a co-founder [[2]](Source Link 2)\n",
    "Citations: (Note: No duplicates allowed in the below list)\n",
    "\n",
    "[1 - Source Title 1](Link 1)\n",
    "[2 - Source Title 2](Link 2)\n",
    "...\n",
    "[n - Source Title n](Link n)'\n",
    " \n",
    "Example answer: \n",
    "'CEO of Insurellm is Avery Lancaster [[1]](knowledge-base\\\\company\\\\about.md). Who is also a co-founder [[2]](knowledge-base\\\\employees\\\\Avery Lancaster.md)\n",
    "Citations:\n",
    "\n",
    "[1 - About Company](knowledge-base\\\\company\\\\about.md)\n",
    "[2 - Avery Lancaster employees](knowledge-base\\\\employees\\\\Avery Lancaster.md)'\n",
    " \n",
    "Important Note: Have unique end of the answer citations. Don't give duplicate citation numbers for the same source link, reuse the same citation number if the same source link is referenced multiple times.\n",
    "'\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee6afbd-03af-448e-b587-991c555930bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the below port if jupyter notebook is running in a different port\n",
    "jupyter_notebook_port = \"8888\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed36033b-a67e-4278-b990-c52255274b63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_user_prompt(message):\n",
    "    retriever = vectorstore.as_retriever()\n",
    "    results = retriever.invoke(message)\n",
    "    doc_chunk_merged = \"\"\n",
    "    for doc_chunk in results:        \n",
    "        source = f\"http://localhost:{jupyter_notebook_port}/lab/tree/week5/\" + doc_chunk.metadata.get(\"source\").replace(\"\\\\\",\"/\")\n",
    "        title = doc_chunk.metadata.get(\"doc_type\") + \" -> \" + source.split('\\\\')[-1][:-3]\n",
    "        doc_chunk_merged += f\"Content: {doc_chunk.page_content}\\n Source title: {title}\\n Source link: {source}\\n\\n\"\n",
    "    return f\"Question: {message}\\n {doc_chunk_merged}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26281283-7792-486f-96b0-c781952b6078",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat(message, history):\n",
    "    messages = [{\"role\": \"system\", \"content\": system_message}] + history + [{\"role\": \"user\", \"content\": generate_user_prompt(message)}]\n",
    "    stream = openai.chat.completions.create(model=MODEL, messages=messages, stream=True, seed=3, max_tokens=1000)\n",
    "    response = \"\"\n",
    "    for chunk in stream:\n",
    "        response += chunk.choices[0].delta.content or ''\n",
    "        yield response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44518359-4e55-4679-8460-7406a85cc26f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Testing the Answer generation - 1\n",
    "user_prompt = \"Please explain what Insurellm is in a couple of sentences\"\n",
    "\n",
    "display_handle = display(Markdown(\"\"), display_id=True)\n",
    "for chunk in chat(user_prompt, []):\n",
    "    update_display(Markdown(chunk), display_id=display_handle.display_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6ad1f9d-fd9e-4846-ad83-a33c96f38b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Testing the Answer generation - 2\n",
    "user_prompt = \"Please explain in short on what products are available in Insurellm\"\n",
    "\n",
    "display_handle = display(Markdown(\"\"), display_id=True)\n",
    "for chunk in chat(user_prompt, []):\n",
    "    update_display(Markdown(chunk), display_id=display_handle.display_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d50adb06-e901-41a8-99ed-6f3b4bfacd40",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Launch Gradio\n",
    "gr.ChatInterface(fn=chat, type=\"messages\").launch()"
   ]
  }
 ],
 "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}