{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd835d40",
   "metadata": {},
   "source": [
    "## Multi-modal outputs: Image & Text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa88e03a",
   "metadata": {},
   "source": [
    "This notebook shows how non-text producing tools can be used to create multi-modal agents.\n",
    "\n",
    "This example is limited to text and image outputs and uses UUIDs to transfer content across tools and agents. \n",
    "\n",
    "This example uses Steamship to generate and store generated images. Generated are auth protected by default. \n",
    "\n",
    "You can get your Steamship api key here: https://steamship.com/account/api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0653da01",
   "metadata": {},
   "outputs": [],
   "source": [
    "from steamship import Block, Steamship\n",
    "import re\n",
    "from IPython.display import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f6933033",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import OpenAI\n",
    "from langchain.agents import initialize_agent\n",
    "from langchain.agents import AgentType\n",
    "from langchain.tools import SteamshipImageGenerationTool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "71e51e53",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9fc769d",
   "metadata": {},
   "source": [
    "## Dall-E "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cd177dfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [SteamshipImageGenerationTool(model_name=\"dall-e\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c71b1e46",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrkl = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "603aeb9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m I need to generate an image of a parrot playing soccer.\n",
      "Action: GenerateImage\n",
      "Action Input: A parrot wearing a soccer uniform, kicking a soccer ball.\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mE28BE7C7-D105-41E0-8A5B-2CE21424DFEC\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3m I now have the UUID of the generated image.\n",
      "Final Answer: The UUID of the generated image is E28BE7C7-D105-41E0-8A5B-2CE21424DFEC.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "output = mrkl.run(\"How would you visualize a parot playing soccer?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "25eb4efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_output(output):\n",
    "    \"\"\"Display the multi-modal output from the agent.\"\"\"\n",
    "    UUID_PATTERN = re.compile(\n",
    "        r\"([0-9A-Za-z]{8}-[0-9A-Za-z]{4}-[0-9A-Za-z]{4}-[0-9A-Za-z]{4}-[0-9A-Za-z]{12})\"\n",
    "    )\n",
    "\n",
    "    outputs = UUID_PATTERN.split(output)\n",
    "    outputs = [\n",
    "        re.sub(r\"^\\W+\", \"\", el) for el in outputs\n",
    "    ]  # Clean trailing and leading non-word characters\n",
    "\n",
    "    for output in outputs:\n",
    "        maybe_block_id = UUID_PATTERN.search(output)\n",
    "        if maybe_block_id:\n",
    "            display(Image(Block.get(Steamship(), _id=maybe_block_id.group()).raw()))\n",
    "        else:\n",
    "            print(output, end=\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "082792a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The UUID of the generated image is \n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 1049852 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_output(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e247b2c4",
   "metadata": {},
   "source": [
    "## StableDiffusion "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "315025e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [SteamshipImageGenerationTool(model_name=\"stable-diffusion\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7930064a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrkl = initialize_agent(\n",
    "    tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "611a833d",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = mrkl.run(\"How would you visualize a parot playing soccer?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d7a3edaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_output(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffdf9c53",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
