Structured data extraction using function calling

View on Google AI Run in Google Colab View source on GitHub

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.

Setup

pip install -U -q google-generativeai
import pathlib
import textwrap

import google.generativeai as genai
import google.ai.generativelanguage as glm


from IPython.display import display
from IPython.display import Markdown

from google.api_core import retry

def to_markdown(text):
  text = text.replace('•', '  *')
  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))

Once you have the API key, pass it to the SDK. You can do this in two ways:

  • Put the key in the GOOGLE_API_KEY environment variable (the SDK will automatically pick it up from there).
  • Pass the key to genai.configure(api_key=...)
try:
    # Used to securely store your API key
    from google.colab import userdata

    # Or use `os.getenv('API_KEY')` to fetch an environment variable.
    GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')
except ImportError:
    import os
    GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']

genai.configure(api_key=GOOGLE_API_KEY)

genai.configure(
  api_key=GOOGLE_API_KEY,
  client_options={'api_endpoint':'autopush-generativelanguage.sandbox.googleapis.com'})

The example task

For this tutorial you'll extract entities from natural language stories. As an example, below is a story written by Gemini.

new_story = False

if new_story:
  model = genai.GenerativeModel(model_name='gemini-1.0-pro')

  response = model.generate_content("""
      Write a long story about a girl with magic backpack, her family, and at
      least one other charater. Make sure everyone has names. Don't forget to
      describe the contents of the backpack, and where everyone and everything
      starts and ends up.""", request_options={'retry': retry.Retry()})
  story = response.text
  print(response.candidates[0].citation_metadata)
