Gemini_Function_Calling_for_Structured_Data_Responses.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rabimba/GDE-ML-Artifacts/blob/main/Gemini_Function_Calling_for_Structured_Data_Responses.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCn46D8L8KSQ"
      },
      "source": [
        "# Gemini Function Calling for Structured Data Responses\n",
        "\n",
        "## Authentication"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dMABgQFJ7tDj"
      },
      "outputs": [],
      "source": [
        "# Authenticate via Colab\n",
        "import sys\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0VJDoMTt8Kh4"
      },
      "source": [
        "# Define project/location and initialize Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HFb9xehi8J44"
      },
      "outputs": [],
      "source": [
        "# Define project information\n",
        "PROJECT_ID = \"gde-all-rk\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "# Initialize Vertex AI\n",
        "import vertexai\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I2xFPsIz8Kuo"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L6R1AfVm4x7d"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from vertexai.preview.generative_models import (\n",
        "    FunctionDeclaration,\n",
        "    GenerativeModel,\n",
        "    Tool,\n",
        ")\n",
        "\n",
        "model = GenerativeModel(\"gemini-1.0-pro\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-jFg13sm9wMq"
      },
      "source": [
        "## Method 1: Prompting for structured JSON output (without Function Calling)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G6GVmE6X7c8P",
        "outputId": "ccab07de-d928-4d77-ab79-11fa21d4629c"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'title': 'close-up portraits of a hummingbird face',\n",
              "  'description': \"the intricate details of hummingbird's facial features, including its eyes, nose, and mouth\"},\n",
              " {'title': 'wide-angle shots of an hummingbird and their herds in its natural habitat',\n",
              "  'description': \"the vastness of the hummingbird's territory, showcasing its freedom and agility\"},\n",
              " {'title': 'a family of hummingbirds',\n",
              "  'description': 'the relationships within a hummingbird family, capturing their interactions and bonds'},\n",
              " {'title': 'the courtship dance of a pair of hummingbirds',\n",
              "  'description': 'the graceful and delicate movements of hummingbirds during their courtship rituals'},\n",
              " {'title': 'the nest of a hummingbird',\n",
              "  'description': \"the intricate structure of a hummingbird's nest, showcasing their engineering skills\"},\n",
              " {'title': 'a hummingbird in flight',\n",
              "  'description': 'the dynamic and graceful movements of a hummingbird in mid-air'},\n",
              " {'title': 'a hummingbird feeding',\n",
              "  'description': 'the delicate and precise way hummingbirds extract nectar from flowers'},\n",
              " {'title': 'a hummingbird bathing',\n",
              "  'description': 'the playful and refreshing nature of hummingbirds as they bathe in water droplets'},\n",
              " {'title': 'the migration of hummingbirds',\n",
              "  'description': 'the long and arduous journey hummingbirds undertake during migration'},\n",
              " {'title': 'a hummingbird in winter',\n",
              "  'description': 'the adaptations and challenges hummingbirds face during the colder months'}]"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "animal_name = \"hummingbird\"\n",
        "\n",
        "prompt = f\"\"\"\n",
        "You are a professional photographer, and will publish a photo book of\n",
        "{animal_name}. What kind of photos and scenes will you choose to make the photo\n",
        "book attractive and fun that focuses on their facial features, habitats,\n",
        "families and herds, in different scenes and seasons?\n",
        "\n",
        "Output the 10 words of titles and 20 words of detailed description of the\n",
        "photos, in a JSON format like this:\n",
        "\n",
        "[\n",
        "  {{\"title\": \"close-up portraits of a {animal_name} face\", \"description\": \"the intricate\n",
        "details of {animal_name}'s facial features, including its eyes, nose, and mouth\"}},\n",
        "  {{\"title\": \"wide-angle shots of an {animal_name} and their herds in its natural habitat\",\n",
        "  \"description\": \"the vastness of the {animal_name}'s territory, showcasing its freedom and\n",
        "  agility\"}}\n",
        "]\n",
        "\"\"\"\n",
        "\n",
        "response = model.generate_content(\n",
        "  prompt,\n",
        "  generation_config={\n",
        "      \"max_output_tokens\": 2048,\n",
        "      \"temperature\": 0.9,\n",
        "      \"top_p\": 1\n",
        "  },\n",
        ")\n",
        "\n",
        "output = json.loads(response.text)\n",
        "output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TnfBuf0m94jh"
      },
      "source": [
        "## Method 2: Using Function Calling to output structured data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "rcgNx9TPARp7"
      },
      "outputs": [],
      "source": [
        "# @title Helper functions to parse function call response\n",
        "from vertexai.preview.generative_models import GenerativeModel, GenerationResponse\n",
        "from proto.marshal.collections import repeated\n",
        "from proto.marshal.collections import maps\n",
        "\n",
        "def recurse_proto_repeated_composite(repeated_object):\n",
        "    repeated_list = []\n",
        "    for item in repeated_object:\n",
        "        if isinstance(item, repeated.RepeatedComposite):\n",
        "            item = recurse_proto_repeated_composite(item)\n",
        "            repeated_list.append(item)\n",
        "        elif isinstance(item, maps.MapComposite):\n",
        "            item = recurse_proto_marshal_to_dict(item)\n",
        "            repeated_list.append(item)\n",
        "        else:\n",
        "            repeated_list.append(item)\n",
        "\n",
        "    return repeated_list\n",
        "\n",
        "def recurse_proto_marshal_to_dict(marshal_object):\n",
        "    new_dict = {}\n",
        "    for k, v in marshal_object.items():\n",
        "      if not v:\n",
        "        continue\n",
        "      elif isinstance(v, maps.MapComposite):\n",
        "          v = recurse_proto_marshal_to_dict(v)\n",
        "      elif isinstance(v, repeated.RepeatedComposite):\n",
        "          v = recurse_proto_repeated_composite(v)\n",
        "      new_dict[k] = v\n",
        "\n",
        "    return new_dict\n",
        "\n",
        "def get_text(response: GenerationResponse):\n",
        "  \"\"\"Returns the Text from the Generation Response object.\"\"\"\n",
        "  part = response.candidates[0].content.parts[0]\n",
        "  try:\n",
        "    text = part.text\n",
        "  except:\n",
        "    text = None\n",
        "\n",
        "  return text\n",
        "\n",
        "def get_function_name(response: GenerationResponse):\n",
        "  return response.candidates[0].content.parts[0].function_call.name\n",
        "\n",
        "def get_function_args(response: GenerationResponse) -> dict:\n",
        "  return recurse_proto_marshal_to_dict(response.candidates[0].content.parts[0].function_call.args)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0448R_fPCiRn"
      },
      "source": [
        "Define a Function Decalaration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iTkg9MG18ch5"
      },
      "outputs": [],
      "source": [
        "generate_photograph_content_func = FunctionDeclaration(\n",
        "    name=\"generate_photograph_content\",\n",
        "    description=\"\"\"\n",
        "    You are a professional photographer, and will publish a photo book of\n",
        "    an {animal_name}. What kind of photos and scenes will you choose to make the\n",
        "    photo book attractive and fun that focuses on their facial features,\n",
        "    habitats, families and herds, in different scenes and seasons?\n",
        "\n",
        "    Generate about 10 words for the title and 20 words for a detailed\n",
        "    description of the photos\"\"\",\n",
        "    parameters={\n",
        "    \"type\": \"object\",\n",
        "    \"properties\": {\n",
        "      \"photos\": {\n",
        "        \"type\": \"array\",\n",
        "        \"description\": \"A list of photo titles and descriptions\",\n",
        "        \"items\": {\n",
        "          \"description\": \"Photo details\",\n",
        "          \"type\": \"object\",\n",
        "          \"properties\": {\n",
        "            \"title\": {\n",
        "                \"type\": \"string\",\n",
        "                \"description\": \"Title of the photo, such as: Close-up portraits of a {animal_name} face. Or wide-angle shots of an {animal_name} and their herds in its natural habitat.\"\n",
        "              },\n",
        "            \"description\": {\n",
        "              \"type\": \"string\",\n",
        "              \"description\": \"Description of the photo, such as: the intricate details of {animal_name}'s facial features, including its eyes, nose, and mouth. Or the vastness of the {animal_name}'s territory, showcasing its freedom and agility\"\n",
        "            }\n",
        "          }\n",
        "        }\n",
        "      }\n",
        "    },\n",
        "         \"required\": [\n",
        "            \"title\",\n",
        "            \"description\",\n",
        "      ]\n",
        "  },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eC8KCn5J_AYm"
      },
      "source": [
        "Register the function as a `Tool`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0PeGeCJo8sTF"
      },
      "outputs": [],
      "source": [
        "generate_photograph_content_tool = Tool(\n",
        "    function_declarations=[generate_photograph_content_func],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VxqS71uO_Aq3"
      },
      "source": [
        "Send a prompt with the animal name and desired number of photos:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_30YkEnY_JyI"
      },
      "outputs": [],
      "source": [
        "prompt = \"\"\"Generate about 10 options for a photo book about hummingbirds\"\"\"\n",
        "\n",
        "response = model.generate_content(\n",
        "    prompt,\n",
        "    tools=[generate_photograph_content_tool],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c4ijIDJ_Bz8p",
        "outputId": "9d615817-3342-4d96-9b72-da48792b545e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'photos': [{'title': 'Hummingbird in Flight',\n",
              "   'description': 'Close-up of a hummingbird in mid-flight, with its wings extended and its beak pointed towards the ground'}]}"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ],
      "source": [
        "output = get_function_args(response)\n",
        "output"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U2V5HpNaDL3Z",
        "outputId": "7716596c-758d-4b28-9ba3-bc8a3db3b604"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Hummingbird in Flight\n",
            "Close-up of a hummingbird in mid-flight, with its wings extended and its beak pointed towards the ground\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for photo in output['photos']:\n",
        "    print(photo[\"title\"])\n",
        "    print(photo[\"description\"])\n",
        "    print()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "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.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}