site/en/gemini-api/tutorials/extract_structured_data.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8968a502d25e"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "906e07f6e562"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NtX45QCEdPaP"
      },
      "source": [
        "# Extract structured data using function calling"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2tO4fP7FFg2V"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemini-api/tutorials/extract_structured_data\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on Google AI</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemini-api/tutorials/extract_structured_data.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemini-api/tutorials/extract_structured_data.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Szkddw5NScW"
      },
      "source": [
        "In this tutorial you'll work through a structured data extraction example, using the Gemini API to extract the lists of characters, relationships, things, and places from a story."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bvrwRlNPdYDr"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "QyW6x11UQHnx"
      },
      "outputs": [],
      "source": [
        "!pip install -U -q google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "TS9l5igubpHO"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import pathlib\n",
        "import textwrap\n",
        "\n",
        "import google.generativeai as genai\n",
        "\n",
        "\n",
        "from IPython.display import display\n",
        "from IPython.display import Markdown\n",
        "\n",
        "from google.api_core import retry\n",
        "\n",
        "def to_markdown(text):\n",
        "  text = text.replace('•', '  *')\n",
        "  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VmSlTHXxb5pV"
      },
      "source": [
        "Once you have the API key, pass it to the SDK. You can do this in two ways:\n",
        "\n",
        "* Put the key in the `GOOGLE_API_KEY` environment variable (the SDK will automatically pick it up from there).\n",
        "* Pass the key to `genai.configure(api_key=...)`\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "ab9ASynfcIZn"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    # Used to securely store your API key\n",
        "    from google.colab import userdata\n",
        "\n",
        "    # Or use `os.getenv('API_KEY')` to fetch an environment variable.\n",
        "    GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
        "    genai.configure(api_key=GOOGLE_API_KEY)\n",
        "except ImportError:\n",
        "    pass\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K6SdtoJCL4pL"
      },
      "source": [
        "## The example task\n",
        "\n",
        "For this tutorial you'll extract entities from natural language stories. As an\n",
        " example, below is a story written by Gemini."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "0THz95wOL4pL"
      },
      "outputs": [],
      "source": [
        "new_story = False\n",
        "\n",
        "if new_story:\n",
        "  model = genai.GenerativeModel(model_name='models/gemini-1.5-pro-latest')\n",
        "\n",
        "  response = model.generate_content(\"\"\"\n",
        "      Write a long story about a girl with magic backpack, her family, and at\n",
        "      least one other charater. Make sure everyone has names. Don't forget to\n",
        "      describe the contents of the backpack, and where everyone and everything\n",
        "      starts and ends up.\"\"\", request_options={'retry': retry.Retry()})\n",
        "  story = response.text\n",
        "  print(response.candidates[0].citation_metadata)\n",
        "else:\n",
        "  story = \"\"\"In the quaint town of Willow Creek, nestled amidst rolling hills and whispering willows, resided a young girl named Anya. As she stepped out of the creaky wooden door of her modest cottage, her heart skipped a beat with excitement and anticipation. Today was her first day of school, and she couldn't wait to show off her prized possession - a magical backpack.\\n\\nHanded down to her from her grandmother, the backpack was no ordinary satchel. Its soft, emerald-green fabric shimmered with an ethereal glow, and its leather straps held secrets that only Anya knew. Within its capacious interior lay an enchanted world, filled with wonders that would ignite her imagination and change her life forever.\\n\\nAnya's parents, kind-hearted Elise and wise-bearded Edward, bid her farewell with warm embraces. \"Remember, my dear,\" whispered her mother, \"use your magic wisely and for good.\" Her father added, \"Always seek knowledge, and let the backpack be your trusted companion.\"\\n\\nWith a skip in her step, Anya set off towards the town's only schoolhouse. On her way, she passed her best friend, Samuel, a curious and adventurous boy with a mischievous grin. \"Hey, Anya,\" he called out. \"Can I see your backpack?\"\\n\\nAnya hesitated for a moment before unzipping the flap and revealing its contents. Samuel's eyes widened in amazement as he peered inside. There, nestled amidst pencils and notebooks, were a shimmering sword, a book of ancient spells, a tiny compass that always pointed north, and a magical key that could open any lock.\\n\\nTogether, they marveled at the backpack's wonders, promising to keep its secrets safe. As they approached the schoolhouse, Anya noticed a group of older children huddled together, their faces etched with fear. Curiosity getting the better of her, she cautiously approached.\\n\\n\"What's wrong?\" she asked.\\n\\nA tall, lanky boy stepped forward. \"There's a monster in the forest,\" he stammered. \"It's been terrorizing the town, attacking animals and even people.\"\\n\\nAnya's heart sank. The town of Willow Creek was small and peaceful, and the thought of a monster brought a shiver down her spine. She knew she had to do something to protect her family and friends.\\n\\nWithout a moment's hesitation, Anya opened her backpack and retrieved the shimmering sword. With a determined gleam in her eye, she turned to her terrified peers. \"Don't worry,\" she said, her voice steady. \"I'll take care of it.\"\\n\\nWith Samuel close behind her, Anya ventured into the shadowy depths of the forest. The trees seemed to whisper secrets as she passed, and the undergrowth rustled with unseen creatures. As they walked deeper into the forest, the air grew heavy and the ground beneath their feet trembled.\\n\\nSuddenly, they came to a clearing, and there before their eyes was the monster - a massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease. The creature roared, a thunderous sound that shook the forest to its core.\\n\\nFear surged through Anya, but she refused to let it consume her. She drew the sword from its sheath and charged towards the monster. The blade shimmered in the sunlight, and as it struck the beast's hide, a blinding light erupted, enveloping everything in its radiance.\\n\\nWhen the light faded, the monster was gone, and in its place was a pile of shattered crystals. Anya had defeated the creature with the magic of her backpack, proving that even the smallest of objects could hold the greatest of powers.\\n\\nAs she and Samuel returned to the town, they were greeted as heroes. The people of Willow Creek rejoiced, and the legend of Anya, the girl with the magic backpack, was passed down through generations. And so, Anya continued her adventures, using the backpack's wonders to make the world a better place, one magical step at a time.\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "yMnxJqubg759"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "> In the quaint town of Willow Creek, nestled amidst rolling hills and whispering willows, resided a young girl named Anya. As she stepped out of the creaky wooden door of her modest cottage, her heart skipped a beat with excitement and anticipation. Today was her first day of school, and she couldn't wait to show off her prized possession - a magical backpack.\n",
              "> \n",
              "> Handed down to her from her grandmother, the backpack was no ordinary satchel. Its soft, emerald-green fabric shimmered with an ethereal glow, and its leather straps held secrets that only Anya knew. Within its capacious interior lay an enchanted world, filled with wonders that would ignite her imagination and change her life forever.\n",
              "> \n",
              "> Anya's parents, kind-hearted Elise and wise-bearded Edward, bid her farewell with warm embraces. \"Remember, my dear,\" whispered her mother, \"use your magic wisely and for good.\" Her father added, \"Always seek knowledge, and let the backpack be your trusted companion.\"\n",
              "> \n",
              "> With a skip in her step, Anya set off towards the town's only schoolhouse. On her way, she passed her best friend, Samuel, a curious and adventurous boy with a mischievous grin. \"Hey, Anya,\" he called out. \"Can I see your backpack?\"\n",
              "> \n",
              "> Anya hesitated for a moment before unzipping the flap and revealing its contents. Samuel's eyes widened in amazement as he peered inside. There, nestled amidst pencils and notebooks, were a shimmering sword, a book of ancient spells, a tiny compass that always pointed north, and a magical key that could open any lock.\n",
              "> \n",
              "> Together, they marveled at the backpack's wonders, promising to keep its secrets safe. As they approached the schoolhouse, Anya noticed a group of older children huddled together, their faces etched with fear. Curiosity getting the better of her, she cautiously approached.\n",
              "> \n",
              "> \"What's wrong?\" she asked.\n",
              "> \n",
              "> A tall, lanky boy stepped forward. \"There's a monster in the forest,\" he stammered. \"It's been terrorizing the town, attacking animals and even people.\"\n",
              "> \n",
              "> Anya's heart sank. The town of Willow Creek was small and peaceful, and the thought of a monster brought a shiver down her spine. She knew she had to do something to protect her family and friends.\n",
              "> \n",
              "> Without a moment's hesitation, Anya opened her backpack and retrieved the shimmering sword. With a determined gleam in her eye, she turned to her terrified peers. \"Don't worry,\" she said, her voice steady. \"I'll take care of it.\"\n",
              "> \n",
              "> With Samuel close behind her, Anya ventured into the shadowy depths of the forest. The trees seemed to whisper secrets as she passed, and the undergrowth rustled with unseen creatures. As they walked deeper into the forest, the air grew heavy and the ground beneath their feet trembled.\n",
              "> \n",
              "> Suddenly, they came to a clearing, and there before their eyes was the monster - a massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease. The creature roared, a thunderous sound that shook the forest to its core.\n",
              "> \n",
              "> Fear surged through Anya, but she refused to let it consume her. She drew the sword from its sheath and charged towards the monster. The blade shimmered in the sunlight, and as it struck the beast's hide, a blinding light erupted, enveloping everything in its radiance.\n",
              "> \n",
              "> When the light faded, the monster was gone, and in its place was a pile of shattered crystals. Anya had defeated the creature with the magic of her backpack, proving that even the smallest of objects could hold the greatest of powers.\n",
              "> \n",
              "> As she and Samuel returned to the town, they were greeted as heroes. The people of Willow Creek rejoiced, and the legend of Anya, the girl with the magic backpack, was passed down through generations. And so, Anya continued her adventures, using the backpack's wonders to make the world a better place, one magical step at a time."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "to_markdown(story)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zldoIzn-MuLE"
      },
      "source": [
        "## Using Natural language\n",
        "\n",
        "Large language models are a powerfuls multitask tools. Often you can just ask Gemini for what you want, and it will do okay.  \n",
        "\n",
        "The Gemini API doesn't have a JSON mode, so there are a few things to watch for when generating data structures this way:\n",
        "\n",
        "- Sometimes parsing fails.\n",
        "- The schema can't be strictly enforced.\n",
        "\n",
        "You'll solve those problems in the next section. First, try a simple natural language prompt with the schema written out as text. This has not been optimized:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "eStTMD6VL4pM"
      },
      "outputs": [],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    model_name='models/gemini-1.5-pro-latest')\n",
        "\n",
        "response = model.generate_content(\n",
        "  textwrap.dedent(\"\"\"\\\n",
        "    Please return JSON describing the the people, places, things and relationships from this story using the following schema:\n",
        "\n",
        "    {\"people\": list[PERSON], \"places\":list[PLACE], \"things\":list[THING], \"relationships\": list[RELATIONSHIP]}\n",
        "\n",
        "    PERSON = {\"name\": str, \"description\": str, \"start_place_name\": str, \"end_place_name\": str}\n",
        "    PLACE = {\"name\": str, \"description\": str}\n",
        "    THING = {\"name\": str, \"description\": str, \"start_place_name\": str, \"end_place_name\": str}\n",
        "    RELATIONSHIP = {\"person_1_name\": str, \"person_2_name\": str, \"relationship\": str}\n",
        "\n",
        "    All fields are required.\n",
        "\n",
        "    Important: Only return a single piece of valid JSON text.\n",
        "\n",
        "    Here is the story:\n",
        "\n",
        "    \"\"\") + story,\n",
        "  generation_config={'response_mime_type':'application/json'}\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "B0b5zHI3uEBm"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'{\"people\":[{\"name\":\"Anya\",\"description\":\"A young girl with a magical backpack\",\"start_place_name\":\"Willow Creek\",\"end_place_name\":\"Willow Creek\"},{\"name\":\"Elise\",\"description\":\"Anya\\'s kind-hearted mother\",\"start_place_name\":\"Willow Creek\",\"end_place_name\":\"Willow Creek\"},{\"name\":\"Edward\",\"description\":\"Anya\\'s wise-bearded father\",\"start_place_name\":\"Willow Creek\",\"end_place_name\":\"Willow Creek\"},{\"name\":\"Samuel\",\"description\":\"Anya\\'s curious and adventurous best friend\",\"start_place_name\":\"Willow Creek\",\"end_place_name\":\"Willow Creek\"}],\"places\":[{\"name\":\"Willow Creek\",\"description\":\"A quaint town nestled amidst rolling hills and whispering willows\"},{\"name\":\"Anya\\'s cottage\",\"description\":\"A modest cottage in Willow Creek\"},{\"name\":\"Schoolhouse\",\"description\":\"The only schoolhouse in Willow Creek\"},{\"name\":\"Forest\",\"description\":\"A shadowy forest near Willow Creek\"}],\"things\":[{\"name\":\"Magical backpack\",\"description\":\"A shimmering emerald-green backpack with enchanted contents\",\"start_place_name\":\"Anya\\'s cottage\",\"end_place_name\":\"Forest\"},{\"name\":\"Shimmering sword\",\"description\":\"A sword found inside the backpack\",\"start_place_name\":\"Anya\\'s cottage\",\"end_place_name\":\"Forest\"},{\"name\":\"Book of ancient spells\",\"description\":\"A book found inside the backpack\",\"start_place_name\":\"Anya\\'s cottage\",\"end_place_name\":\"Forest\"},{\"name\":\"Tiny compass\",\"description\":\"A compass found inside the backpack\",\"start_place_name\":\"Anya\\'s cottage\",\"end_place_name\":\"Forest\"},{\"name\":\"Magical key\",\"description\":\"A key found inside the backpack\",\"start_place_name\":\"Anya\\'s cottage\",\"end_place_name\":\"Forest\"}],\"relationships\":[{\"person_1_name\":\"Anya\",\"person_2_name\":\"Elise\",\"relationship\":\"Mother-daughter\"},{\"person_1_name\":\"Anya\",\"person_2_name\":\"Edward\",\"relationship\":\"Father-daughter\"},{\"person_1_name\":\"Anya\",\"person_2_name\":\"Samuel\",\"relationship\":\"Best friends\"}]}\\n'"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "response.text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScEdqKq1lhmQ"
      },
      "source": [
        "That returned a json string. Try parsing it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "xSdj50czL4pM"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "    \"people\": [\n",
            "        {\n",
            "            \"name\": \"Anya\",\n",
            "            \"description\": \"A young girl with a magical backpack\",\n",
            "            \"start_place_name\": \"Willow Creek\",\n",
            "            \"end_place_name\": \"Willow Creek\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Elise\",\n",
            "            \"description\": \"Anya's kind-hearted mother\",\n",
            "            \"start_place_name\": \"Willow Creek\",\n",
            "            \"end_place_name\": \"Willow Creek\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Edward\",\n",
            "            \"description\": \"Anya's wise-bearded father\",\n",
            "            \"start_place_name\": \"Willow Creek\",\n",
            "            \"end_place_name\": \"Willow Creek\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Samuel\",\n",
            "            \"description\": \"Anya's curious and adventurous best friend\",\n",
            "            \"start_place_name\": \"Willow Creek\",\n",
            "            \"end_place_name\": \"Willow Creek\"\n",
            "        }\n",
            "    ],\n",
            "    \"places\": [\n",
            "        {\n",
            "            \"name\": \"Willow Creek\",\n",
            "            \"description\": \"A quaint town nestled amidst rolling hills and whispering willows\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Anya's cottage\",\n",
            "            \"description\": \"A modest cottage in Willow Creek\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Schoolhouse\",\n",
            "            \"description\": \"The only schoolhouse in Willow Creek\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Forest\",\n",
            "            \"description\": \"A shadowy forest near Willow Creek\"\n",
            "        }\n",
            "    ],\n",
            "    \"things\": [\n",
            "        {\n",
            "            \"name\": \"Magical backpack\",\n",
            "            \"description\": \"A shimmering emerald-green backpack with enchanted contents\",\n",
            "            \"start_place_name\": \"Anya's cottage\",\n",
            "            \"end_place_name\": \"Forest\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Shimmering sword\",\n",
            "            \"description\": \"A sword found inside the backpack\",\n",
            "            \"start_place_name\": \"Anya's cottage\",\n",
            "            \"end_place_name\": \"Forest\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Book of ancient spells\",\n",
            "            \"description\": \"A book found inside the backpack\",\n",
            "            \"start_place_name\": \"Anya's cottage\",\n",
            "            \"end_place_name\": \"Forest\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Tiny compass\",\n",
            "            \"description\": \"A compass found inside the backpack\",\n",
            "            \"start_place_name\": \"Anya's cottage\",\n",
            "            \"end_place_name\": \"Forest\"\n",
            "        },\n",
            "        {\n",
            "            \"name\": \"Magical key\",\n",
            "            \"description\": \"A key found inside the backpack\",\n",
            "            \"start_place_name\": \"Anya's cottage\",\n",
            "            \"end_place_name\": \"Forest\"\n",
            "        }\n",
            "    ],\n",
            "    \"relationships\": [\n",
            "        {\n",
            "            \"person_1_name\": \"Anya\",\n",
            "            \"person_2_name\": \"Elise\",\n",
            "            \"relationship\": \"Mother-daughter\"\n",
            "        },\n",
            "        {\n",
            "            \"person_1_name\": \"Anya\",\n",
            "            \"person_2_name\": \"Edward\",\n",
            "            \"relationship\": \"Father-daughter\"\n",
            "        },\n",
            "        {\n",
            "            \"person_1_name\": \"Anya\",\n",
            "            \"person_2_name\": \"Samuel\",\n",
            "            \"relationship\": \"Best friends\"\n",
            "        }\n",
            "    ]\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "print(json.dumps(json.loads(response.text), indent=4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TgC_wkHPmkHn"
      },
      "source": [
        "That's relatively simple and often works, but you can potentially make this more strict/robust by defining the schema using the API's function calling feature."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CxMC28LAOfUf"
      },
      "source": [
        "## Use function calling"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x-V6PJn83Kh9"
      },
      "source": [
        "If you haven't gone through the [Function calling basics](https://ai.google.dev/tutorials/function_calling_python_quickstart) tutorial yet, make sure you do that first.\n",
        "\n",
        "With function calling your function and its parameters are described to the API as a `genai.protos.FunctionDeclaration`. In basic cases the SDK can build the `FunctionDeclaration` from the function and its annotations. The SDK doesn't currently handle the description of nested `OBJECT` (`dict`) parameters. So you'll need to define them explicitly, for now."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k83LZ5MCBfTJ"
      },
      "source": [
        "### Define the schema\n",
        "\n",
        "Start by defining `person` as an object with string fields `name`, `description`, `start_place_name`, `end_place_name`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "d6293fed386a"
      },
      "outputs": [],
      "source": [
        "from typing_extensions import TypedDict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "e48ace2a3ded"
      },
      "outputs": [],
      "source": [
        "class Person(TypedDict):\n",
        "    name: str\n",
        "    description: str\n",
        "    start_place_name: str\n",
        "    end_place_name: str"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N6uD63sBBJ3i"
      },
      "source": [
        "Then do the same for each of the entities you're trying to extract:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "7wd3jTqj_bVi"
      },
      "outputs": [],
      "source": [
        "class Place(TypedDict):\n",
        "    name:str\n",
        "    description:str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "45cLwvCd_vg_"
      },
      "outputs": [],
      "source": [
        "class Thing(TypedDict):\n",
        "    name:str\n",
        "    description:str\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "8DdVSZJfADDY"
      },
      "outputs": [],
      "source": [
        "class Relationship(TypedDict):\n",
        "    person_1_name: str\n",
        "    person_2_name: str\n",
        "    relationship: str"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mJwqEUqjBToJ"
      },
      "source": [
        "Now build the `FunctionDeclaration`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "b1fe9734d8c3"
      },
      "outputs": [],
      "source": [
        "def add_to_database(\n",
        "    people: list[Person],\n",
        "    places: list[Place],\n",
        "    things: list[Thing],\n",
        "    relationships: list[Relationship]\n",
        "):\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1_QSwD9Bmy5"
      },
      "source": [
        "### Call the API\n",
        "\n",
        "Like you saw in [Function calling basics](https://ai.google.dev/tutorials/function_calling_python_quickstart) now you can pass this `FunctionDeclaration` to the `tools` argument of the `genai.GenerativeModel` constructor (the constructor would also accept an equivalent JSON representation of the function declaration):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "5PGAPRDJP4Qx"
      },
      "outputs": [],
      "source": [
        "model = genai.GenerativeModel(\n",
        "    model_name='models/gemini-1.5-pro-latest',\n",
        "    tools = [add_to_database])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1uTYW5cVCDST"
      },
      "source": [
        "Each time you call the API the SDK will send the tools along with your prompt, and the model should call that function you defined:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "bAPA7fNtSUwN"
      },
      "outputs": [],
      "source": [
        "result = model.generate_content(f\"\"\"\n",
        "{story}\n",
        "\n",
        "Please add the people, places, things, and relationships from this story to the database:\n",
        "\"\"\",\n",
        "# Force a function call\n",
        "# tool_config={'function_calling_config':'ANY'}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oSG7r6IBCL7S"
      },
      "source": [
        "Now there is no text to parse. The result _is_ a datastructure."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "07n3wXzFOZ4x"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "'text' in result.candidates[0].content.parts[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "i-8hm1HPI5Ce"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "'function_call' in result.candidates[0].content.parts[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "n8BTs6ogDEkq"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'google.ai.generativelanguage_v1beta.types.content.FunctionCall'>\n"
          ]
        }
      ],
      "source": [
        "fc = result.candidates[0].content.parts[0].function_call\n",
        "print(type(fc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kILNHmG2IED3"
      },
      "source": [
        "The `genai.protos.FunctionCall` class is based on Google Protocol Buffers, convert it to a more familiar JSON compatible object:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "5GKHtT4-F3qa"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{\n",
            "    \"name\": \"add_to_database\",\n",
            "    \"args\": {\n",
            "        \"people\": [\n",
            "            {\n",
            "                \"description\": \"A young girl with a magical backpack.\",\n",
            "                \"name\": \"Anya\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"Anya\\\\'s kind-hearted mother.\",\n",
            "                \"name\": \"Elise\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"Anya\\\\'s wise-bearded father.\",\n",
            "                \"name\": \"Edward\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"Anya\\\\'s curious and adventurous best friend.\",\n",
            "                \"name\": \"Samuel\"\n",
            "            }\n",
            "        ],\n",
            "        \"things\": [\n",
            "            {\n",
            "                \"description\": \"A backpack with magical contents, including a shimmering sword, a book of ancient spells, a tiny compass, and a magical key.\",\n",
            "                \"name\": \"Magical backpack\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"A sword that emits a blinding light when it strikes.\",\n",
            "                \"name\": \"Shimmering sword\"\n",
            "            }\n",
            "        ],\n",
            "        \"places\": [\n",
            "            {\n",
            "                \"description\": \"A quaint town nestled amidst rolling hills and whispering willows.\",\n",
            "                \"name\": \"Willow Creek\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"A modest cottage in Willow Creek.\",\n",
            "                \"name\": \"Anya\\\\'s cottage\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"The town\\\\'s only schoolhouse.\",\n",
            "                \"name\": \"Schoolhouse\"\n",
            "            },\n",
            "            {\n",
            "                \"description\": \"A shadowy forest near Willow Creek where the monster lived.\",\n",
            "                \"name\": \"Forest\"\n",
            "            }\n",
            "        ],\n",
            "        \"relationships\": [\n",
            "            {\n",
            "                \"person_1_name\": \"Anya\",\n",
            "                \"person_2_name\": \"Elise\",\n",
            "                \"relationship\": \"daughter\"\n",
            "            },\n",
            "            {\n",
            "                \"person_1_name\": \"Anya\",\n",
            "                \"person_2_name\": \"Edward\",\n",
            "                \"relationship\": \"daughter\"\n",
            "            },\n",
            "            {\n",
            "                \"person_1_name\": \"Anya\",\n",
            "                \"person_2_name\": \"Samuel\",\n",
            "                \"relationship\": \"friend\"\n",
            "            }\n",
            "        ]\n",
            "    }\n",
            "}\n"
          ]
        }
      ],
      "source": [
        "print(json.dumps(type(fc).to_dict(fc), indent=4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4m8FakjCIKmI"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "While the API can handle structured data extraction problems with pure text input and text output, using function calling is likely more reliable since it lets you define a strict schema, and eliminates a potentially error-prone parsing step."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "extract_structured_data.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}