else:
  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."""
to_markdown(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.

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.

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."

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?"

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.

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.

"What's wrong?" she asked.

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."

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.

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."

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.

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.

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.

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.

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.

Using Natural language

Large language models are a powerfuls multitask tools. Often you can just ask Gemini for what you want, and it will do okay.

The Gemini API doesn't have a JSON mode, so there are a few things to watch for when generating data structures this way:

  • Sometimes parsing fails.
  • The schema can't be strictly enforced.

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:

model = model = model = genai.GenerativeModel(
    model_name='gemini-1.0-pro')

response = model.generate_content(textwrap.dedent("""\
    Please return JSON describing the the people, places, things and relationships from this story using the following schema:

    {"people": list[PERSON], "places":list[PLACE], "things":list[THING], "relationships": list[RELATIONSHIP]}

    PERSON = {"name": str, "description": str, "start_place_name": str, "end_place_name": str}
    PLACE = {"name": str, "description": str}
    THING = {"name": str, "description": str, "start_place_name": str, "end_place_name": str}
    RELATIONSHIP = {"person_1_name": str, "person_2_name": str, "relationship": str}

    All fields are required.

    Important: Only return a single piece of valid JSON text.

    Here is the story:

    """) + story)
response.text
'{"people": [{"name": "Anya", "description": "A young girl who possesses a magical backpack.", "start_place_name": "Willow Creek", "end_place_name": null}, {"name": "Elise", "description": "Anya\'s kind-hearted mother.", "start_place_name": "Willow Creek", "end_place_name": null}, {"name": "Edward", "description": "Anya\'s wise-bearded father.", "start_place_name": "Willow Creek", "end_place_name": null}, {"name": "Samuel", "description": "Anya\'s curious and adventurous best friend.", "start_place_name": "Willow Creek", "end_place_name": null}], "places": [{"name": "Willow Creek", "description": "A quaint town nestled amidst rolling hills and whispering willows."}, {"name": "The forest", "description": "A shadowy place with whispering trees and unseen creatures."}], "things": [{"name": "Magical backpack", "description": "A magical backpack that holds an enchanted world filled with wonders.", "start_place_name": "Anya\'s grandmother\'s house", "end_place_name": "Willow Creek"}, {"name": "Shimmering sword", "description": "A sword that Anya uses to defeat the monster.", "start_place_name": "Anya\'s backpack", "end_place_name": null}, {"name": "Book of ancient spells", "description": "A book that contains ancient spells.", "start_place_name": "Anya\'s backpack", "end_place_name": null}, {"name": "Tiny compass", "description": "A compass that always points north.", "start_place_name": "Anya\'s backpack", "end_place_name": null}, {"name": "Magical key", "description": "A key that can open any lock.", "start_place_name": "Anya\'s backpack", "end_place_name": null}, {"name": "Monster", "description": "A massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease.", "start_place_name": "The forest", "end_place_name": null}], "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"}]}'

That returned a json string. Try parsing it:

import json

json_text = response.text.strip('`\r\n ').removeprefix('json')
print(json.dumps(json.loads(json_text), indent=4))
{
    "people": [
        {
            "name": "Anya",
            "description": "A young girl who possesses a magical backpack.",
            "start_place_name": "Willow Creek",
            "end_place_name": null
        },
        {
            "name": "Elise",
            "description": "Anya's kind-hearted mother.",
            "start_place_name": "Willow Creek",
            "end_place_name": null
        },
        {
            "name": "Edward",
            "description": "Anya's wise-bearded father.",
            "start_place_name": "Willow Creek",
            "end_place_name": null
        },
        {
            "name": "Samuel",
            "description": "Anya's curious and adventurous best friend.",
            "start_place_name": "Willow Creek",
            "end_place_name": null
        }
    ],
    "places": [
        {
            "name": "Willow Creek",
            "description": "A quaint town nestled amidst rolling hills and whispering willows."
        },
        {
            "name": "The forest",
            "description": "A shadowy place with whispering trees and unseen creatures."
        }
    ],
    "things": [
        {
            "name": "Magical backpack",
            "description": "A magical backpack that holds an enchanted world filled with wonders.",
            "start_place_name": "Anya's grandmother's house",
            "end_place_name": "Willow Creek"
        },
        {
            "name": "Shimmering sword",
            "description": "A sword that Anya uses to defeat the monster.",
            "start_place_name": "Anya's backpack",
            "end_place_name": null
        },
        {
            "name": "Book of ancient spells",
            "description": "A book that contains ancient spells.",
            "start_place_name": "Anya's backpack",
            "end_place_name": null
        },
        {
            "name": "Tiny compass",
            "description": "A compass that always points north.",
            "start_place_name": "Anya's backpack",
            "end_place_name": null
        },
        {
            "name": "Magical key",
            "description": "A key that can open any lock.",
            "start_place_name": "Anya's backpack",
            "end_place_name": null
        },
        {
            "name": "Monster",
            "description": "A massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease.",
            "start_place_name": "The forest",
            "end_place_name": null
        }
    ],
    "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"
        }
    ]
}

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.

Use function calling

If you haven't gone through the Function calling basics tutorial yet, make sure you do that first.

With function calling your function and its parameters are described to the API as a glm.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.

Define the schema

Start by defining person as an object with string fields name, description, start_place_name, end_place_name.

person = glm.Schema(
    type = glm.Type.OBJECT,
    properties = {
        'name':  glm.Schema(type=glm.Type.STRING),
        'description':  glm.Schema(type=glm.Type.STRING),
        'start_place_name': glm.Schema(type=glm.Type.STRING),
        'end_place_name': glm.Schema(type=glm.Type.STRING)
    },
    required=['name', 'description', 'start_place_name', 'end_place_name']
)

Then define people as an ARRAY of person objects:

people = glm.Schema(
    type=glm.Type.ARRAY,
    items=person
)

Then do the same for each of the entities you're trying to extract:

place = glm.Schema(
    type = glm.Type.OBJECT,
    properties = {
        'name':  glm.Schema(type=glm.Type.STRING),
        'description':  glm.Schema(type=glm.Type.STRING),
    }
)

places = glm.Schema(
    type=glm.Type.ARRAY,
    items=place
)
thing = glm.Schema(
  type = glm.Type.OBJECT,
  properties = {
      'name':  glm.Schema(type=glm.Type.STRING),
      'description':  glm.Schema(type=glm.Type.STRING),
  }
)

things = glm.Schema(
    type=glm.Type.ARRAY,
    items=thing
)
relationship = glm.Schema(
    type = glm.Type.OBJECT,
    properties = {
        'person_1_name':  glm.Schema(type=glm.Type.STRING),
        'person_2_name':  glm.Schema(type=glm.Type.STRING),
        'relationship':  glm.Schema(type=glm.Type.STRING),
    }
)

relationships = glm.Schema(
    type=glm.Type.ARRAY,
    items=relationship
)

Now build the FunctionDeclaration:

add_to_database = glm.FunctionDeclaration(
    name="add_to_database",
    description=textwrap.dedent("""\
        Adds entities to the database.
        """),
    parameters=glm.Schema(
        type=glm.Type.OBJECT,
        properties = {
            'people': people,
            'places': places,
            'things': things,
            'relationships': relationships
        }
    )
)

Call the API

Like you saw in Function calling basics 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):

model = model = genai.GenerativeModel(
    model_name='gemini-1.0-pro',
    tools = [add_to_database])

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:

result = model.generate_content(f"""
Please add the people, places, things, and relationships from this story to the database:

