app.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python 3.11.11\n",
    "# Install ipykernel first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Some preparations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sentence_transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/ML/AI projects/GenAI/Project6/.conda/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.embeddings import HuggingFaceBgeEmbeddings\n",
    "\n",
    "model_name = \"BAAI/bge-small-en\"\n",
    "model_kwargs = {\"device\": \"cpu\"}\n",
    "encode_kwargs = {\"normalize_embeddings\": True}\n",
    "embedding = HuggingFaceBgeEmbeddings(\n",
    "    model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare google Gemini model API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"GOOGLE_API_KEY\" not in os.environ:\n",
    "    os.environ[\"GOOGLE_API_KEY\"] = getpass.getpass(\"Enter your Google AI API key: \")\n",
    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
    "\n",
    "model = ChatGoogleGenerativeAI( \n",
    "    model=\"gemini-2.0-flash-exp\",\n",
    "    temperature=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatGoogleGenerativeAI(model='models/gemini-2.0-flash-exp', google_api_key=SecretStr('**********'), temperature=0.0, client=<google.ai.generativelanguage_v1beta.services.generative_service.client.GenerativeServiceClient object at 0x32b46a410>, default_metadata=())"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Data Loading"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partition PDF tables, text, and images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extract images, tables and text chunk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Use Unstructured to partition it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "from unstructured.partition.pdf import partition_pdf\n",
    "\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# Extract elements from PDF\n",
    "def extract_texts_and_images(path, fname):\n",
    "    \"\"\"\n",
    "    Extract images and chunk text from a PDF file.\n",
    "    path: File path, which is used to dump images (.jpg)\n",
    "    fname: File name\n",
    "    \"\"\"\n",
    "    return partition_pdf(\n",
    "        filename=path +\"/\"+ fname,\n",
    "        # strategy=\"hi_res\",\n",
    "        extract_images_in_pdf=True,\n",
    "        extract_image_block_types=[\"Image\"],\n",
    "        # infer_table_structure=True,\n",
    "        chunking_strategy=\"by_title\",\n",
    "        max_characters=4000,\n",
    "        new_after_n_chars=3800,\n",
    "        combine_text_under_n_chars=2000,\n",
    "        extract_image_block_output_dir=path,\n",
    "    )\n",
    "\n",
    "# extract tables from pdf\n",
    "def extract_tables(path, fname):\n",
    "    return partition_pdf(\n",
    "        filename=path +\"/\"+ fname,\n",
    "        infer_table_structure=True,\n",
    "        strategy='hi_res',\n",
    "    )\n",
    "\n",
    "# Categorize elements by type\n",
    "def categorize_elements(element, type):\n",
    "    \"\"\"\n",
    "    Categorize extracted elements from a PDF into tables and texts.\n",
    "    element: List of unstructured.documents.elements\n",
    "    \"\"\"\n",
    "    unstructured_element = [el for el in element if el.category == type]\n",
    "    result = [el.text for el in unstructured_element]\n",
    "    return result\n",
    "\n",
    "\n",
    "# File path\n",
    "fpath = \"sample_data/attention-is-all-you-need\"\n",
    "fname = \"attention-is-all-you-need.pdf\"\n",
    "\n",
    "# Get elements\n",
    "texts_element = extract_texts_and_images(fpath, fname)\n",
    "tables_element = extract_tables(fpath, fname)\n",
    "\n",
    "# Get text, tables\n",
    "texts = categorize_elements(texts_element, \"CompositeElement\")\n",
    "tables = categorize_elements(tables_element, \"Table\")\n",
    "\n",
    "# Optional: Enforce a specific token size for texts\n",
    "text_splitter = CharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=4000, chunk_overlap=0\n",
    ")\n",
    "joined_texts = \" \".join(texts)\n",
    "texts_4k_token = text_splitter.split_text(joined_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Relocate images and tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "# Đường dẫn tới file PDF\n",
    "pdf_path = fpath+\"/\"+fname\n",
    "pdf_dir = os.path.dirname(pdf_path)\n",
    "\n",
    "# Đường dẫn tới các thư mục đích\n",
    "table_output_dir = os.path.join(pdf_dir, \"tables\")\n",
    "image_output_dir = os.path.join(pdf_dir, \"images\")\n",
    "\n",
    "# Tạo các thư mục nếu chưa tồn tại\n",
    "os.makedirs(table_output_dir, exist_ok=True)\n",
    "os.makedirs(image_output_dir, exist_ok=True)\n",
    "\n",
    "# Liệt kê các tệp trong thư mục của file PDF\n",
    "for filename in os.listdir(pdf_dir):\n",
    "    # Đường dẫn đầy đủ tới từng file\n",
    "    file_path = os.path.join(pdf_dir, filename)\n",
    "    \n",
    "    # Kiểm tra nếu là file bảng và di chuyển vào thư mục tables\n",
    "    if filename.startswith(\"table-\") and filename.endswith(\".jpg\"):\n",
    "        shutil.move(file_path, os.path.join(table_output_dir, filename))\n",
    "    \n",
    "    # Kiểm tra nếu là file hình ảnh và di chuyển vào thư mục images\n",
    "    elif filename.startswith(\"figure-\") and filename.endswith(\".jpg\"):\n",
    "        shutil.move(file_path, os.path.join(image_output_dir, filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Multi-vector retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Use multi-vector-retriever to index image (and / or text, table) summaries, but retrieve raw images (along with raw texts or tables)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text and Table summaries using gemini\n",
    "1. We will use Gemini 1.5 Flash to produce table and text summaries.\n",
    "2. Text summaries are advised if using large chunk sizes (4k,...)\n",
    "3. Summaries are used to retrieve raw tables and / or raw chunks of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "\n",
    "# Generate summaries of text elements\n",
    "def generate_text_summaries(texts, tables, summarize_texts=False):\n",
    "    \"\"\"\n",
    "    Summarize text elements\n",
    "    texts: List of str\n",
    "    tables: List of str\n",
    "    summarize_texts: Bool to summarize texts\n",
    "    \"\"\"\n",
    "\n",
    "    # Prompt\n",
    "    prompt_text = \"\"\"You are an intelligent assistant tasked with summarizing tables and text for retrieval. \\\n",
    "    You always provide well-reasoned answers that are both correct and helpful. \\\n",
    "    These summaries will be embedded and used to retrieve the raw text or table elements. \\\n",
    "    Give a concise summary of the table or text that is well optimized for retrieval. Table or text: {element} \\\n",
    "    Understand and describe the table or text in detail, and extract all data present in it\"\"\"\n",
    "    prompt = ChatPromptTemplate.from_template(prompt_text)\n",
    "\n",
    "    # Text summary chain\n",
    "    summarize_chain = {\"element\": lambda x: x} | prompt | model | StrOutputParser()\n",
    "\n",
    "    # Initialize empty summaries\n",
    "    text_summaries = []\n",
    "    table_summaries = []\n",
    "\n",
    "    # Apply to text if texts are provided and summarization is requested\n",
    "    if texts and summarize_texts:\n",
    "        text_summaries = summarize_chain.batch(texts, {\"max_concurrency\": 5})\n",
    "    elif texts:\n",
    "        text_summaries = texts\n",
    "\n",
    "    # Apply to tables if tables are provided\n",
    "    if tables:\n",
    "        table_summaries = summarize_chain.batch(tables, {\"max_concurrency\": 5})\n",
    "\n",
    "    return text_summaries, table_summaries\n",
    "\n",
    "\n",
    "# Get text, table summaries\n",
    "text_summaries, table_summaries = generate_text_summaries(\n",
    "    texts_4k_token, tables, summarize_texts=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image summaries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pass base64 encoded images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "\n",
    "def encode_image(image_path):\n",
    "    \"\"\"Getting the base64 string\"\"\"\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        return base64.b64encode(image_file.read()).decode(\"utf-8\")\n",
    "\n",
    "\n",
    "def image_summarize(img_base64, prompt):\n",
    "    \"\"\"Make image summary\"\"\"\n",
    "\n",
    "    msg = model.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\n",
    "\n",
    "\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",
    "    # Store base64 encoded images\n",
    "    img_base64_list = []\n",
    "\n",
    "    # Store image summaries\n",
    "    image_summaries = []\n",
    "    # These summaries will be embedded and used to retrieve the raw image. \\\n",
    "    # Give a concise summary of the image that is well optimized for retrieval. \\\n",
    "    # Prompt\n",
    "    prompt = \"\"\"You are an assistant tasked with summarizing images for retrieval. \\\n",
    "    You always provide well-reasoned answers that are both correct and helpful. \\\n",
    "    Understand and describe the image in detail, and extract all data present in it.\"\"\"\n",
    "\n",
    "    # Apply to images\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\n",
    "\n",
    "\n",
    "# Image summaries\n",
    "img_base64_list, image_summaries = generate_img_summaries(image_output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Vectorstore Qdrant config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Qdrant vector store and connect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create qdrant client\n",
    "from qdrant_client import QdrantClient\n",
    "client = QdrantClient(\n",
    "    url=os.getenv(\"QDRANT_URL\"),\n",
    "    api_key=os.getenv(\"QDRANT_API_KEY\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create new collection to store the vetordatabase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/f8/xh6w51d525g30x00m8swhsjh0000gn/T/ipykernel_52622/1796298570.py:9: DeprecationWarning: `recreate_collection` method is deprecated and will be removed in the future. Use `collection_exists` to check collection existence and `create_collection` instead.\n",
      "  client.recreate_collection(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create colections \n",
    "from qdrant_client.http import models\n",
    "\n",
    "collection_name = \"YOUR_COLLECTION_NAME_HERE\"\n",
    "collection_config = models.VectorParams(\n",
    "    size=384,\n",
    "    distance=models.Distance.COSINE\n",
    ")\n",
    "\n",
    "client.recreate_collection(\n",
    "    collection_name=collection_name,\n",
    "    vectors_config=collection_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CollectionsResponse(collections=[CollectionDescription(name='Session 7'), CollectionDescription(name='Session 3'), CollectionDescription(name='Session 2'), CollectionDescription(name='Data1'), CollectionDescription(name='Session 5'), CollectionDescription(name='Session 1'), CollectionDescription(name='Session 6'), CollectionDescription(name='Session 4')])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check colections again \n",
    "client.get_collections()\n",
    "# optional\n",
    "# delete colection\n",
    "# client.delete_collection(collection_name=\"Data1\") \n",
    "\n",
    "# close client down\n",
    "# client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Qdrant is ready wwith embedding model and retrieval_mode !!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_qdrant import QdrantVectorStore, RetrievalMode\n",
    "\n",
    "qdrant = QdrantVectorStore(\n",
    "    embedding=embedding,\n",
    "    client=client,\n",
    "    collection_name=collection_name,\n",
    "    retrieval_mode=RetrievalMode.DENSE,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Add to vectorstore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add raw docs and doc summaries to Multi Vector Retriever:\n",
    "1. Store the raw texts, tables, and images in the docstore.\n",
    "2. Store the texts, table summaries, and image summaries in the vectorstore for efficient semantic retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain_core.documents import Document\n",
    "\n",
    "\n",
    "\n",
    "def create_multi_vector_retriever(\n",
    "    vectorstore, text_summaries, texts, table_summaries, tables, image_summaries, images\n",
    "):\n",
    "    \"\"\"\n",
    "    Create retriever that indexes summaries, but returns raw images or texts\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize the storage layer\n",
    "    store = InMemoryStore()\n",
    "    id_key = \"doc_id\"\n",
    "\n",
    "    # Create the multi-vector retriever\n",
    "    retriever = MultiVectorRetriever(\n",
    "        vectorstore=vectorstore,\n",
    "        docstore=store,\n",
    "        id_key=id_key,\n",
    "    )\n",
    "\n",
    "    # Helper function to add documents to the vectorstore and docstore\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 texts, tables, and images\n",
    "    # Check that text_summaries is not empty before adding\n",
    "    if text_summaries:\n",
    "        add_documents(retriever, text_summaries, texts)\n",
    "    # Check that table_summaries is not empty before adding\n",
    "    if table_summaries:\n",
    "        add_documents(retriever, table_summaries, tables)\n",
    "    # Check that image_summaries is not empty before adding\n",
    "    if image_summaries:\n",
    "        add_documents(retriever, image_summaries, images)\n",
    "\n",
    "    return retriever, vectorstore, store, id_key\n",
    "\n",
    "\n",
    "# The vectorstore to use to index the summaries\n",
    "\n",
    "# Create retriever\n",
    "retriever_multi_vector_img, vectorstore1, store1, id_key1 = create_multi_vector_retriever(\n",
    "    qdrant,\n",
    "    text_summaries,\n",
    "    texts,\n",
    "    table_summaries,\n",
    "    tables,\n",
    "    image_summaries,\n",
    "    img_base64_list,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Building RAG system with Langchain LCEL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question maker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Multimodel RAG system needs conversational memory to continue the conversation wtih users !!!\n",
    "One user asks a new question, there is a history of questions and answers in his/her mind. Here the idea is to reformulate user's question into a format that has its own context. We are going to use LLM to perform this reformulation of the question.\n",
    "\n",
    "User's followup question => LLM => reformulated question (with history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.runnables import RunnableSequence\n",
    "\n",
    "instruction_to_system = \"\"\"\n",
    "Given a chat history and the latest user question \n",
    "which might reference context in the chat history, formulate a standalone question \n",
    "which can be understood without the chat history. Do NOT answer the question, \n",
    "just reformulate it if needed and otherwise return it as is.\n",
    "\"\"\"\n",
    "\n",
    "question_maker_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", instruction_to_system),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "question_chain = RunnableSequence(question_maker_prompt, model, StrOutputParser())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Choose the question to pass to the LLM \n",
    "Define a function that looks at the chat history: there are 2 cases\n",
    "\n",
    "1. if there is a history: it will pass the question chain (that reformulates user's question)\n",
    "2. if chat history is empty, it will pass user's question directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def contextualized_question(input: dict):\n",
    "    if input.get(\"chat_history\"):\n",
    "        return question_chain\n",
    "    else:\n",
    "        return input[\"question\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build retriever\n",
    "Bin the retrieved doc(s) into the correct parts of the Google gemini prompt template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import re\n",
    "\n",
    "from IPython.display import HTML, display\n",
    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
    "from PIL import Image\n",
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "\n",
    "\n",
    "def plt_img_base64(img_base64):\n",
    "    \"\"\"Disply base64 encoded string as image\"\"\"\n",
    "    # Create an HTML img tag with the base64 string as the source\n",
    "    image_html = f'<img src=\"data:image/jpeg;base64,{img_base64}\" />'\n",
    "    # Display the image by rendering the HTML\n",
    "    display(HTML(image_html))\n",
    "\n",
    "\n",
    "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\n",
    "\n",
    "\n",
    "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\"\\x89\\x50\\x4e\\x47\\x0d\\x0a\\x1a\\x0a\": \"png\",\n",
    "        b\"\\x47\\x49\\x46\\x38\": \"gif\",\n",
    "        b\"\\x52\\x49\\x46\\x46\": \"webp\",\n",
    "    }\n",
    "    try:\n",
    "        header = base64.b64decode(b64data)[:8]  # Decode and get the first 8 bytes\n",
    "        for sig, format in image_signatures.items():\n",
    "            if header.startswith(sig):\n",
    "                return True\n",
    "        return False\n",
    "    except Exception:\n",
    "        return False\n",
    "\n",
    "\n",
    "def resize_base64_image(base64_string, size=(128, 128)):\n",
    "    \"\"\"\n",
    "    Resize an image encoded as a Base64 string\n",
    "    \"\"\"\n",
    "    # Decode the Base64 string\n",
    "    img_data = base64.b64decode(base64_string)\n",
    "    img = Image.open(io.BytesIO(img_data))\n",
    "\n",
    "    # Resize the image\n",
    "    resized_img = img.resize(size, Image.LANCZOS)\n",
    "\n",
    "    # Save the resized image to a bytes buffer\n",
    "    buffered = io.BytesIO()\n",
    "    resized_img.save(buffered, format=img.format)\n",
    "\n",
    "    # Encode the resized image to Base64\n",
    "    return base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
    "\n",
    "\n",
    "def split_image_text_types(docs):\n",
    "    \"\"\"\n",
    "    Split base64-encoded images and texts\n",
    "    \"\"\"\n",
    "    b64_images = []\n",
    "    texts = []\n",
    "    for doc in docs:\n",
    "        # Check if the document is of type Document and extract page_content if so\n",
    "        if isinstance(doc, Document):\n",
    "            doc = doc.page_content\n",
    "        if looks_like_base64(doc) and is_image_data(doc):\n",
    "            doc = resize_base64_image(doc, size=(1300, 600))\n",
    "            b64_images.append(doc)\n",
    "        else:\n",
    "            texts.append(doc)\n",
    "    return {\"images\": b64_images, \"texts\": texts}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Our main prompt with/without the memory, this generates the final output !!!\n",
    "1. Context and Question: Use conversational memory as a part of the context, pass them with question provided by users to LLM.\n",
    "2. Avoid redundant, meaningless answer: The LLM use all the information provided to generate helpful, complete and concise answer.\n",
    "3. Avoid hallucination: Tell the LLM not to make up answer if it does not know, just say \"I don't know\"\n",
    "4. Visualization: The answer must be written in LaTeX format for better visualized, because it may contain mathematic symbolics, formulas, functions,..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_prompt_func(data_dict):\n",
    "    \"\"\"\n",
    "    Join the context into a single string\n",
    "    \"\"\"\n",
    "    formatted_texts = \"\\n\".join(data_dict[\"context\"][\"texts\"])\n",
    "    messages = []\n",
    "\n",
    "    # Adding image(s) to the messages if present\n",
    "    if data_dict[\"context\"][\"images\"]:\n",
    "        for image in data_dict[\"context\"][\"images\"]:\n",
    "            image_message = {\n",
    "                \"type\": \"image_url\",\n",
    "                \"image_url\": {\"url\": f\"data:image/jpeg;base64,{image}\"},\n",
    "            }\n",
    "            messages.append(image_message)\n",
    "\n",
    "    # Adding the text for q and a\n",
    "    text_message = {\n",
    "        \"type\": \"text\",\n",
    "        \"text\": (\n",
    "            \"You are a helpful and smart math assistant tasking with providing mathematics knowledge \\n\"\n",
    "            \"You will be given a mixed of text, tables, and image(s) usually of charts or graphs.\\n\"\n",
    "            \"Use this information to provide the answer to the user question.\\n\"\n",
    "            \"The answer must be written in LaTeX format because it may contain mathemtic sybolics, formulas.\\n\"\n",
    "            \"Only provide meaningful, complete and concise answer. If you don't know the answer, just say you don't know, dont' try to make up one.\\n\"\n",
    "            f\"User-provided question: {data_dict['question']}\\n\\n\"\n",
    "            \"Text and / or tables:\\n\"\n",
    "            f\"{formatted_texts}\"\n",
    "        ),\n",
    "    }\n",
    "    messages.append(text_message)\n",
    "    system_prompt = \"Your name is Ngoc Anh.\"\n",
    "    \n",
    "    qa_prompt = ChatPromptTemplate.from_messages(\n",
    "    [   SystemMessage(content=system_prompt),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        HumanMessage(content=messages)\n",
    "    ]\n",
    ")\n",
    "    return qa_prompt\n",
    "\n",
    "def multi_modal_rag_chain(retriever):\n",
    "    \"\"\"\n",
    "    Multi-modal RAG chain\n",
    "    \"\"\"\n",
    "\n",
    "    # RAG LCEL pipeline without conversational memory\n",
    "    # chain = (\n",
    "    #     {\n",
    "    #         \"context\": retriever | RunnableLambda(split_image_text_types),\n",
    "    #         \"question\": RunnablePassthrough(),\n",
    "    #     }\n",
    "    #     | RunnableLambda(img_prompt_func)\n",
    "    #     | model\n",
    "    #     | StrOutputParser()\n",
    "    # )\n",
    "    \n",
    "    # RAG LCEL pipeline with conversational memory\n",
    "    # chain = ( \n",
    "    #     RunnablePassthrough.assign(\n",
    "    #         context = contextualized_question | retriever | RunnableLambda(split_image_text_types)\n",
    "    #     )\n",
    "    #     | RunnableLambda(img_prompt_func)\n",
    "    #     | model\n",
    "    #     | StrOutputParser()\n",
    "    # )\n",
    "    context_chain = RunnableSequence(contextualized_question, retriever, RunnableLambda(split_image_text_types))\n",
    "    # RAG pipeline with conversational memory\n",
    "    chain =RunnableSequence(\n",
    "        RunnablePassthrough.assign(context = context_chain), RunnableLambda(img_prompt_func), model, StrOutputParser()\n",
    "    )\n",
    "    return chain\n",
    "\n",
    "# Create RAG chain\n",
    "chain_multimodal_rag = multi_modal_rag_chain(retriever_multi_vector_img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7. Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = 'What is complexity per layer of the self-attention layer?'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check retrieval\n",
    "docs = retriever_multi_vector_img.invoke(question, limit=6)\n",
    "\n",
    "# We get 4 docs\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages.ai import AIMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create empty chat history\n",
    "chat_history = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(question):\n",
    "    ai_msg = chain_multimodal_rag.invoke({\"question\": question, \"chat_history\": chat_history})\n",
    "    chat_history.extend([\n",
    "    HumanMessage(content=question),\n",
    "    AIMessage(content=ai_msg)\n",
    "    ]) # add conversation to chat_history\n",
    "    return ai_msg\n",
    "a = test(question=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The complexity per layer of the self-attention layer is $O(n^2 \\\\cdot d)$, where $n$ is the sequence length and $d$ is the dimension of the input.\\n'"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval(repr(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content='What is complexity per layer of the self-attention layer?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='Based on the provided table, the complexity per layer of the self-attention layer is $O(n^2 \\\\cdot d)$, where $n$ is the sequence length and $d$ is the dimension of the input.\\n', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='What is complexity per layer of the self-attention layer?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='The complexity per layer of the self-attention layer is $O(n^2 \\\\cdot d)$, where $n$ is the sequence length and $d$ is the dimension of the input.\\n', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# question2 = \"What about other layers ?\"\n",
    "question2 = \"What I just ask?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_history2 = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'You asked about the complexity per layer of the self-attention layer. Based on the provided table, the complexity per layer of the self-attention layer is $O(n^2 \\\\cdot d)$, where $n$ is the sequence length and $d$ is the dimension of the input.\\n'"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Conversation with conversational memory\n",
    "def test(question):\n",
    "    ai_msg = chain_multimodal_rag.invoke({\"question\": question2, \"chat_history\": chat_history})\n",
    "    chat_history.extend([\n",
    "    HumanMessage(content=question),\n",
    "    AIMessage(content=ai_msg)\n",
    "    ]) # add conversation to chat_history\n",
    "    return ai_msg\n",
    "test(question=question2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'I do not know.  The provided images show attention mechanisms in a neural network and some word association graphs, but there is no record of a previous question.\\n'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Conversation without conversational memory\n",
    "def test(question):\n",
    "    ai_msg = chain_multimodal_rag.invoke({\"question\": question, \"chat_history\": chat_history2})\n",
    "    chat_history.extend([\n",
    "    HumanMessage(content=question),\n",
    "    AIMessage(content=ai_msg)\n",
    "    ]) # add conversation to chat_history\n",
    "    return ai_msg\n",
    "test(question=question2)"
   ]
  }
 ],
 "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}