{
 "cells": [
  {
   "attachments": {
    "9bbbcfe4-2b85-4e76-996a-ce8d1497d34e.png": {
     "image/png": "[encoded data, 406052 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "812a4dbc-fe04-4b84-bdf9-390045e30806",
   "metadata": {},
   "source": [
    "## Multi-modal RAG\n",
    "\n",
    "Many documents contain a mixture of content types, including text and images. \n",
    "\n",
    "Yet, information captured in images is lost in most RAG applications.\n",
    "\n",
    "With the emergence of multimodal LLMs, like [GPT-4V](https://openai.com/research/gpt-4v-system-card), it is worth considering how to utilize images in RAG:\n",
    "\n",
    "`Option 1:` \n",
    "\n",
    "* Use multimodal embeddings (such as [CLIP](https://openai.com/research/clip)) to embed images and text\n",
    "* Retrieve both using similarity search\n",
    "* Pass raw images and text chunks to a multimodal LLM for answer synthesis \n",
    "\n",
    "`Option 2:` \n",
    "\n",
    "* Use a multimodal LLM (such as [GPT-4V](https://openai.com/research/gpt-4v-system-card), [LLaVA](https://llava.hliu.cc/), or [FUYU-8b](https://www.adept.ai/blog/fuyu-8b)) to produce text summaries from images\n",
    "* Embed and retrieve text \n",
    "* Pass text chunks to an LLM for answer synthesis \n",
    "\n",
    "`Option 3`\n",
    "\n",
    "* Use a multimodal LLM (such as [GPT-4V](https://openai.com/research/gpt-4v-system-card), [LLaVA](https://llava.hliu.cc/), or [FUYU-8b](https://www.adept.ai/blog/fuyu-8b)) to produce text summaries from images\n",
    "* Embed and retrieve image summaries with a reference to the raw image \n",
    "* Pass raw images and text chunks to a multimodal LLM for answer synthesis   \n",
    "\n",
    "---\n",
    "\n",
    "This cookbook highlights `Option 3`. \n",
    "\n",
    "* We will use [Unstructured](https://unstructured.io/) to parse images, text, and tables from documents (PDFs).\n",
    "* We will use the [multi-vector retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector) with [Chroma](https://www.trychroma.com/) to store raw text and images along with their summaries for retrieval.\n",
    "* We will use GPT-4V for both image summarization (for retrieval) as well as final answer synthesis from join review of images and texts (or tables).\n",
    "\n",
    "---\n",
    "\n",
    "A separate cookbook highlights `Option 1` [here](https://github.com/langchain-ai/langchain/blob/master/cookbook/multi_modal_RAG_chroma.ipynb).\n",
    "\n",
    "And option `Option 2` is appropriate for cases when a multi-modal LLM cannot be used for answer synthesis (e.g., cost, etc).\n",
    "\n",
    "![ss_mm_rag.png](attachment:9bbbcfe4-2b85-4e76-996a-ce8d1497d34e.png)\n",
    "\n",
    "## Packages\n",
    "\n",
    "In addition to the below pip packages, you will also need `poppler` ([installation instructions](https://pdf2image.readthedocs.io/en/latest/installation.html)) and `tesseract` ([installation instructions](https://tesseract-ocr.github.io/tessdoc/Installation.html)) in your system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f9ee74-395f-4aa4-9695-c00ade01195a",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U langchain openai langchain-chroma langchain-experimental # (newest versions required for multi-modal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140580ef-5db0-43cc-a524-9c39e04d4df0",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install \"unstructured[all-docs]\" pillow pydantic lxml pillow matplotlib chromadb tiktoken"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74b56bde-1ba0-4525-a11d-cab02c5659e4",
   "metadata": {},
   "source": [
    "## Data Loading\n",
    "\n",
    "### Partition PDF tables, text, and images\n",
    "  \n",
    "Let's look at a [popular blog](https://cloudedjudgement.substack.com/p/clouded-judgement-111023) by Jamin Ball.\n",
    "\n",
    "This is a great use-case because much of the information is captured in images (of tables or charts).\n",
    "\n",
    "We use `Unstructured` to partition it (see [blog post](https://blog.langchain.dev/semi-structured-multi-modal-rag/)).\n",
    "\n",
    "---\n",
    "\n",
    "To skip `Unstructured` extraction:\n",
    "\n",
    "[Here](https://drive.google.com/file/d/1QlhGFIFwEkNEjQGOvV_hQe4bnOLDJwCR/view?usp=sharing) is a zip file with a sub-set of the extracted images and pdf.\n",
    "\n",
    "If you want to use the provided folder, then simply opt for a [pdf loader](https://python.langchain.com/docs/modules/data_connection/document_loaders/pdf) for the document:\n",
    "\n",
    "```\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "loader = PyPDFLoader(path + fname)\n",
    "docs = loader.load()\n",
    "tables = [] # Ignore w/ basic pdf loader\n",
    "texts = [d.page_content for d in docs]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59df23c-86f7-4e5d-8b8c-de92a92f6637",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "from unstructured.partition.pdf import partition_pdf\n",
    "\n",
    "\n",
    "# Extract elements from PDF\n",
    "def extract_pdf_elements(path, fname):\n",
    "    \"\"\"\n",
    "    Extract images, tables, 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",
    "        extract_images_in_pdf=False,\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",
    "        image_output_dir_path=path,\n",
    "    )\n",
    "\n",
    "\n",
    "# Categorize elements by type\n",
    "def categorize_elements(raw_pdf_elements):\n",
    "    \"\"\"\n",
    "    Categorize extracted elements from a PDF into tables and texts.\n",
    "    raw_pdf_elements: List of unstructured.documents.elements\n",
    "    \"\"\"\n",
    "    tables = []\n",
    "    texts = []\n",
    "    for element in raw_pdf_elements:\n",
    "        if \"unstructured.documents.elements.Table\" in str(type(element)):\n",
    "            tables.append(str(element))\n",
    "        elif \"unstructured.documents.elements.CompositeElement\" in str(type(element)):\n",
    "            texts.append(str(element))\n",
    "    return texts, tables\n",
    "\n",
    "\n",
    "# File path\n",
    "fpath = \"/Users/rlm/Desktop/cj/\"\n",
    "fname = \"cj.pdf\"\n",
    "\n",
    "# Get elements\n",
    "raw_pdf_elements = extract_pdf_elements(fpath, fname)\n",
    "\n",
    "# Get text, tables\n",
    "texts, tables = categorize_elements(raw_pdf_elements)\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",
   "id": "0aa7f52f-bf5c-4ba4-af72-b2ccba59a4cf",
   "metadata": {},
   "source": [
    "## Multi-vector retriever\n",
    "\n",
    "Use [multi-vector-retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector#summary) to index image (and / or text, table) summaries, but retrieve raw images (along with raw texts or tables).\n",
    "\n",
    "### Text and Table summaries\n",
    "\n",
    "We will use GPT-4 to produce table and, optionall, text summaries.\n",
    "\n",
    "Text summaries are advised if using large chunk sizes (e.g., as set above, we use 4k token chunks).\n",
    "\n",
    "Summaries are used to retrieve raw tables and / or raw chunks of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "523e6ed2-2132-4748-bdb7-db765f20648d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\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 assistant tasked with summarizing tables and text for retrieval. \\\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",
    "    prompt = ChatPromptTemplate.from_template(prompt_text)\n",
    "\n",
    "    # Text summary chain\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4\")\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",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b1feadda-8171-4aed-9a60-320a88dc9ee1",
   "metadata": {},
   "source": [
    "### Image summaries \n",
    "\n",
    "We will use [GPT-4V](https://openai.com/research/gpt-4v-system-card) to produce the image summaries.\n",
    "\n",
    "The API docs [here](https://platform.openai.com/docs/guides/vision):\n",
    "\n",
    "* We pass base64 encoded images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9e6b1d97-4245-45ac-95ba-9bc1cfd10182",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from langchain_core.messages import HumanMessage\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",
    "    chat = ChatOpenAI(model=\"gpt-4-vision-preview\", max_tokens=1024)\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\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",
    "\n",
    "    # Prompt\n",
    "    prompt = \"\"\"You are an assistant tasked with summarizing images for retrieval. \\\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",
    "\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(fpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67b030d4-2ac5-41b6-9245-fc3ba5771d87",
   "metadata": {},
   "source": [
    "### Add to vectorstore\n",
    "\n",
    "Add raw docs and doc summaries to [Multi Vector Retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector#summary): \n",
    "\n",
    "* Store the raw texts, tables, and images in the `docstore`.\n",
    "* Store the texts, table summaries, and image summaries in the `vectorstore` for efficient semantic retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "24a0a289-b970-49fe-b04f-5d857a4c159b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_core.documents import Document\n",
    "from langchain_openai import OpenAIEmbeddings\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\n",
    "\n",
    "\n",
    "# The vectorstore to use to index the summaries\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"mm_rag_cj_blog\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "\n",
    "# Create retriever\n",
    "retriever_multi_vector_img = create_multi_vector_retriever(\n",
    "    vectorstore,\n",
    "    text_summaries,\n",
    "    texts,\n",
    "    table_summaries,\n",
    "    tables,\n",
    "    image_summaries,\n",
    "    img_base64_list,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69060724-e390-4dda-8250-5f86025c874a",
   "metadata": {},
   "source": [
    "## RAG\n",
    "\n",
    "### Build retriever\n",
    "\n",
    "We need to bin the retrieved doc(s) into the correct parts of the GPT-4V prompt template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "771a47fa-1267-4db8-a6ae-5fde48bbc069",
   "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",
    "\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}\n",
    "\n",
    "\n",
    "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 analysis\n",
    "    text_message = {\n",
    "        \"type\": \"text\",\n",
    "        \"text\": (\n",
    "            \"You are financial analyst tasking with providing investment advice.\\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 investment advice related to the user question. \\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",
    "    return [HumanMessage(content=messages)]\n",
    "\n",
    "\n",
    "def multi_modal_rag_chain(retriever):\n",
    "    \"\"\"\n",
    "    Multi-modal RAG chain\n",
    "    \"\"\"\n",
    "\n",
    "    # Multi-modal LLM\n",
    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-vision-preview\", max_tokens=1024)\n",
    "\n",
    "    # RAG pipeline\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",
    "    return chain\n",
    "\n",
    "\n",
    "# Create RAG chain\n",
    "chain_multimodal_rag = multi_modal_rag_chain(retriever_multi_vector_img)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "087ab1e2-fc9a-42af-9d93-a35dd172b130",
   "metadata": {},
   "source": [
    "### Check\n",
    "\n",
    "Examine retrieval; we get back images that are relevant to our question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "9f4695c6-7374-4284-b2fe-a94ac17b630f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check retrieval\n",
    "query = \"Give me company names that are interesting investments based on EV / NTM and NTM rev growth. Consider EV / NTM multiples vs historical?\"\n",
    "docs = retriever_multi_vector_img.invoke(query, limit=6)\n",
    "\n",
    "# We get 4 docs\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "b7a2b0e0-87eb-4e1b-a3f0-067cbf288ef6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check retrieval\n",
    "query = \"What are the EV / NTM and NTM rev growth for MongoDB, Cloudflare, and Datadog?\"\n",
    "docs = retriever_multi_vector_img.invoke(query, limit=6)\n",
    "\n",
    "# We get 4 docs\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "94c74413-9dd7-4337-bdca-05e9ee151f27",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"data:image/jpeg;base64,[encoded data, 164832 characters]\" />"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# We get back relevant images\n",
    "plt_img_base64(docs[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bdfd863-c756-4cb4-b7be-ea00284687d2",
   "metadata": {},
   "source": [
    "### Sanity Check\n",
    "\n",
    "Why does this work? Let's look back at the image that we stored ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "bab422d9-104f-4e47-9760-96bdfdd3a9cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"data:image/jpeg;base64,[encoded data, 164832 characters]\" />"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_img_base64(img_base64_list[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e3afd3b-4482-49af-995a-083a8af8eb57",
   "metadata": {},
   "source": [
    "... here is the corresponding summary, which we embedded and used in similarity search.\n",
    "\n",
    "It's pretty reasonable that this image is indeed retrieved from our `query` based on it's similarity to this summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "bfb944f3-712b-4bdb-9396-6d4afdc62af4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The image is a data table comparing key financial metrics of ten technology companies. Metrics include Enterprise Value to Next Twelve Months Revenue (EV/NTM Rev), EV to 2024 Revenue (EV/2024 Rev), EV to NTM Free Cash Flow (EV/NTM FCF), NTM Revenue Growth, Gross Margin, Operating Margin, Free Cash Flow Margin (FCF Margin), and the percentage in Top 10 Multiple Last Twelve Months (LTM). The table lists averages and medians for these metrics, including an overall median for reference. It features the logo of Altimeter and the watermark \"@jaminball\" at the bottom.'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_summaries[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a60c457c-a675-4689-a6c4-a843f28a9c23",
   "metadata": {},
   "source": [
    "### RAG\n",
    "\n",
    "Now let's run RAG and test the ability to synthesize an answer to our question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "9c64b19e-5a89-4dda-af38-fcc4a36a1b44",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Based on the first image provided, which appears to be a table of financial metrics for various companies, we can extract the following information for MongoDB, Cloudflare, and Datadog:\\n\\nMongoDB:\\n- EV / NTM Rev: 14.6x\\n- NTM Rev Growth: 17%\\n\\nCloudflare:\\n- EV / NTM Rev: 13.4x\\n- NTM Rev Growth: 28%\\n\\nDatadog:\\n- EV / NTM Rev: 13.1x\\n- NTM Rev Growth: 19%\\n\\nThese figures represent the enterprise value to next twelve months' revenue (EV / NTM Rev) multiple and the projected revenue growth for the next twelve months (NTM Rev Growth) for each company. These metrics are often used by investors to assess the valuation and growth prospects of companies, particularly in the technology sector.\""
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run RAG chain\n",
    "chain_multimodal_rag.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea241f1-bd11-45cb-bb33-c4e2e8286855",
   "metadata": {},
   "source": [
    "Here is the trace where we can see what is passed to the LLM:\n",
    " \n",
    "* Question 1 [Trace focused on investment advice](https://smith.langchain.com/public/d77b7b52-4128-4772-82a7-c56eb97e8b97/r)\n",
    "* Question 2 [Trace focused on table extraction](https://smith.langchain.com/public/4624f086-1bd7-4284-9ca9-52fd7e7a4568/r)\n",
    "\n",
    "For question 1, we can see that we pass 3 images along with a text chunk:"
   ]
  },
  {
   "attachments": {
    "2f72d65f-e9b5-4e2e-840a-8d111792d20b.png": {
     "image/png": "[encoded data, 896688 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "2352cfc7-ef05-4257-87f5-32ee0d89ef12",
   "metadata": {},
   "source": [
    "![trace.png](attachment:2f72d65f-e9b5-4e2e-840a-8d111792d20b.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "857e6c08-8798-4159-b9d1-af2f048448b2",
   "metadata": {},
   "source": [
    "### Considerations\n",
    "\n",
    "**Retrieval**\n",
    " \n",
    "* Retrieval is performed based upon similarity to image summaries as well as text chunks.\n",
    "* This requires some careful consideration because image retrieval can fail if there are competing text chunks.\n",
    "* To mitigate this, I produce larger (4k token) text chunks and summarize them for retrieval.\n",
    "\n",
    "**Image Size**\n",
    "\n",
    "* The quality of answer synthesis appears to be sensitive to image size, [as expected](https://platform.openai.com/docs/guides/vision).\n",
    "* I'll do evals soon to test this more carefully."
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
