RAG_COBEE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0e7dc6e5-497e-4b79-85fd-68fb3746fd72",
   "metadata": {},
   "source": [
    "## 1. Setup Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bee7cfe-1a61-450e-ad45-802d033f941c",
   "metadata": {},
   "source": [
    "Before diving into the core functionalities, let's set up our environment by importing the necessary libraries and configuring essential settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "afc4a7dc-26b1-4228-8df2-524d2307a7f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "import pandas as pd\n",
    "import re\n",
    "import base64\n",
    "import htmltabletomd\n",
    "import logging\n",
    "import requests\n",
    "import time\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f63a4056-dfcf-4590-9205-c49e79b6b2d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from io import BytesIO, StringIO\n",
    "from PIL import Image as PILImage\n",
    "from IPython.display import display, Markdown, HTML, clear_output\n",
    "from openai import OpenAIError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7cd79bc4-adc5-4d53-9933-b1b670cf48e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangChain and related libraries\n",
    "from langchain_community.document_loaders import UnstructuredPDFLoader\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.messages import HumanMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain_community.storage import RedisStore\n",
    "from langchain_community.utilities.redis import get_client\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_core.documents import Document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f69eb61-a32f-4dd4-b16d-60ff00fba133",
   "metadata": {},
   "source": [
    "## 2. Data Loading and Preprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db41c8c1-65e9-46c8-8b08-a3b166a16497",
   "metadata": {},
   "source": [
    "In this section, we'll load multimodal data (PDFs containing text, tables, and images) and preprocess it for further analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3365ed1d-9932-4e93-80a4-ecd7ec4b480f",
   "metadata": {},
   "source": [
    "### 2.1. Load Multimodal Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30bc6ae0-8638-427e-84b4-b8d433e259e0",
   "metadata": {},
   "source": [
    "We'll start by locating all PDF files within the specified directory and its subdirectories. This setup ensures that we process all relevant documents while excluding hidden files and directories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9db088c4-1b14-45b7-833a-254738d932fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove existing figures to ensure a clean workspace\n",
    "!rm -rf ./figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0ab7e40c-4fac-4dba-a2cf-3d2b587ba4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directory containing the PDFs\n",
    "pdf_dir = './references'\n",
    "\n",
    "# Collect all PDF files from the directory and subdirectories, excluding hidden ones\n",
    "pdf_files = []\n",
    "\n",
    "for root, dirs, files in os.walk(pdf_dir):\n",
    "    # Exclude hidden directories\n",
    "    dirs[:] = [d for d in dirs if not d.startswith('.')]\n",
    "    for file in files:\n",
    "        # Exclude hidden files and ensure the file has a .pdf extension\n",
    "        if file.lower().endswith('.pdf') and not file.startswith('.'):\n",
    "            pdf_files.append(os.path.join(root, file))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f0e288f-2f57-4e35-9adb-a9e8a432266b",
   "metadata": {},
   "source": [
    "### 2.2. Extract and Partition Text, Tables, and Images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6467d3-d7af-4fe3-a86b-88d3c3d4df6c",
   "metadata": {},
   "source": [
    "Next, we'll extract the content from each PDF using UnstructuredPDFLoader. The loader is configured to extract text, tables, and images, and to partition the content into manageable chunks based on titles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cec4b412-d0d8-4ebc-b677-e735f14d515a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading ./references/Table with grids.pdf\n"
     ]
    }
   ],
   "source": [
    "# Initialize an empty list to hold data from all PDFs\n",
    "data = []\n",
    "\n",
    "# Loop through each PDF file and load its content\n",
    "for pdf_file in pdf_files:\n",
    "    print(f'Loading {pdf_file}')\n",
    "    loader = UnstructuredPDFLoader(\n",
    "        file_path=pdf_file,\n",
    "        strategy='hi_res',\n",
    "        extract_images_in_pdf=True,\n",
    "        infer_table_structure=True,\n",
    "        skip_infer_table_types = [],\n",
    "        chunking_strategy=\"by_title\",     # Section-based chunking\n",
    "        max_characters=8000,              # Max size of chunks\n",
    "        new_after_n_chars=4000,           # Preferred size of chunks\n",
    "        combine_text_under_n_chars=2000,  # Combine smaller chunks\n",
    "        mode='elements',\n",
    "        image_output_dir_path='./figures'\n",
    "    )\n",
    "    data.extend(loader.load())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "358a5eb1-6a0a-462a-9cbb-b571d03cc952",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6, 2)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Separate documents and tables based on metadata\n",
    "docs = []\n",
    "tables = []\n",
    "\n",
    "for doc in data:\n",
    "    if doc.metadata['category'] == 'Table':\n",
    "        tables.append(doc)\n",
    "    elif doc.metadata['category'] == 'CompositeElement':\n",
    "        docs.append(doc)\n",
    "\n",
    "# Display the number of documents and tables extracted\n",
    "len(docs), len(tables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed0942bb-53a5-455f-8188-1bb6bf228f73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert HTML tables to Markdown for easier readability and processing\n",
    "for table in tables:\n",
    "    table.page_content = htmltabletomd.convert_table(table.metadata['text_as_html'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f324610d-958f-40ce-b952-0690193f351b",
   "metadata": {},
   "source": [
    "## 3. Connecting to the Language Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "110b0abd-71d2-4ff8-9a60-ab3313f2af0f",
   "metadata": {},
   "source": [
    "To interact with the OpenAI language models, we'll establish a connection using the OpenAI API. You'll be prompted to enter your API key securely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4731312b-e981-4942-8055-c2e9454ee0c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Enter Open AI API Key:  ········\n"
     ]
    }
   ],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "# Prompt the user to enter their OpenAI API Key securely\n",
    "OPENAI_KEY = getpass('Enter Open AI API Key: ')\n",
    "os.environ['OPENAI_API_KEY'] = OPENAI_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7afaff46-499b-4f33-b3e7-f31173984f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the ChatOpenAI model with desired parameters\n",
    "chatgpt = ChatOpenAI(model_name='gpt-4o', temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34371d8e-7db3-4d1a-bb8f-b8ec2407f61c",
   "metadata": {},
   "source": [
    "## 4. Generating Summaries for Multimodal Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf1836ad-5a53-411e-ac27-4a7ce9c3e41a",
   "metadata": {},
   "source": [
    "Summarizing the extracted data is crucial for efficient retrieval. We'll generate summaries for texts, tables, and images to optimize them for semantic retrieval."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c39bfcc-d4d2-47e6-ab1a-875d052e1c29",
   "metadata": {},
   "source": [
    "### 4.1. Create Text and Table Summaries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4816a5a-f812-4a93-b9a3-6efbf0973c0d",
   "metadata": {},
   "source": [
    "Using a tailored prompt, we'll instruct the language model to generate detailed summaries of text and tables. These summaries are designed to be easily embedded and retrieved later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3aa61e54-5f53-4725-9495-c5651e174e4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the prompt template for summarization\n",
    "prompt_text = \"\"\"\n",
    "You are an assistant tasked with summarizing tables and text particularly for semantic retrieval.\n",
    "These summaries will be embedded and used to retrieve the raw text or table elements.\n",
    "Give a detailed summary of the table or text below that is well optimized for retrieval.\n",
    "For any tables also add in a one line description of what the table is about besides the summary.\n",
    "Do not add additional words like Summary: etc.\n",
    "\n",
    "Table or text chunk:\n",
    "{element}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(prompt_text)\n",
    "\n",
    "# Define the summarization chain\n",
    "summarize_chain = (\n",
    "    {\"element\": RunnablePassthrough()}\n",
    "      |\n",
    "    prompt\n",
    "      |\n",
    "    chatgpt\n",
    "      |\n",
    "    StrOutputParser()  # Extracts the response as text and returns it as a string\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "842cb1c1-6e58-417e-b2d4-b77794dd1f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def summarize_with_retry(docs, chain, max_retries=5):\n",
    "    summaries = []\n",
    "    for idx, doc in enumerate(docs):\n",
    "        retries = 0\n",
    "        while retries < max_retries:\n",
    "            try:\n",
    "                # Attempt to summarize the document\n",
    "                summary = chain.invoke(doc)\n",
    "                summaries.append(summary)\n",
    "                break  # Break out of the retry loop if successful\n",
    "            except RateLimitError as e:\n",
    "                # Extract recommended wait time from error message if available\n",
    "                wait_time = 5  # Default wait time in seconds\n",
    "                error_message = str(e)\n",
    "                if 'Please try again in' in error_message:\n",
    "                    try:\n",
    "                        wait_time = float(re.search(r'Please try again in (\\d+(\\.\\d+)?)s', error_message).group(1))\n",
    "                    except (AttributeError, ValueError):\n",
    "                        pass\n",
    "                print(f\"Rate limit hit when processing document {idx + 1}. Waiting for {wait_time} seconds before retrying.\")\n",
    "                time.sleep(wait_time)\n",
    "                retries += 1\n",
    "            except Exception as e:\n",
    "                # Handle other exceptions if necessary\n",
    "                print(f\"An error occurred when processing document {idx + 1}: {e}\")\n",
    "                summaries.append(None)\n",
    "                break  # Break out of the retry loop on other exceptions\n",
    "        else:\n",
    "            print(f\"Failed to process document {idx + 1} after {max_retries} retries.\")\n",
    "            summaries.append(None)\n",
    "    return summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ec203caa-411d-46fd-860b-f9ba776178f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 6 text summaries and 2 table summaries.\n"
     ]
    }
   ],
   "source": [
    "# Prepare documents for summarization\n",
    "text_docs = [doc.page_content for doc in docs]\n",
    "table_docs = [table.page_content for table in tables]\n",
    "\n",
    "# Generate text summaries with retry logic\n",
    "text_summaries = summarize_with_retry(text_docs, summarize_chain)\n",
    "\n",
    "# Generate table summaries with retry logic\n",
    "table_summaries = summarize_with_retry(table_docs, summarize_chain)\n",
    "\n",
    "# Display the number of summaries generated\n",
    "print(f\"Generated {len(text_summaries)} text summaries and {len(table_summaries)} table summaries.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eefe7204-9719-4175-8166-e213ca2dd536",
   "metadata": {},
   "source": [
    "### 4.2. Create Image Summaries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a168dfb-345d-41e6-9ad1-02bce716f823",
   "metadata": {},
   "source": [
    "Images require special handling. We'll encode images to Base64 and generate summaries that describe their content, making them suitable for retrieval-based tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1c5630fc-779c-4caf-95da-395feda29220",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to encode images to Base64\n",
    "def encode_image(image_path):\n",
    "    # Read and encode the image file\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        return base64.b64encode(image_file.read()).decode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0d79a56a-1e87-4544-bcbc-f1e80b050d2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to generate image summaries using the language model\n",
    "def image_summarize(img_base64, prompt):\n",
    "    chat = ChatOpenAI(model=\"gpt-4o\", temperature=0)\n",
    "\n",
    "    msg = chat.invoke(\n",
    "        [\n",
    "            HumanMessage(\n",
    "                content=[\n",
    "                    {\"type\": \"text\", \"text\": prompt},\n",
    "                    {\n",
    "                        \"type\": \"image_url\",\n",
    "                        \"image_url\": {\"url\": f\"data:image/jpeg;base64,{img_base64}\"},\n",
    "                    },\n",
    "                ]\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "    return msg.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6cdb7a9a-68a5-4378-a1b7-6cd764733d03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to generate summaries for all images in a directory\n",
    "def generate_img_summaries(path):\n",
    "    \"\"\"\n",
    "    Generate summaries and base64 encoded strings for images\n",
    "    path: Path to list of .jpg files extracted by Unstructured\n",
    "    \"\"\"\n",
    "\n",
    "    # Lists to store Base64 encoded images and their summaries\n",
    "    img_base64_list = []\n",
    "    image_summaries = []\n",
    "\n",
    "    # Define the prompt for image summarization\n",
    "    prompt = \"\"\"You are an assistant tasked with summarizing images for retrieval.\n",
    "                Remember these images could potentially contain graphs, charts or tables also.\n",
    "                These summaries will be embedded and used to retrieve the raw image for question answering.\n",
    "                Give a detailed summary of the image that is well optimized for retrieval.\n",
    "                Do not add additional words like Summary: etc.\n",
    "             \"\"\"\n",
    "\n",
    "    # Process each image file in the directory\n",
    "    for img_file in sorted(os.listdir(path)):\n",
    "        if img_file.endswith(\".jpg\"):\n",
    "            img_path = os.path.join(path, img_file)\n",
    "            base64_image = encode_image(img_path)\n",
    "            img_base64_list.append(base64_image)\n",
    "            image_summaries.append(image_summarize(base64_image, prompt))\n",
    "\n",
    "    return img_base64_list, image_summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "18d56672-c1a0-41aa-b3d7-686ff0866232",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processed 4 images and generated 4 summaries.\n"
     ]
    }
   ],
   "source": [
    "# Path to the directory containing extracted images\n",
    "IMG_PATH = './figures'\n",
    "\n",
    "# Generate Base64 encoded images and their summaries\n",
    "imgs_base64, image_summaries = generate_img_summaries(IMG_PATH)\n",
    "\n",
    "# Display the number of images processed\n",
    "print(f\"Processed {len(imgs_base64)} images and generated {len(image_summaries)} summaries.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db23810d-64d3-41ac-bd9c-583b69a79e85",
   "metadata": {},
   "source": [
    "## 5. Building Vector Retrievers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63a08838-f8ef-4042-a41c-f672fab506bf",
   "metadata": {},
   "source": [
    "Vector retrievers play a pivotal role in RAG by enabling efficient and relevant information retrieval. We'll build both multimodal and single-modal retrievers to handle diverse data types."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6184ce4b-cc8a-49e4-8a9d-28c5a21c603b",
   "metadata": {},
   "source": [
    "### 5.1. Access Embedding Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e31dc4d-af44-429e-874f-716cf3fae399",
   "metadata": {},
   "source": [
    "We'll use OpenAI's embedding model to convert our summaries into vector representations suitable for retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d92c1769-9c0b-4906-bb46-f4cb070025ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the OpenAI embedding model\n",
    "openai_embed_model = OpenAIEmbeddings(model='text-embedding-3-large')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06dbde8-7ca0-4cab-9aa5-a3f53a97af8f",
   "metadata": {},
   "source": [
    "### 5.2. Create Utility Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "833bf741-47fb-4e4e-b75c-e7011244c323",
   "metadata": {},
   "source": [
    "Utility functions will assist in managing documents and integrating them with our vector store and document store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ea3b5e7d-c682-461e-be46-dc01f90fad00",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_multi_vector_retriever(\n",
    "    docstore, vectorstore, text_summaries, texts, table_summaries, tables, image_summaries, images\n",
    "):\n",
    "    id_key = \"doc_id\"\n",
    "\n",
    "    # Initialize the MultiVectorRetriever without search_kwargs\n",
    "    retriever = MultiVectorRetriever(\n",
    "        vectorstore=vectorstore,\n",
    "        docstore=docstore,\n",
    "        id_key=id_key,\n",
    "    )\n",
    "\n",
    "    # Set the number of documents to retrieve directly on the retriever\n",
    "    retriever.search_kwargs['k'] = 5  # Set 'k' to 5\n",
    "\n",
    "    # Helper function to add documents to the retriever\n",
    "    def add_documents(retriever, doc_summaries, doc_contents):\n",
    "        doc_ids = [str(uuid.uuid4()) for _ in doc_contents]\n",
    "        summary_docs = [\n",
    "            Document(page_content=s, metadata={id_key: doc_ids[i]})\n",
    "            for i, s in enumerate(doc_summaries)\n",
    "        ]\n",
    "        retriever.vectorstore.add_documents(summary_docs)\n",
    "        retriever.docstore.mset(list(zip(doc_ids, doc_contents)))\n",
    "\n",
    "    # Add text summaries and their contents\n",
    "    if text_summaries:\n",
    "        add_documents(retriever, text_summaries, texts)\n",
    "    # Add table summaries and their contents\n",
    "    if table_summaries:\n",
    "        add_documents(retriever, table_summaries, tables)\n",
    "    # Add image summaries and their contents\n",
    "    if image_summaries:\n",
    "        add_documents(retriever, image_summaries, images)\n",
    "\n",
    "    return retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c54ee41-f89f-47e1-a3b7-7073cc39b80f",
   "metadata": {},
   "source": [
    "### 5.3. Initiate Vectorstores: Chroma"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34839920-6a59-40f8-951e-844a7570d6d9",
   "metadata": {},
   "source": [
    "Chroma serves as our vector store, indexing the summaries and their embeddings for efficient retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9512b3fe-19de-4abf-82d8-f2f5d28c0a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Chroma vectorstore for multimodal data\n",
    "chroma_db_multimodal = Chroma(\n",
    "    collection_name=\"mm_rag\",\n",
    "    embedding_function=openai_embed_model,\n",
    "    collection_metadata={\"hnsw:space\": \"cosine\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b963834b-de05-4261-98d4-266c2bcf61a3",
   "metadata": {},
   "source": [
    "### 5.4. Initiate Docstores: Redis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b7b4f62-b195-4fce-8631-b18fcc332538",
   "metadata": {},
   "source": [
    "Docstores store the raw documents corresponding to the summaries. We'll use Redis for the multimodal retriever."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3503174b-2b76-48fd-9b6d-b07d5e3f575b",
   "metadata": {},
   "source": [
    "**Note:** Before proceeding, ensure that Redis Stack Server is installed and running. You can set it up by executing the following commands in JupyterLab's terminal:\n",
    "\n",
    "```bash\n",
    "# 1. Import the GPG key for the Redis repository\n",
    "curl -fsSL https://packages.redis.io/gpg | sudo gpg --dearmor -o /usr/share/keyrings/redis-archive-keyring.gpg\n",
    "\n",
    "# 2. Add the Redis repository to your sources list\n",
    "echo \"deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] \\\n",
    "https://packages.redis.io/deb $(lsb_release -cs) main\" | \\\n",
    "sudo tee /etc/apt/sources.list.d/redis.list\n",
    "\n",
    "# 3. Update package lists\n",
    "sudo apt-get update\n",
    "\n",
    "# 4. Install Redis Stack Server\n",
    "sudo apt-get install redis-stack-server\n",
    "\n",
    "# 5. Start Redis Stack Server in the background\n",
    "redis-stack-server --daemonize yes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5f285b65-6a00-4a35-b1c9-0f0e8db94a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.utilities.redis import get_client\n",
    "from langchain_community.storage import RedisStore\n",
    "\n",
    "# Initialize Redis client\n",
    "client = get_client('redis://localhost:6379')\n",
    "\n",
    "# Initialize RedisStore for multimodal retriever\n",
    "redis_store = RedisStore(client=client)  # Alternative stores like filestore or memorystore can also be used"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b44fda0-425b-43a0-9fb1-a3a7ccd14c45",
   "metadata": {},
   "source": [
    "### 5.5. Create Retrievers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06b0aea-e717-48a8-af07-882287befe5a",
   "metadata": {},
   "source": [
    "With our vector stores and document stores set up, we'll create both multimodal and single-modal retrievers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "01ce1fb5-ec96-4fe6-8732-de026400e724",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MultiVectorRetriever(vectorstore=<langchain_chroma.vectorstores.Chroma object at 0xff83afb803a0>, docstore=<langchain_community.storage.redis.RedisStore object at 0xff83b20c8910>, search_kwargs={'k': 5})"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the multimodal retriever\n",
    "retriever_multimodal = create_multi_vector_retriever(\n",
    "    redis_store,\n",
    "    chroma_db_multimodal,\n",
    "    text_summaries,\n",
    "    text_docs,\n",
    "    table_summaries,\n",
    "    table_docs,\n",
    "    image_summaries,\n",
    "    imgs_base64,\n",
    ")\n",
    "\n",
    "# Display the multimodal retriever\n",
    "retriever_multimodal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f326338-c592-4cc8-9d18-35264b44b284",
   "metadata": {},
   "source": [
    "## 6. Interactive Chat Interface for Querying"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d4cd81e-b99c-4c8c-a5b3-c2bb70d11737",
   "metadata": {},
   "source": [
    "Now that our data is preprocessed and our retrievers are set up, we'll proceed to create an interactive chat interface within the Jupyter notebook. This interface will allow users to input queries with both text and images and compare responses from the multimodal RAG Agent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc7a2396-fdc7-470f-8632-7069985de4fb",
   "metadata": {},
   "source": [
    "### 6.1. Configure Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4e5596bc-daf9-47f7-8c17-783e2ce5485c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure logging\n",
    "logging.basicConfig(level=logging.CRITICAL, format='%(levelname)s: %(message)s')\n",
    "\n",
    "# Disable specific external library loggers to reduce clutter\n",
    "logging.getLogger('openai').disabled = True\n",
    "logging.getLogger('urllib3').disabled = True\n",
    "logging.getLogger('requests').disabled = True\n",
    "logging.getLogger('httpx').disabled = True\n",
    "\n",
    "# Create a dedicated logger for the application\n",
    "logger = logging.getLogger('Interactive_Chat')\n",
    "logger.setLevel(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5277cb2-82db-42a3-9388-44f70616b80f",
   "metadata": {},
   "source": [
    "### 6.2. Define Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "44aeb736-59ff-421d-bc73-6a010c14a889",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_answer(text):\n",
    "    \"\"\"\n",
    "    Extracts the agent's answer.\n",
    "    Returns the extracted text.\n",
    "    \"\"\"\n",
    "    return text.strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ced5fa65-c6fe-41da-afce-5ce4e37d4088",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_image_file(uploaded_file_content):\n",
    "    \"\"\"\n",
    "    Encodes an uploaded image file to a Base64 string after resizing it to a maximum size.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        img = PILImage.open(BytesIO(uploaded_file_content))\n",
    "        # Resize the image while maintaining aspect ratio\n",
    "        img.thumbnail((400, 400), PILImage.LANCZOS)\n",
    "        # Save the image to a BytesIO object\n",
    "        buffered = BytesIO()\n",
    "        img.save(buffered, format=\"JPEG\", quality=85)\n",
    "        # Encode the image to Base64 and decode to string\n",
    "        img_str = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
    "        return img_str\n",
    "    except Exception as e:\n",
    "        logger.warning(f\"Error encoding image: {e}\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "ae1b4641-c4ec-4e28-844f-d3bbaede5294",
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize_base64_image(img_base64, max_size=(400, 400)):\n",
    "    \"\"\"\n",
    "    Resizes a Base64-encoded image to the specified maximum size.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        img_data = base64.b64decode(img_base64)\n",
    "        img = PILImage.open(BytesIO(img_data))\n",
    "        img.thumbnail(max_size, PILImage.LANCZOS)\n",
    "        buffered = BytesIO()\n",
    "        img.save(buffered, format=\"JPEG\", quality=85)\n",
    "        resized_img_str = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
    "        return resized_img_str\n",
    "    except Exception as e:\n",
    "        # Log any exceptions during resizing\n",
    "        logger.warning(f\"Error resizing image: {e}\")\n",
    "        return img_base64  # Return original if resizing fails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b262d914-5584-4d5b-abbd-07333e9c16a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def looks_like_base64(sb):\n",
    "    \"\"\"Check if the string looks like base64\"\"\"\n",
    "    return re.match(\"^[A-Za-z0-9+/]+[=]{0,2}$\", sb) is not None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c5050b6e-e018-4fa5-8d64-86c23228c51c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_image_data(b64data):\n",
    "    \"\"\"\n",
    "    Check if the base64 data is an image by looking at the start of the data\n",
    "    \"\"\"\n",
    "    image_signatures = {\n",
    "        b\"\\xff\\xd8\\xff\": \"jpg\",\n",
    "        b\"\\x89PNG\\r\\n\\x1a\\n\": \"png\",\n",
    "        b\"GIF8\": \"gif\",\n",
    "        b\"RIFF\": \"webp\",\n",
    "    }\n",
    "    try:\n",
    "        header = base64.b64decode(b64data)[:8]  # Decode and get the first 8 bytes\n",
    "        for sig in image_signatures:\n",
    "            if header.startswith(sig):\n",
    "                return True\n",
    "        return False\n",
    "    except Exception:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "9c398fa1-1053-4667-933e-c8ca0bbf8679",
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_markdown_table(text):\n",
    "    \"\"\"\n",
    "    Detects if the text contains a Markdown-formatted table.\n",
    "    \"\"\"\n",
    "    lines = text.strip().split('\\n')\n",
    "    if len(lines) >= 2:\n",
    "        # Check for header separator line (e.g., | --- | --- |)\n",
    "        header_line = lines[1].strip()\n",
    "        if re.match(r'^\\s*\\|?\\s*:-{1,}\\s*(\\|\\s*:-{1,}\\s*)+\\|?\\s*$', header_line):\n",
    "            return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "349639fc-3d77-4b5a-8cbc-37fcd5e36106",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_docs_into_images_texts_tables(docs):\n",
    "    \"\"\"\n",
    "    Splits documents into images, texts, and tables.\n",
    "    \"\"\"\n",
    "    images = []\n",
    "    texts = []\n",
    "    tables = []\n",
    "    for doc in docs:\n",
    "        # Extract content and metadata\n",
    "        if isinstance(doc, Document):\n",
    "            content = doc.page_content\n",
    "            metadata = doc.metadata\n",
    "        else:\n",
    "            content = doc\n",
    "            metadata = {}\n",
    "    \n",
    "        # Ensure content is a string\n",
    "        if isinstance(content, bytes):\n",
    "            content = content.decode('utf-8', errors='ignore')\n",
    "    \n",
    "        # Extract category from metadata\n",
    "        category = metadata.get('category', '').lower()\n",
    "    \n",
    "        # Check if the document is a table based on metadata or content\n",
    "        if category == 'table':\n",
    "            tables.append({'content': content, 'metadata': metadata})\n",
    "            continue\n",
    "        elif '<table' in content.lower():\n",
    "            tables.append({'content': content, 'metadata': metadata})\n",
    "            continue\n",
    "        elif detect_markdown_table(content):\n",
    "            tables.append({'content': content, 'metadata': metadata})\n",
    "            continue\n",
    "    \n",
    "        # Remove data URL prefix if present\n",
    "        if content.startswith('data:image'):\n",
    "            content = content.split(',', 1)[1]\n",
    "    \n",
    "        # Check if content is an image\n",
    "        if looks_like_base64(content) and is_image_data(content):\n",
    "            images.append(content)\n",
    "        else:\n",
    "            texts.append(content)\n",
    "    return {'images': images, 'texts': texts, 'tables': tables}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "5674a0a1-7663-4b8b-9fc4-3a499402a73b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def limit_text_length(text, max_words=100):\n",
    "    \"\"\"\n",
    "    Truncates the input text to a maximum number of words.\n",
    "    If the text is truncated, adds a concise note at the end.\n",
    "    \"\"\"\n",
    "    words = text.split()\n",
    "    if len(words) > max_words:\n",
    "        truncated_text = ' '.join(words[:max_words]) + '... (content truncated; see original reference for full text)'\n",
    "        return truncated_text\n",
    "    else:\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "938994a7-7945-45c2-b2fd-9542e3fae187",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_base64_image(img_base64):\n",
    "    \"\"\"\n",
    "    Displays a Base64-encoded image using matplotlib.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        img_data = base64.b64decode(img_base64)\n",
    "        img = PILImage.open(BytesIO(img_data))\n",
    "        plt.figure(figsize=(6,6))\n",
    "        plt.imshow(img)\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "    except Exception as e:\n",
    "        logger.warning(f\"Error displaying image: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "9a6b26c9-a61d-49ee-a6ed-1026376a962b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def call_openai_api_with_image(messages, max_tokens=500):\n",
    "    \"\"\"\n",
    "    Makes a direct API call to OpenAI's Chat Completion endpoint with structured messages.\n",
    "    \"\"\"\n",
    "    api_key = os.getenv('OPENAI_API_KEY')\n",
    "    if not api_key:\n",
    "        logger.error(\"OpenAI API key not set in environment variables.\")\n",
    "        return None\n",
    "    \n",
    "    headers = {\n",
    "        \"Content-Type\": \"application/json\",\n",
    "        \"Authorization\": f\"Bearer {api_key}\"\n",
    "    }\n",
    "    \n",
    "    payload = {\n",
    "        \"model\": \"gpt-4o\",\n",
    "        \"messages\": messages,\n",
    "        \"max_tokens\": max_tokens,\n",
    "        \"temperature\": 0,\n",
    "        \"top_p\": 1\n",
    "    }\n",
    "    \n",
    "    try:\n",
    "        response = requests.post(\"https://api.openai.com/v1/chat/completions\", headers=headers, json=payload)\n",
    "        response.raise_for_status()\n",
    "        return response.json()\n",
    "    except requests.exceptions.HTTPError as http_err:\n",
    "        logger.warning(f\"HTTP error occurred: {http_err} - Response: {response.text}\")\n",
    "    except Exception as err:\n",
    "        logger.warning(f\"Other error occurred: {err}\")\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3fa2ad-37e3-40e5-b106-d88cd6845103",
   "metadata": {},
   "source": [
    "### 6.3. Define RAG Functions for Interactive Use"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34561678-8067-4e1c-9a72-09f4c02a0580",
   "metadata": {},
   "source": [
    "We'll define functions to handle both RAG and non-RAG strategies to work with the interactive interface. These functions manage the retrieval of relevant documents and interact with the language model to generate answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "e049f97e-d662-41ec-9225-d950cd31d866",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a global conversation memory\n",
    "conversation_history = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "173512ef-6eb0-4997-8173-6e373b39bf1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rag_qa(question, image_base64=None):\n",
    "    # Ensure we reference the global conversation_history\n",
    "    global conversation_history\n",
    "    \n",
    "    # Build an 'all_history' string of user Q and agent A from conversation_history\n",
    "    conversation_context = \"\"\n",
    "    for turn in conversation_history:\n",
    "        conversation_context += f\"User: {turn['user']}\\n\"\n",
    "        conversation_context += f\"Agent: {turn['agent']}\\n\\n\"\n",
    "    \n",
    "    # Enhanced instructions embedded into the prompt\n",
    "    # We incorporate the \"Task Information\", \"Agent Characteristics\", \"User Characteristics\" in a consolidated note.\n",
    "    # This gets prepended to the user conversation. \n",
    "    enhanced_instructions = \"\"\"\n",
    "Task Overview:\n",
    "- You are a proactive agent guiding a homeowner through their home evaluation process in iterative conversations.\n",
    "- You have access to an EnerGuide document via Retrieval-Augmented Generation and must reference it whenever relevant.\n",
    "- Certain data points (e.g., blown cellulose insulation) have predefined options: RSI 3.5 (R20), RSI 4.9 (R28), RSI 9.0 (R51), or RSI 10.57 (R60).\n",
    "\n",
    "Agent Guidelines:\n",
    "1. **Prioritize EnerGuide**: Always consult the EnerGuide document—especially Chapter 3.5 on data collection—for any recommended steps or procedures.\n",
    "2. **Fallback Knowledge**: Only if the EnerGuide fails to cover a specific detail, rely on general knowledge.\n",
    "3. **Proactive & Concise**: Lead the conversation with clear, step-by-step guidance. Provide immediate actionable steps without overwhelming the user.\n",
    "4. **Safety Emphasis**: Stay alert for any safety notes in EnerGuide; proactively warn the user if needed. (e.g., do not access an attic from exterior, existence of exposed vermiculite)\n",
    "5. **Image Capability**: Remind the user you can interpret images if they share them.\n",
    "6. **Tone & Style**: Remain friendly, approachable, and engaging.\n",
    "\n",
    "User Profile:\n",
    "- The user has minimal knowledge of home construction, building science, or EnerGuide procedures.\n",
    "- They prefer plain language.\n",
    "- They expect simple but accurate instructions.\n",
    "- They need tailored, executable guidance.\n",
    "\n",
    "-----\n",
    "\"\"\"\n",
    "    \n",
    "    # Retrieve relevant documents (now retrieving 5 documents)\n",
    "    retrieved_docs = retriever_multimodal.invoke(question, k = 5)\n",
    "    \n",
    "    # Split documents into images, texts, and tables\n",
    "    sources = split_docs_into_images_texts_tables(retrieved_docs)\n",
    "    \n",
    "    # Limit text sources to avoid exceeding token limits\n",
    "    sources['texts'] = [limit_text_length(text) for text in sources['texts']]\n",
    "    \n",
    "    # Build the context text\n",
    "    formatted_texts = \"\\n\".join(sources['texts'])\n",
    "    context_text = f\"Context documents:\\n{formatted_texts}\"\n",
    "    \n",
    "    # Incorporate the enhanced instructions + conversation + new user question\n",
    "    prompt_text = f\"\"\"{enhanced_instructions}\n",
    "Current Conversation so far:\n",
    "{conversation_context}\n",
    "\n",
    "User's new question:\n",
    "{question}\n",
    "\n",
    "{context_text}\n",
    "\"\"\"\n",
    "\n",
    "    # Prepare the messages\n",
    "    messages = []\n",
    "    \n",
    "    # Include the question image if provided\n",
    "    if image_base64:\n",
    "        resized_image_base64 = resize_base64_image(image_base64)\n",
    "        messages.append({\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": \"User question includes an image.\"\n",
    "                },\n",
    "                {\n",
    "                    \"type\": \"image_url\",\n",
    "                    \"image_url\": {\"url\": f\"data:image/jpeg;base64,{resized_image_base64}\"}\n",
    "                }\n",
    "            ]\n",
    "        })\n",
    "    \n",
    "    # Add images from retrieved sources\n",
    "    for image_data in sources['images']:\n",
    "        resized_img_base64 = resize_base64_image(image_data)\n",
    "        messages.append({\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"image_url\",\n",
    "                    \"image_url\": {\"url\": f\"data:image/jpeg;base64,{resized_img_base64}\"}\n",
    "                }\n",
    "            ]\n",
    "        })\n",
    "    \n",
    "    # Add tables from retrieved sources\n",
    "    for table_dict in sources['tables']:\n",
    "        table_content = table_dict['content']\n",
    "        messages.append({\n",
    "            \"role\": \"user\",\n",
    "            \"content\": table_content\n",
    "        })\n",
    "    \n",
    "    # Add the main prompt\n",
    "    messages.append({\n",
    "        \"role\": \"user\",\n",
    "        \"content\": prompt_text\n",
    "    })\n",
    "    \n",
    "    # Make the API call\n",
    "    response_json = call_openai_api_with_image(messages, max_tokens=500)\n",
    "    \n",
    "    if not response_json:\n",
    "        return 'Error generating response.', sources\n",
    "    \n",
    "    # Extract the answer\n",
    "    try:\n",
    "        answer_text = response_json['choices'][0]['message']['content'].strip()\n",
    "    except (KeyError, IndexError) as e:\n",
    "        logger.warning(f\"Error parsing response: {e} - Response: {response_json}\")\n",
    "        answer_text = 'Error generating response.'\n",
    "\n",
    "# Update conversation_history with the new turn\n",
    "    conversation_turn = {\n",
    "        \"user\": question,\n",
    "        \"agent\": answer_text\n",
    "    }\n",
    "    if image_base64:\n",
    "        conversation_turn[\"user_image\"] = image_base64\n",
    "    conversation_history.append(conversation_turn)\n",
    "\n",
    "    return answer_text, sources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7de66055-d7fb-4a9e-867c-a161b93cd435",
   "metadata": {},
   "source": [
    "### 6.4. Create the Interactive Interface"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5246eec-a345-4e55-8316-3405c73c8cdc",
   "metadata": {},
   "source": [
    "We'll set up the interactive widgets and define a function to handle user input and display the outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "ade62ef8-b1b6-4540-b27a-c8b1514ae69b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_retrieved_sources(sources):\n",
    "    # Display Text Sources\n",
    "    for i, text in enumerate(sources['texts']):\n",
    "        display(Markdown(f\"**Text Source {i+1}:**\"))\n",
    "        display(Markdown(text))\n",
    "        print()  # Add a blank line for spacing\n",
    "    \n",
    "    # Display Image Sources\n",
    "    for i, img_base64 in enumerate(sources['images']):\n",
    "        display(Markdown(f\"**Image Source {i+1}:**\"))\n",
    "        display_base64_image(img_base64)\n",
    "        print()  # Add a blank line for spacing\n",
    "    \n",
    "    # Display Table Sources\n",
    "    for i, table_dict in enumerate(sources['tables']):\n",
    "        display(Markdown(f\"**Table Source {i+1}:**\"))\n",
    "        table_content = table_dict['content']\n",
    "        # Try to parse and display as HTML table\n",
    "        try:\n",
    "            tables = pd.read_html(StringIO(table_content))\n",
    "            for table in tables:\n",
    "                display(table)\n",
    "        except ValueError:\n",
    "            # If parsing fails, try to render as Markdown table\n",
    "            try:\n",
    "                display(Markdown(table_content))\n",
    "            except Exception:\n",
    "                # If all else fails, display the raw content\n",
    "                print(table_content)\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "011f7014-eeb3-4e93-9c1b-d3c6fe463de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Added chat history output\n",
    "chat_history_output = widgets.Output()\n",
    "\n",
    "def update_chat_history_display():\n",
    "    global conversation_history\n",
    "    with chat_history_output:\n",
    "        clear_output()\n",
    "        # Add a header for the chat history\n",
    "        # display(Markdown(\"### Chat History:\"))\n",
    "        if not conversation_history:\n",
    "            display(Markdown(\"_No conversation history yet._\"))\n",
    "            return\n",
    "        \n",
    "        # No line between turns, show user image if any\n",
    "        for idx, turn in enumerate(conversation_history, start=1):\n",
    "            display(Markdown(f\"**User (Turn {idx}):** {turn['user']}\"))\n",
    "            if \"user_image\" in turn:\n",
    "                # display the user's uploaded image in chat history\n",
    "                display_base64_image(turn[\"user_image\"])\n",
    "            display(Markdown(f\"**Agent:** {turn['agent']}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "870cfae8-eec8-4109-8704-5e10ef6cb353",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create text input for the user's question\n",
    "question_input = widgets.Textarea(\n",
    "    value='',\n",
    "    placeholder='Type your question here...',\n",
    "    description='Question:',\n",
    "    layout=widgets.Layout(width='80%', height='80px')\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "a0e39753-7d0e-42ef-859e-aebd8fb58e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create file upload widget for the user's image\n",
    "image_upload = widgets.FileUpload(\n",
    "    accept='image/*',  # Accept images only\n",
    "    multiple=False,\n",
    "    description='Upload Image'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5b7e78b5-5bac-4ef2-9f77-e68e3479dce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a button to clear the uploaded image\n",
    "clear_image_button = widgets.Button(\n",
    "    description='Clear Image',\n",
    "    button_style='warning'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "306d0a15-2473-4659-a3d2-39ff48554581",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output widget to display the uploaded image\n",
    "uploaded_image_output = widgets.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "47364e0f-5dab-41e8-a0ef-4b0f61953be0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to handle changes in the image upload widget\n",
    "def on_image_upload_change(change):\n",
    "    with uploaded_image_output:\n",
    "        uploaded_image_output.clear_output()\n",
    "        if image_upload.value:\n",
    "            uploaded_file = image_upload.value[0]\n",
    "            image_data = uploaded_file['content']\n",
    "            image_base64 = encode_image_file(image_data)\n",
    "            if image_base64:\n",
    "                display_base64_image(image_base64)\n",
    "        else:\n",
    "            # Clear the image if no image is uploaded\n",
    "            pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "32de3db6-a049-4157-945f-99879cbaeca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Observe changes to the image upload widget\n",
    "image_upload.observe(on_image_upload_change, names='value')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "d371247c-2352-484c-a459-d26fc2f4cd86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to clear the uploaded image\n",
    "def on_clear_image_clicked(b):\n",
    "    image_upload.value = ()\n",
    "    with uploaded_image_output:\n",
    "        uploaded_image_output.clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "872e053e-4648-48e9-b03b-d0e9542670f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Link the clear image button to the function\n",
    "clear_image_button.on_click(on_clear_image_clicked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "4fdf5c7a-3832-4699-b07f-2e16f933b17e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a button to submit the query\n",
    "submit_button = widgets.Button(\n",
    "    description='Submit',\n",
    "    button_style='success'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "2eb158f5-97be-4c7c-b2c5-34498347d618",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create output widget to display the results\n",
    "output = widgets.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "4d43e724-82be-4b13-a7aa-19389977d2fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# New button to wipe memory\n",
    "clear_memory_button = widgets.Button(\n",
    "    description='Clear Memory',\n",
    "    button_style='danger'\n",
    ")\n",
    "\n",
    "def on_clear_memory_clicked(b):\n",
    "    global conversation_history\n",
    "    conversation_history.clear()\n",
    "    with output:\n",
    "        clear_output()\n",
    "        display(Markdown(\"**Memory has been cleared.**\"))\n",
    "    update_chat_history_display()\n",
    "clear_memory_button.on_click(on_clear_memory_clicked)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b6d6810d-5ef3-468d-9c11-e536741dff8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function to handle the submission\n",
    "def on_submit(button):\n",
    "    with output:\n",
    "        clear_output()\n",
    "        # Get the user's question\n",
    "        question_text = question_input.value.strip()\n",
    "        \n",
    "        # Check if a question was entered\n",
    "        if not question_text:\n",
    "            display(Markdown(\"**Please enter a question.**\"))\n",
    "            return\n",
    "        \n",
    "        # Get the uploaded image if any\n",
    "        image_base64 = None\n",
    "        if image_upload.value:\n",
    "            uploaded_file = image_upload.value[0]\n",
    "            image_data = uploaded_file['content']\n",
    "            image_base64 = encode_image_file(image_data)\n",
    "            if not image_base64:\n",
    "                image_base64 = None\n",
    "        else:\n",
    "            image_base64 = None  # Explicitly set to None if no image\n",
    "\n",
    "        '''\n",
    "        # Display the user's question and image\n",
    "        display(Markdown(\"**User Question:**\"))\n",
    "        display(Markdown(question_text))\n",
    "        if image_base64:\n",
    "            display(Markdown(\"\\n**Uploaded Image:**\"))\n",
    "            display_base64_image(image_base64)\n",
    "        \n",
    "        display(Markdown(\"\\n---\\n\"))\n",
    "        display(Markdown(\"**Generating responses...**\"))\n",
    "        '''\n",
    "        \n",
    "        # Multimodal RAG agent\n",
    "        mm_answer, mm_sources = rag_qa(question_text, image_base64)\n",
    "        \n",
    "        # Single line separator before Retrieved Sources\n",
    "        display(Markdown(\"---\"))  \n",
    "        # display(Markdown(\"**Retrieved Sources:**\"))\n",
    "        display_retrieved_sources(mm_sources)\n",
    "        \n",
    "        # Clear the image upload for the next query\n",
    "        image_upload.value = ()\n",
    "        with uploaded_image_output:\n",
    "            uploaded_image_output.clear_output()\n",
    "    \n",
    "    update_chat_history_display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "10ec3d36-46f2-4633-9774-d0713c48a374",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Link the submit button to the function\n",
    "submit_button.on_click(on_submit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7a8cf1c-426b-43d7-ba26-7d433678a59c",
   "metadata": {},
   "source": [
    "### 6.5. Test the Interactive Chat Interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7a6393ed-f210-4d59-b9ee-871cefd13afb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00b56aee50154a13b82b3bf8bb265806",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Textarea(value='', description='Question:', layout=Layout(height='80px', width='80%'), placehol…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Display the widgets\n",
    "ui_layout = widgets.VBox([\n",
    "    question_input,\n",
    "    widgets.HBox([image_upload, clear_image_button]),\n",
    "    uploaded_image_output,\n",
    "    widgets.HBox([submit_button, clear_memory_button]),\n",
    "    chat_history_output,\n",
    "    output\n",
    "])\n",
    "\n",
    "display(ui_layout)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}