notebooks/04_Agent_use_tools_leveraging_NVIDIA_AI_endpoints.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "481943fc",
   "metadata": {},
   "source": [
    "# Multimodal Models from NVIDIA AI Catelog and AI Catalog with LangChain Agent \n",
    "\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "To run this notebook, you need the following:\n",
    "\n",
    "1. Performed the [setup](https://python.langchain.com/docs/integrations/text_embedding/nvidia_ai_endpoints#setup) and generated an API key.\n",
    "\n",
    "2. Installed Python dependencies from [requirements.txt](https://github.com/NVIDIA/GenerativeAIExamples/blob/main/notebooks/requirements.txt).\n",
    "\n",
    "3. Installed additional packages for this notebook: \n",
    "\n",
    "        pip install gradio matplotlib scikit-image\n",
    "\n",
    "\n",
    "This notebook covers the following custom plug-in components:\n",
    "\n",
    "- LLM using **ai-mixtral-8x7b-instruct**\n",
    "    \n",
    "- A NVIDIA AI Catalog **Deplot** as one of the tool\n",
    "\n",
    "- A NVIDIA AI Catalog **Fuyu** as one of the tool\n",
    "    \n",
    "- Gradio as the simply User Interface where we will upload a few images\n",
    "\n",
    "At the end of the day, as below illustrated, we would like to have a UI which allow user to upload image of their choice and have the agent choose tools to do visual reasoning. \n",
    "\n",
    "![interactive UI](./imgs/visual_reasoning.png)    \n",
    "Note: As one can see, since we are using NVIDIA AI Catalog as an API, there is no further requirement in the prerequisites about GPUs as compute hardware\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86843453",
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment the below to install additional python packages.\n",
    "#!pip install unstructured\n",
    "#!pip install matplotlib scikit-image\n",
    "!pip install gradio==3.48.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a20a5967",
   "metadata": {},
   "source": [
    "### Step 1  - Export the NVIDIA_API_KEY\n",
    "You can supply the NVIDIA_API_KEY directly in this notebook when you run the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb05477",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# del os.environ['NVIDIA_API_KEY']  ## delete key and reset\n",
    "if os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    print(\"Valid NVIDIA_API_KEY already in environment. Delete to reset\")\n",
    "else:\n",
    "    nvapi_key = getpass.getpass(\"NVAPI Key (starts with nvapi-): \")\n",
    "    assert nvapi_key.startswith(\"nvapi-\"), f\"{nvapi_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvapi_key\n",
    "global nvapi_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de94bb16",
   "metadata": {},
   "source": [
    "### Step 2 - wrap the Fuyu API call into a function and verify by supplying an image to get a respond"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61212696",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai, httpx, sys\n",
    "\n",
    "import base64, io\n",
    "from PIL import Image\n",
    "import requests, json\n",
    "\n",
    "def fetch_outputs(output):\n",
    "    collect_streaming_outputs=[]\n",
    "    for o in output:\n",
    "        try:\n",
    "            start = o.index('{')\n",
    "            jsonString=o[start:]\n",
    "            d = json.loads(jsonString)\n",
    "            temp=d['choices'][0]['delta']['content']\n",
    "            collect_streaming_outputs.append(temp)\n",
    "        except:\n",
    "            pass\n",
    "    outputs=''.join(collect_streaming_outputs)\n",
    "    return outputs.replace('\\\\','').replace('\\'','')\n",
    "\n",
    "def img2base64_string(img_path):\n",
    "    image = Image.open(img_path)\n",
    "    if image.width > 800 or image.height > 800:\n",
    "        image.thumbnail((800, 800))\n",
    "    buffered = io.BytesIO()\n",
    "    image.convert(\"RGB\").save(buffered, format=\"JPEG\", quality=85)\n",
    "    image_base64 = base64.b64encode(buffered.getvalue()).decode()\n",
    "    return image_base64\n",
    "\n",
    "\n",
    "\n",
    "def fuyu(prompt,img_path):\n",
    "    invoke_url = \"https://ai.api.nvidia.com/v1/vlm/adept/fuyu-8b\"\n",
    "    stream = True\n",
    "    \n",
    "    \n",
    "    image_b64=img2base64_string(img_path)\n",
    "    \n",
    "    \n",
    "    assert len(image_b64) < 200_000, \\\n",
    "      \"To upload larger images, use the assets API (see docs)\"\n",
    "\n",
    "    headers = {\n",
    "      \"Authorization\": f\"Bearer {nvapi_key}\",\n",
    "      \"Accept\": \"text/event-stream\" if stream else \"application/json\"\n",
    "    }\n",
    "    \n",
    "    payload = {\n",
    "      \"messages\": [\n",
    "        {\n",
    "          \"role\": \"user\",\n",
    "          \"content\": f'{prompt} <img src=\"data:image/png;base64,{image_b64}\" />'\n",
    "        }\n",
    "      ],\n",
    "      \"max_tokens\": 1024,\n",
    "      \"temperature\": 0.20,\n",
    "      \"top_p\": 0.70,\n",
    "      \"seed\": 0,\n",
    "      \"stream\": stream\n",
    "    }\n",
    "    \n",
    "    response = requests.post(invoke_url, headers=headers, json=payload)\n",
    "    \n",
    "    if stream:\n",
    "        output=[]\n",
    "        for line in response.iter_lines():\n",
    "            if line:\n",
    "                output.append(line.decode(\"utf-8\"))\n",
    "    else:\n",
    "        output=response.json()\n",
    "    out=fetch_outputs(output)\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b637395f",
   "metadata": {},
   "source": [
    "fetch a test image of a pair of white sneakers and verify the function works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7504a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget \"https://docs.google.com/uc?export=download&id=12ZpBBFkYu-jzz1iz356U5kMikn4uN9ww\" -O ./toy_data/jordan.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6181b9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path=\"./toy_data/jordan.png\"\n",
    "prompt=\"describe the image\"\n",
    "out=fuyu(prompt,img_path)\n",
    "out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bab48356",
   "metadata": {},
   "source": [
    "### Step 3 - we are gonna use mixtral_8x7b model as our main LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d823ef1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test run and see that you can genreate a respond successfully\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "llm = ChatNVIDIA(model=\"ai-mixtral-8x7b-instruct\", nvidia_api_key=nvapi_key, max_tokens=1024)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "288e14ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set up Prerequisites for Image Captioning App User Interface\n",
    "import os\n",
    "import io\n",
    "import IPython.display\n",
    "from PIL import Image\n",
    "import base64\n",
    "import requests\n",
    "import gradio as gr\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8fffc74",
   "metadata": {},
   "source": [
    "### Step 4- wrap Deplot and Fuyu as tools for later usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03ce445a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set up Prerequisites for Image Captioning App User Interface\n",
    "import os\n",
    "import io\n",
    "import IPython.display\n",
    "from PIL import Image\n",
    "import base64\n",
    "import requests\n",
    "import gradio as gr\n",
    "\n",
    "from langchain.tools import BaseTool\n",
    "from transformers import BlipProcessor, BlipForConditionalGeneration, DetrImageProcessor, DetrForObjectDetection\n",
    "from PIL import Image\n",
    "import torch\n",
    "#\n",
    "import os\n",
    "from tempfile import NamedTemporaryFile\n",
    "from langchain.agents import initialize_agent\n",
    "from langchain.chains.conversation.memory import ConversationBufferWindowMemory\n",
    "    \n",
    "def fetch_outputs(output):\n",
    "    collect_streaming_outputs=[]\n",
    "    for o in output:\n",
    "        try:\n",
    "            start = o.index('{')\n",
    "            jsonString=o[start:]\n",
    "            d = json.loads(jsonString)\n",
    "            temp=d['choices'][0]['delta']['content']\n",
    "            collect_streaming_outputs.append(temp)\n",
    "        except:\n",
    "            pass\n",
    "    outputs=''.join(collect_streaming_outputs)\n",
    "    return outputs.replace('\\\\','').replace('\\'','')\n",
    "\n",
    "def img2base64_string(img_path):\n",
    "    image = Image.open(img_path)\n",
    "    if image.width > 800 or image.height > 800:\n",
    "        image.thumbnail((800, 800))\n",
    "    buffered = io.BytesIO()\n",
    "    image.convert(\"RGB\").save(buffered, format=\"JPEG\", quality=85)\n",
    "    image_base64 = base64.b64encode(buffered.getvalue()).decode()\n",
    "    return image_base64\n",
    "\n",
    "\n",
    "class ImageCaptionTool(BaseTool):\n",
    "    name = \"Image captioner from Fuyu\"\n",
    "    description = \"Use this tool when given the path to an image that you would like to be described. \" \\\n",
    "                  \"It will return a simple caption describing the image.\"\n",
    "\n",
    "    def _run(self, img_path):\n",
    "        invoke_url = \"https://ai.api.nvidia.com/v1/vlm/adept/fuyu-8b\"\n",
    "        stream = True\n",
    "\n",
    "\n",
    "        image_b64=img2base64_string(img_path)\n",
    "\n",
    "\n",
    "        assert len(image_b64) < 200_000, \\\n",
    "          \"To upload larger images, use the assets API (see docs)\"\n",
    "        headers = {\n",
    "          \"Authorization\": f\"Bearer {nvapi_key}\",\n",
    "          \"Accept\": \"text/event-stream\" if stream else \"application/json\"\n",
    "        }\n",
    "\n",
    "        payload = {\n",
    "          \"messages\": [\n",
    "            {\n",
    "              \"role\": \"user\",\n",
    "              \"content\": f'what is in this image <img src=\"data:image/png;base64,{image_b64}\" />'\n",
    "            }\n",
    "          ],\n",
    "          \"max_tokens\": 1024,\n",
    "          \"temperature\": 0.20,\n",
    "          \"top_p\": 0.70,\n",
    "          \"seed\": 0,\n",
    "          \"stream\": stream\n",
    "        }\n",
    "\n",
    "        response = requests.post(invoke_url, headers=headers, json=payload)\n",
    "\n",
    "        if stream:\n",
    "            output=[]\n",
    "            for line in response.iter_lines():\n",
    "                if line:\n",
    "                    output.append(line.decode(\"utf-8\"))\n",
    "        else:\n",
    "            output=response.json()\n",
    "        out=fetch_outputs(output)\n",
    "        return out\n",
    "\n",
    "    def _arun(self, query: str):\n",
    "        raise NotImplementedError(\"This tool does not support async\")\n",
    "\n",
    "\n",
    "class TabularPlotTool(BaseTool):\n",
    "    name = \"Tabular Plot reasoning tool\"\n",
    "    description = \"Use this tool when given the path to an image that contain bar, pie chart objects. \" \\\n",
    "                  \"It will extract and return the tabular data \"\n",
    "\n",
    "\n",
    "    def _run(self, img_path):\n",
    "        invoke_url = \"https://ai.api.nvidia.com/v1/vlm/google/deplot\"\n",
    "        stream = True\n",
    "\n",
    "        image_b64=img2base64_string(img_path)\n",
    "\n",
    "        assert len(image_b64) < 180_000, \\\n",
    "          \"To upload larger images, use the assets API (see docs)\"\n",
    "\n",
    "        headers = {\n",
    "          \"Authorization\": f\"Bearer {nvapi_key}\",\n",
    "          \"Accept\": \"text/event-stream\" if stream else \"application/json\"\n",
    "        }\n",
    "\n",
    "        payload = {\n",
    "          \"messages\": [\n",
    "            {\n",
    "              \"role\": \"user\",\n",
    "              \"content\": f'Generate underlying data table of the figure below: <img src=\"data:image/png;base64,{image_b64}\" />'\n",
    "            }\n",
    "          ],\n",
    "          \"max_tokens\": 1024,\n",
    "          \"temperature\": 0.20,\n",
    "          \"top_p\": 0.20,\n",
    "          \"stream\": stream\n",
    "        }\n",
    "\n",
    "        response = requests.post(invoke_url, headers=headers, json=payload)\n",
    "\n",
    "        if stream:\n",
    "            output=[]\n",
    "            for line in response.iter_lines():\n",
    "                if line:\n",
    "                    temp=line.decode(\"utf-8\")\n",
    "                    output.append(temp)\n",
    "                    #print(temp)\n",
    "        else:\n",
    "            output=response.json()\n",
    "        outputs=fetch_outputs(output)\n",
    "        return outputs\n",
    "    def _arun(self, query: str):\n",
    "        raise NotImplementedError(\"This tool does not support async\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "467237c3",
   "metadata": {},
   "source": [
    "### Step 5 - initaite the agent with tools we previously defined "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a95d59ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "#initialize the gent\n",
    "tools = [ImageCaptionTool(),TabularPlotTool()]\n",
    "\n",
    "conversational_memory = ConversationBufferWindowMemory(\n",
    "    memory_key='chat_history',\n",
    "    k=5,\n",
    "    return_messages=True\n",
    ")\n",
    "\n",
    "\n",
    "agent = initialize_agent(\n",
    "    agent=\"chat-conversational-react-description\",\n",
    "    tools=tools,\n",
    "    llm=llm,\n",
    "    max_iterations=5,\n",
    "    verbose=True,\n",
    "    memory=conversational_memory,\n",
    "    handle_parsing_errors=True,\n",
    "    early_stopping_method='generate'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "990470dd",
   "metadata": {},
   "source": [
    "### Step 6 - verify the agent can indeed use the tools with the supplied image and query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf5bfb97",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path=\"./toy_data/jordan.png\"\n",
    "response = agent.invoke({\"input\":f' this is the image path: {img_path}'})\n",
    "print(response['output'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6e66101",
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_agent(img_path):\n",
    "    response = agent.invoke({\"input\":f'this is the image path: {img_path}'})\n",
    "    return response['output']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4062503",
   "metadata": {},
   "source": [
    "### Step 7 - wrap the agent into a simple gradio UI so we can interactively upload arbitrary image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbaaf746",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "ImageCaptionApp = gr.Interface(fn=my_agent,\n",
    "                    inputs=[gr.Image(label=\"Upload image\", type=\"filepath\")],\n",
    "                    outputs=[gr.Textbox(label=\"Caption\")],\n",
    "                    title=\"Image Captioning with langchain agent\",\n",
    "                    description=\"combine langchain agent using tools for image reasoning\",\n",
    "                    allow_flagging=\"never\")\n",
    "\n",
    "ImageCaptionApp.launch(share=True)\n"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}