DEV Community

Cover image for Building Knowledge Graphs with Gemini
Google AI profile image Laurent Picard
Laurent Picard for Google AI

Posted on Originally published at hackernoon.com

Building Knowledge Graphs with Gemini

#ai #gemini #knowledgegraph #dataextraction

✨ Overview

In this exploration, we'll see how to turn raw, unstructured documents into structured knowledge graphs using Gemini. We'll start by prototyping to develop our intuition. Then, we'll optimize our prompts and outputs, and finally scale up to process entire books or dense legal contracts. By the end, we'll even visualize extracted book narratives and contractual network graphs!

A few notes before we start:

  • I'm a software engineer, developer advocate at Google Cloud, and hope you'll learn a few things. Thoughts and opinions are entirely my own.
  • The complete source code is available in this notebook (including setup details and future updates) under the Apache 2.0 license. You can also directly open the notebook in Colab. This article reproduces all the results generated by a click on “Run all”.
  • You can experiment and build for free with Gemini in the following platforms:

🔥 Challenge

Documents are everywhere. We use them for business, daily operations, legal matters, technical docs, education, and even just for fun. However, documents are not databases. They're generally unstructured, and fully understanding them requires multiple reading passes.

So, can we extract structured knowledge from documents using only the following?

  • 1 document
  • 1 prompt
  • 1 request

Let's try with Gemini…

🏁 Setup

🐍 Python packages

We'll use the following packages:

We'll also need:

  • tenacity for request management (a dependency of google-genai)
  • matplotlib and pillow for data visualization (dependencies of networkx) 
%pip install --quiet "google-genai>=2.6.0" "networkx[default]"
Enter fullscreen mode Exit fullscreen mode

🤝 Gemini API

To use the Gemini API, we have two main options:

  1. Via Agent Platform (formerly Vertex AI) with a Google Cloud project
  2. Via Google AI Studio with a Gemini API key

The Google Gen AI SDK provides a unified interface to these APIs and we can use environment variables for the configuration. 🔽

🛠️ Option 1 - Gemini API via Agent Platform

Requirements:

Gen AI SDK environment variables:

  • GOOGLE_GENAI_USE_ENTERPRISE="True"
  • GOOGLE_CLOUD_PROJECT="<PROJECT_ID>"
  • GOOGLE_CLOUD_LOCATION="<LOCATION>"

💡 For preview models, the location must be set to global. For generally available models, we can choose the closest location among the Google model endpoint locations.

ℹ️ Learn more about setting up a project and a development environment.

🛠️ Option 2 - Gemini API via Google AI Studio

Requirement:

  • A Gemini API key

Gen AI SDK environment variables:

  • GOOGLE_GENAI_USE_ENTERPRISE="False"
  • GOOGLE_API_KEY="<API_KEY>"

ℹ️ Learn more about getting a Gemini API key from Google AI Studio.

💡 You can store your environment configuration outside of the source code:

Environment Method
IDE .env file (or equivalent)
Colab Colab Secrets (🗝️ icon in left panel, see code below)
Colab Enterprise Google Cloud project and location are automatically defined
Workbench Google Cloud project and location are automatically defined

Define the following environment detection functions. You can also define your configuration manually if needed. 🔽 
import os
import sys
from collections.abc import Callable

from google import genai

# Manual setup (leave unchanged if setup is environment-defined)

# @markdown **Which API: Agent Platform (formerly Vertex AI) or Google AI Studio?**
GOOGLE_GENAI_USE_ENTERPRISE = True  # @param {type: "boolean"}

# @markdown **Option A - Google Cloud project [+location]**
GOOGLE_CLOUD_PROJECT = ""  # @param {type: "string"}
GOOGLE_CLOUD_LOCATION = "global"  # @param {type: "string"}

# @markdown **Option B - Google AI Studio API key**
GOOGLE_API_KEY = ""  # @param {type: "string"}


def check_environment() -> bool:
    check_colab_user_authentication()
    return check_manual_setup() or check_enterprise() or check_colab() or check_local()


def check_manual_setup() -> bool:
    return check_define_env_vars(
        GOOGLE_GENAI_USE_ENTERPRISE,
        GOOGLE_CLOUD_PROJECT.strip(),  # Might have been pasted with a newline
        GOOGLE_CLOUD_LOCATION,
        GOOGLE_API_KEY,
    )


def check_enterprise() -> bool:
    # Workbench and Colab Enterprise
    match os.getenv("VERTEX_PRODUCT", ""):