{story}
""")

Now there is no text to parse. The result is a datastructure.

'text' in result.candidates[0].content.parts[0]
False
'function_call' in result.candidates[0].content.parts[0]
True
fc = result.candidates[0].content.parts[0].function_call
print(type(fc))
<class 'google.ai.generativelanguage_v1beta.types.content.FunctionCall'>

The glm.FunctionCall class is based on Google Protocol Buffers, convert it to a more familiar JSON compatible object:

print(json.dumps(type(fc).to_dict(fc), indent=4))
{
    "name": "add_to_database",
    "args": {
        "relationships": [
            {
                "relationship": "mother-daughter",
                "person_2_name": "Elise",
                "person_1_name": "Anya"
            },
            {
                "person_1_name": "Anya",
                "relationship": "father-daughter",
                "person_2_name": "Edward"
            },
            {
                "relationship": "best friends",
                "person_1_name": "Anya",
                "person_2_name": "Samuel"
            }
        ],
        "places": [
            {
                "name": "Willow Creek",
                "description": "a quaint town nestled amidst rolling hills and whispering willows"
            },
            {
                "name": "forest",
                "description": "a shadowy place with rustling undergrowth"
            }
        ],
        "things": [
            {
                "description": "a backpack with a shimmering emerald-green fabric and leather straps, containing a magical sword, a book of ancient spells, a tiny compass that always points north, and a magical key that could open any lock",
                "name": "magical backpack"
            },
            {
                "description": "a weapon that can defeat monsters",
                "name": "shimmering sword"
            },
            {
                "description": "a book containing magical spells",
                "name": "book of ancient spells"
            },
            {
                "name": "tiny compass",
                "description": "a compass that always points north"
            },
            {
                "name": "magical key",
                "description": "a key that can open any lock"
            },
            {
                "name": "monster",
                "description": "a massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease"
            }
        ],
        "people": [
            {
                "description": "a young girl",
                "start_place_name": "Willow Creek",
                "name": "Anya",
                "end_place_name": null
            },
            {
                "start_place_name": "Willow Creek",
                "description": "Anya's mother",
                "name": "Elise",
                "end_place_name": null
            },
            {
                "start_place_name": "Willow Creek",
                "name": "Edward",
                "end_place_name": null,
                "description": "Anya's father"
            },
            {
                "name": "Samuel",
                "end_place_name": null,
                "start_place_name": "Willow Creek",
                "description": "Anya's best friend"
            },
            {
                "name": "tall, lanky boy",
                "description": "a boy who warned Anya about the monster",
                "start_place_name": "Willow Creek",
                "end_place_name": null
            }
        ]
    }
}

Conclusion

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.