agent_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "TzTH-nKn5EJy",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@title Imports\n",
    "import sys\n",
    "import agent.utility_classes as util\n",
    "import agent.agent_classes as ac\n",
    "from typing import TypeVar\n",
    "import vertexai\n",
    "\n",
    "import tenacity\n",
    "from vertexai.generative_models import GenerationConfig\n",
    "from vertexai.generative_models import GenerativeModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Vertex AI model ids\n",
    "# https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference\n",
    "DEFAULT_LLM_VERTEX_ID = 'gemini-1.5-pro-002' \n",
    "vertexai.init(project=<your_project_id>, location=\"us-central1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "oeIQSeTuBXSR",
    "tags": []
   },
   "outputs": [],
   "source": [
    "llm_client = util.LLM()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "DQGJ34gOrdy9",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#@title User prompt\n",
    "prompt = 'generate an image of a pig flying in the sky with an eagle' # @param\n",
    "state = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "siYg31Oub-CT",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"A shimmering dragonfly, its wings like stained glass, landed on a dew-kissed blade of grass.  Sunlight, fractured through the droplets, painted fleeting rainbows on the insect's iridescent body. It paused, antennae twitching, sensing the subtle vibrations of the morning. Then, with a whir of its delicate wings, it launched itself into the air, a flash of sapphire against the emerald backdrop of the meadow.  A gentle breeze carried it towards a stand of ancient willows, their weeping branches trailing in the crystal clear stream.  The dragonfly, a tiny jewel in the vastness of nature, continued its silent dance, a fleeting moment of beauty in the unfolding day.\\n\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Ensure the LLM client is running\n",
    "llm_client.generate('generate something')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "z3-CuJPkarXn",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing the agent. This can take a while...\n"
     ]
    }
   ],
   "source": [
    "agent_type = 'attribute' # @param ['belief', 'attribute', 'llm'] {isTemplate: true}\n",
    "print('Initializing the agent. This can take a while...')\n",
    "if agent_type == 'belief':\n",
    "  agent = ac.BeliefAgent(\n",
    "      prompt, llm_client, state=state\n",
    "  )\n",
    "elif agent_type == 'attribute':\n",
    "  agent = ac.AttributeAgent(\n",
    "      prompt, llm_client, state=state\n",
    "  )\n",
    "elif agent_type == 'llm':\n",
    "  agent = ac.LLMAgent(\n",
    "      prompt, llm_client, state=state\n",
    "  )\n",
    "else:\n",
    "  raise ValueError(f'Unknown agent type: {agent_type}')\n",
    "state = agent.state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qoKpldZgjFG9"
   },
   "source": [
    "# Run multi-turn dialogue between user and agent\n",
    "\n",
    "Enter \"\\_print_agent_state\\_\" to print the state of the agent.\n",
    "\n",
    "Enter \"\\_print_new_prompt\\_\" to print the new T2I prompt that incorporates all user provided information so far."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "vXfUW4X66W8k",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User: generate an image of a pig flying in the sky with an eagle\n",
      "Agent: What species of eagle are you thinking of? a. Bald Eagle, b. Golden Eagle, c. Harpy Eagle, d. Philippine Eagle, e. Steller's Sea Eagle f. unkown. If none of these options, what species of eagle are you thinking of?\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "User:  stellar sea eagle\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Updating agent state... This can take a while...\n",
      "\n",
      "\n",
      "\n",
      "Agent: What should the relation pig-eagle be? a. in front of, b. behind, c. left of, d. right of, e. above f. unkown. If none of these options, what should the relation pig-eagle be?\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "User:  behind\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Updating agent state... This can take a while...\n",
      "\n",
      "\n",
      "\n",
      "Agent: What should the relation pig-eagle be? a. in front of, b. behind, c. left of, d. right of, e. above f. unkown. If none of these options, what should the relation pig-eagle be?\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "User:  behind\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Updating agent state... This can take a while...\n",
      "\n",
      "\n",
      "\n",
      "Agent: What should the image style be? a. cartoon, b. photorealistic, c. painting, d. digital art, e. drawing f. unkown. If none of these options, what should the image style be?\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "User:  digital art\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Updating agent state... This can take a while...\n",
      "\n",
      "\n",
      "\n",
      "Agent: What should the time of day be? a. daytime, b. sunset, c. sunrise d. unkown. If none of these options, what should the time of day be?\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "User:  sunset\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Updating agent state... This can take a while...\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('User:', prompt)\n",
    "MAX_TURN = 5\n",
    "for _ in range(MAX_TURN):\n",
    "  action = agent.select_action(verbose=False)\n",
    "  if action is None:\n",
    "    print('Done')\n",
    "    break\n",
    "  print('Agent:', agent.verbalize_action(action))\n",
    "  observation = input('User: ')\n",
    "  if observation == '_print_agent_state_':\n",
    "    print(agent.state)\n",
    "  elif observation == '_print_new_prompt_':\n",
    "    print(agent.prompt)\n",
    "    sys.stdout.flush()\n",
    "  else:\n",
    "    print('\\nUpdating agent state... This can take a while...')\n",
    "    sys.stdout.flush()\n",
    "    agent.transition(action, observation)\n",
    "  print('\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "UrjOLPhNsKtS",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Refined new prompt:\n",
      "A Steller's Sea Eagle soars through the sunset sky, a pig flying behind it. The pig is behind the eagle. This scene is depicted in a digital art style.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"Refined new prompt:\")\n",
    "print(f\"{agent.prompt}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "last_runtime": {
    "build_target": "//learning/deepmind/research/funsearch/colabs:colab",
    "kind": "private"
   },
   "private_outputs": true,
   "provenance": [
    {
     "file_id": "1X4X6Rnap2VQYbvkRiy3KjyKKLdjI-3UH",
     "timestamp": 1729053276421
    },
    {
     "file_id": "/piper/depot/google3/experimental/deepmind/quest_for_agency/colabs/simple_rl_agent.ipynb?workspaceId=wangzi:q4a::citc",
     "timestamp": 1729049901549
    },
    {
     "file_id": "/piper/depot/google3/experimental/deepmind/quest_for_agency/simple_rl_agent.ipynb",
     "timestamp": 1718398702600
    },
    {
     "file_id": "/piper/depot/google3/experimental/deepmind/quest_for_agency/rl_agent.ipynb?workspaceId=wangzi:create::citc",
     "timestamp": 1718145428623
    },
    {
     "file_id": "/piper/depot/google3/experimental/deepmind/quest_for_agency/rl_agent.ipynb?workspaceId=wangzi:fig-export-create-1888-change-4::citc",
     "timestamp": 1718145356407
    },
    {
     "file_id": "/piper/depot/google3/experimental/deepmind/quest_for_agency/GDM_Create_2024_End_to_End_Colab.ipynb",
     "timestamp": 1718143501057
    },
    {
     "file_id": "1SetZhYmKZu8eumMeAz_ncA5SzVVlpqd5",
     "timestamp": 1718087217063
    }
   ],
   "toc_visible": true
  },
  "environment": {
   "kernel": "python3",
   "name": ".m120",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/:m120"
  },
  "kernelspec": {
   "display_name": "Python 3 (Local)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}