notebook/agentchat_dalle_and_gpt4v.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2c75da30",
   "metadata": {},
   "source": [
    "# Agent Chat with Multimodal Models: DALLE  and GPT-4V\n",
    "\n",
    "Requires: OpenAI V1. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f51914c",
   "metadata": {},
   "source": [
    "### Before everything starts, install AutoGen with the `lmm` option\n",
    "```bash\n",
    "pip install \"pyautogen[lmm]>=0.2.3\"\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d45964",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "from typing import Any, Optional, Union\n",
    "\n",
    "import PIL\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from diskcache import Cache\n",
    "from openai import OpenAI\n",
    "\n",
    "from autogen import Agent, AssistantAgent, ConversableAgent, LLMConfig, UserProxyAgent\n",
    "from autogen.agentchat.contrib.img_utils import _to_pil, get_image_data, get_pil_image\n",
    "from autogen.agentchat.contrib.multimodal_conversable_agent import MultimodalConversableAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1db6f5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = \"OAI_CONFIG_LIST\"\n",
    "\n",
    "llm_config_4v = LLMConfig.from_json(\n",
    "    path=config_path,\n",
    "    max_tokens=1000,\n",
    ").where(\n",
    "    model=[\"gpt-4-vision-preview\"],\n",
    ")\n",
    "\n",
    "llm_config_dalle = LLMConfig.from_json(\n",
    "    path=config_path,\n",
    ").where(model=[\"dalle\"])\n",
    "\n",
    "\n",
    "gpt4_llm_config = LLMConfig.from_json(path=config_path, cache_seed=42).where(\n",
    "    model=[\"gpt-4\", \"gpt-4-0314\", \"gpt4\", \"gpt-4-32k\", \"gpt-4-32k-0314\", \"gpt-4-32k-v0314\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54de3579",
   "metadata": {},
   "source": [
    "The `llm_config_dalle` should be something like:\n",
    "\n",
    "```python\n",
    "LLMConfig(\n",
    "    config_list=[\n",
    "        {'api_type': 'openai', 'model': 'dalle', 'api_key': '**********', 'api_version': '2024-02-01', 'tags': []}\n",
    "    ]\n",
    ")\n",
    " ```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24850276",
   "metadata": {},
   "source": [
    "## Helper Functions\n",
    "\n",
    "We first create a wrapper for DALLE call, make the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a5e2f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dalle_call(client: OpenAI, model: str, prompt: str, size: str, quality: str, n: int) -> str:\n",
    "    \"\"\"Generate an image using OpenAI's DALL-E model and cache the result.\n",
    "\n",
    "    This function takes a prompt and other parameters to generate an image using OpenAI's DALL-E model.\n",
    "    It checks if the result is already cached; if so, it returns the cached image data. Otherwise,\n",
    "    it calls the DALL-E API to generate the image, stores the result in the cache, and then returns it.\n",
    "\n",
    "    Args:\n",
    "        client (OpenAI): The OpenAI client instance for making API calls.\n",
    "        model (str): The specific DALL-E model to use for image generation.\n",
    "        prompt (str): The text prompt based on which the image is generated.\n",
    "        size (str): The size specification of the image. TODO: This should allow specifying landscape, square, or portrait modes.\n",
    "        quality (str): The quality setting for the image generation.\n",
    "        n (int): The number of images to generate.\n",
    "\n",
    "    Returns:\n",
    "    str: The image data as a string, either retrieved from the cache or newly generated.\n",
    "\n",
    "    Note:\n",
    "    - The cache is stored in a directory named '.cache/'.\n",
    "    - The function uses a tuple of (model, prompt, size, quality, n) as the key for caching.\n",
    "    - The image data is obtained by making a secondary request to the URL provided by the DALL-E API response.\n",
    "    \"\"\"\n",
    "    # Function implementation...\n",
    "    cache = Cache(\".cache/\")  # Create a cache directory\n",
    "    key = (model, prompt, size, quality, n)\n",
    "    if key in cache:\n",
    "        return cache[key]\n",
    "\n",
    "    # If not in cache, compute and store the result\n",
    "    response = client.images.generate(\n",
    "        model=model,\n",
    "        prompt=prompt,\n",
    "        size=size,\n",
    "        quality=quality,\n",
    "        n=n,\n",
    "    )\n",
    "    image_url = response.data[0].url\n",
    "    img_data = get_image_data(image_url)\n",
    "    cache[key] = img_data\n",
    "\n",
    "    return img_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f4f43d",
   "metadata": {},
   "source": [
    "Here is a helper function to extract image from a DALLE agent. We will show the DALLE agent later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fee2643",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_img(agent: Agent) -> PIL.Image:\n",
    "    \"\"\"Extracts an image from the last message of an agent and converts it to a PIL image.\n",
    "\n",
    "    This function searches the last message sent by the given agent for an image tag,\n",
    "    extracts the image data, and then converts this data into a PIL (Python Imaging Library) image object.\n",
    "\n",
    "    Parameters:\n",
    "        agent (Agent): An instance of an agent from which the last message will be retrieved.\n",
    "\n",
    "    Returns:\n",
    "        PIL.Image: A PIL image object created from the extracted image data.\n",
    "\n",
    "    Note:\n",
    "    - The function assumes that the last message contains an <img> tag with image data.\n",
    "    - The image data is extracted using a regular expression that searches for <img> tags.\n",
    "    - It's important that the agent's last message contains properly formatted image data for successful extraction.\n",
    "    - The `_to_pil` function is used to convert the extracted image data into a PIL image.\n",
    "    - If no <img> tag is found, or if the image data is not correctly formatted, the function may raise an error.\n",
    "    \"\"\"\n",
    "    last_message = agent.last_message()[\"content\"]\n",
    "\n",
    "    if isinstance(last_message, str):\n",
    "        img_data = re.findall(\"<img (.*)>\", last_message)[0]\n",
    "    elif isinstance(last_message, list):\n",
    "        # The GPT-4V format, where the content is an array of data\n",
    "        assert isinstance(last_message[0], dict)\n",
    "        img_data = last_message[0][\"image_url\"][\"url\"]\n",
    "\n",
    "    pil_img = get_pil_image(img_data)\n",
    "    return pil_img"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af86569",
   "metadata": {},
   "source": [
    "## The DALLE Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5558caa2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DALLEAgent(ConversableAgent):\n",
    "    def __init__(self, name, llm_config: dict[str, Any], **kwargs: Any):\n",
    "        super().__init__(name, llm_config=llm_config, **kwargs)\n",
    "\n",
    "        try:\n",
    "            config_list = llm_config[\"config_list\"]\n",
    "            api_key = config_list[0][\"api_key\"]\n",
    "        except Exception as e:\n",
    "            print(\"Unable to fetch API Key, because\", e)\n",
    "            api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "        self._dalle_client = OpenAI(api_key=api_key)\n",
    "        self.register_reply([Agent, None], DALLEAgent.generate_dalle_reply)\n",
    "\n",
    "    def send(\n",
    "        self,\n",
    "        message: Union[dict[str, Any], str],\n",
    "        recipient: Agent,\n",
    "        request_reply: Optional[bool] = None,\n",
    "        silent: Optional[bool] = False,\n",
    "    ):\n",
    "        # override and always \"silent\" the send out message;\n",
    "        # otherwise, the print log would be super long!\n",
    "        super().send(message, recipient, request_reply, silent=True)\n",
    "\n",
    "    def generate_dalle_reply(self, messages: Optional[list[dict[str, Any]]], sender: \"Agent\", config):\n",
    "        \"\"\"Generate a reply using OpenAI DALLE call.\"\"\"\n",
    "        client = self._dalle_client if config is None else config\n",
    "        if client is None:\n",
    "            return False, None\n",
    "        if messages is None:\n",
    "            messages = self._oai_messages[sender]\n",
    "\n",
    "        prompt = messages[-1][\"content\"]\n",
    "        # TODO: integrate with autogen.oai. For instance, with caching for the API call\n",
    "        img_data = dalle_call(\n",
    "            client=client,\n",
    "            model=\"dall-e-3\",\n",
    "            prompt=prompt,\n",
    "            size=\"1024x1024\",  # TODO: the size should be flexible, deciding landscape, square, or portrait mode.\n",
    "            quality=\"standard\",\n",
    "            n=1,\n",
    "        )\n",
    "\n",
    "        img_data = _to_pil(img_data)  # Convert to PIL image\n",
    "\n",
    "        # Return the OpenAI message format\n",
    "        return True, {\"content\": [{\"type\": \"image_url\", \"image_url\": {\"url\": img_data}}]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd3f9e56",
   "metadata": {},
   "source": [
    "## Simple Example: Call directly from User"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4095796",
   "metadata": {},
   "outputs": [],
   "source": [
    "dalle = DALLEAgent(name=\"Dalle\", llm_config=llm_config_dalle)\n",
    "\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"User_proxy\", system_message=\"A human admin.\", human_input_mode=\"NEVER\", max_consecutive_auto_reply=0\n",
    ")\n",
    "\n",
    "# Ask the question with an image\n",
    "user_proxy.initiate_chat(\n",
    "    dalle,\n",
    "    message=\"\"\"Create an image with black background, a happy robot is showing a sign with \"I Love AG2\".\"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77ae209",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = extract_img(dalle)\n",
    "\n",
    "plt.imshow(img)\n",
    "plt.axis(\"off\")  # Turn off axis numbers\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e59f055",
   "metadata": {},
   "source": [
    "## Example With Critics: Iterate several times to improve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72214592",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DalleCreator(AssistantAgent):\n",
    "    def __init__(self, n_iters=2, **kwargs):\n",
    "        \"\"\"Initializes a DalleCreator instance.\n",
    "\n",
    "        This agent facilitates the creation of visualizations through a collaborative effort among\n",
    "        its child agents: dalle and critics.\n",
    "\n",
    "        Parameters:\n",
    "            - n_iters (int, optional): The number of \"improvement\" iterations to run. Defaults to 2.\n",
    "            - **kwargs: keyword arguments for the parent AssistantAgent.\n",
    "        \"\"\"\n",
    "        super().__init__(**kwargs)\n",
    "        self.register_reply([Agent, None], reply_func=DalleCreator._reply_user, position=0)\n",
    "        self._n_iters = n_iters\n",
    "\n",
    "    def _reply_user(self, messages=None, sender=None, config=None):\n",
    "        if all((messages is None, sender is None)):\n",
    "            error_msg = f\"Either {messages=} or {sender=} must be provided.\"\n",
    "            logger.error(error_msg)  # noqa: F821\n",
    "            raise AssertionError(error_msg)\n",
    "\n",
    "        if messages is None:\n",
    "            messages = self._oai_messages[sender]\n",
    "\n",
    "        img_prompt = messages[-1][\"content\"]\n",
    "\n",
    "        # Define the agents\n",
    "        self.critics = MultimodalConversableAgent(\n",
    "            name=\"Critics\",\n",
    "            system_message=\"\"\"You need to improve the prompt of the figures you saw.\n",
    "How to create a figure that is better in terms of color, shape, text (clarity), and other things.\n",
    "Reply with the following format:\n",
    "\n",
    "CRITICS: the image needs to improve...\n",
    "PROMPT: here is the updated prompt!\n",
    "\n",
    "\"\"\",\n",
    "            llm_config=llm_config_4v,\n",
    "            human_input_mode=\"NEVER\",\n",
    "            max_consecutive_auto_reply=3,\n",
    "        )\n",
    "\n",
    "        self.dalle = DALLEAgent(name=\"Dalle\", llm_config=llm_config_dalle, max_consecutive_auto_reply=0)\n",
    "\n",
    "        # Data flow begins\n",
    "        self.send(message=img_prompt, recipient=self.dalle, request_reply=True)\n",
    "        img = extract_img(self.dalle)\n",
    "        plt.imshow(img)\n",
    "        plt.axis(\"off\")  # Turn off axis numbers\n",
    "        plt.show()\n",
    "        print(\"Image PLOTTED\")\n",
    "\n",
    "        for i in range(self._n_iters):\n",
    "            # Downsample the image s.t. GPT-4V can take\n",
    "            img = extract_img(self.dalle)\n",
    "            smaller_image = img.resize((128, 128), Image.Resampling.LANCZOS)\n",
    "            smaller_image.save(\"result.png\")\n",
    "\n",
    "            self.msg_to_critics = f\"\"\"Here is the prompt: {img_prompt}.\n",
    "            Here is the figure <img result.png>.\n",
    "            Now, critic and create a prompt so that DALLE can give me a better image.\n",
    "            Show me both \"CRITICS\" and \"PROMPT\"!\n",
    "            \"\"\"\n",
    "            self.send(message=self.msg_to_critics, recipient=self.critics, request_reply=True)\n",
    "            feedback = self._oai_messages[self.critics][-1][\"content\"]\n",
    "            img_prompt = re.findall(\"PROMPT: (.*)\", feedback)[0]\n",
    "\n",
    "            self.send(message=img_prompt, recipient=self.dalle, request_reply=True)\n",
    "            img = extract_img(self.dalle)\n",
    "            plt.imshow(img)\n",
    "            plt.axis(\"off\")  # Turn off axis numbers\n",
    "            plt.show()\n",
    "            print(f\"Image {i} PLOTTED\")\n",
    "\n",
    "        return True, \"result.jpg\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5883009",
   "metadata": {},
   "outputs": [],
   "source": [
    "creator = DalleCreator(\n",
    "    name=\"DALLE Creator!\",\n",
    "    max_consecutive_auto_reply=0,\n",
    "    system_message=\"Help me coordinate generating image\",\n",
    "    llm_config=gpt4_llm_config,\n",
    ")\n",
    "\n",
    "user_proxy = UserProxyAgent(name=\"User\", human_input_mode=\"NEVER\", max_consecutive_auto_reply=0)\n",
    "\n",
    "user_proxy.initiate_chat(\n",
    "    creator, message=\"\"\"Create an image with black background, a happy robot is showing a sign with \"I Love AG2\".\"\"\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Multimodal agent chat with DALL-E and GPT-4v.",
   "tags": [
    "multimodal",
    "gpt-4v"
   ]
  },
  "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": 5
}