Building Knowledge Graphs with Gemini

A developer has demonstrated how to build structured knowledge graphs from unstructured documents using Google's Gemini AI model. The approach involves prototyping with Gemini to extract relationships and entities from raw text, then optimizing prompts and scaling up to process entire books or legal contracts. The method can visualize extracted narratives and contractual network graphs from documents.

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: 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? Let's try with Gemini… We'll use the following packages: google-genai for calling Gemini with the networkx for graph managementWe'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 " To use the Gemini API, we have two main options: 🛠️ 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: 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 | python 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", "" : case "WORKBENCH INSTANCE": pass case "COLAB ENTERPRISE": if not running in colab env : return False case : return False return check define env vars True, os.getenv "GOOGLE CLOUD PROJECT", "" , os.getenv "GOOGLE CLOUD REGION", "" , "", def check colab - bool: if not running in colab env : return False Colab Enterprise was checked before, so this is Colab only from google.colab import auth as colab auth type: ignore colab auth.authenticate user Use Colab Secrets 🗝️ icon in left panel to store the environment variables Secrets are private, visible only to you and the notebooks that you select - Agent Platform: Store your settings as secrets - Google AI: Directly import your Gemini API key from the UI enterprise, project, location, api key = get vars get colab secret return check define env vars enterprise, project, location, api key def check local - bool: enterprise, project, location, api key = get vars os.getenv return check define env vars enterprise, project, location, api key def running in colab env - bool: Colab or Colab Enterprise return "google.colab" in sys.modules def check colab user authentication - None: if running in colab env : from google.colab import auth as colab auth type: ignore colab auth.authenticate user def get colab secret secret name: str, default: str - str: from google.colab import errors, userdata type: ignore try: return userdata.get secret name except errors.SecretNotFoundError: return default def disable colab cell scrollbar - None: if running in colab env : from google.colab import output type: ignore output.no vertical scroll def get vars getenv: Callable str, str , str - tuple bool, str, str, str : Limit getenv calls to the minimum may trigger UI confirmation for secret access enterprise str = getenv "GOOGLE GENAI USE ENTERPRISE", "" if not enterprise str: enterprise str = getenv "GOOGLE GENAI USE VERTEXAI", "" if enterprise str: enterprise = enterprise str.lower in "true", "1" else: enterprise = bool getenv "GOOGLE CLOUD PROJECT", "" project = getenv "GOOGLE CLOUD PROJECT", "" if enterprise else "" location = getenv "GOOGLE CLOUD LOCATION", "" if project else "" api key = getenv "GOOGLE API KEY", "" if not project else "" return enterprise, project, location, api key def check define env vars enterprise: bool, project: str, location: str, api key: str, - bool: match enterprise, bool project , bool location , bool api key : case True, True, , : Agent Platform - Google Cloud project +location location = location or "global" define env vars enterprise, project, location, "" case True, False, , True : Agent Platform - API key define env vars enterprise, "", "", api key case False, , , True : Google AI Studio - API key define env vars enterprise, "", "", api key case : return False return True def define env vars enterprise: bool, project: str, location: str, api key: str, - None: os.environ "GOOGLE GENAI USE ENTERPRISE" = str enterprise os.environ "GOOGLE GENAI USE VERTEXAI" = str enterprise os.environ "GOOGLE CLOUD PROJECT" = project os.environ "GOOGLE CLOUD LOCATION" = location os.environ "GOOGLE API KEY" = api key def check configuration client: genai.Client - None: service = "Agent Platform" if client.vertexai else "Google AI Studio" print f"✅ Using the {service} API", end="" if client. api client.project: print f' with project "{client. api client.project :7 }…"', end="" print f' in location "{client. api client.location}"' elif client. api client.api key: api key = client. api client.api key print f' with API key "{api key :5 }…{api key -5: }"', end="" print f" in case of error, make sure it was created for {service} " print "✅ Environment functions defined" ✅ Environment functions defined To send Gemini requests, we'll use a google.genai client: python from google import genai check environment client = genai.Client check configuration client ✅ Using the Agent Platform API with project "lpdemo-…" in location "global" We need a suite of test data to develop our solution. Multimodality We'll test the following types: text/plain : Classic books are good text sources of varying lengths and languages. application/pdf : Legal agreements are also great examples of complex and dense documents.Gemini is natively multimodal, which means it can process different types of inputs. Once we've built knowledge graphs from text or PDF inputs, the solution will also naturally support the following formats: image/ audio/ video/ General knowledge ⚠️ LLMs are trained on general knowledge, which becomes part of their "long-term memory". To avoid generating memorized information, we'll explicitly instruct the model to use only the provided inputs. Multilinguality Gemini is also natively multilingual https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models expandable-1 , which lets us process inputs and generate outputs in 100+ languages. To keep things general, we'll use English for prompts and knowledge graphs, but you can use any of the 100+ supported languages, as long as your prompts remain clear and explicit. Let's define a few data sources and helpers: 🔽 python import mimetypes from collections.abc import Iterator from enum import Enum from pathlib import Path from google.genai.types import Part GOOGLE CLOUD STORAGE PREFIX = "gs://" HTTPS PREFIX = "https://" FILE PREFIX = "file://" LOCAL FOLDER = "./" class Source Enum : def yield contents self - Iterator Part : file uri = self.value if not client.vertexai: file uri = convert to https url if cloud storage uri file uri mime type, = mimetypes.guess type file uri assert mime type is not None, f"❌ Could not determine mime type: {file uri=}" if file uri.startswith GOOGLE CLOUD STORAGE PREFIX, HTTPS PREFIX : yield Part.from uri file uri=file uri, mime type=mime type return if file uri.startswith FILE PREFIX : file = Path file uri.removeprefix FILE PREFIX assert file.exists , f"❌ File does not exist: {file=}" if mime type == "text/plain": yield Part.from text text=file.read text encoding="utf-8" else: yield Part.from bytes data=file.read bytes , mime type=mime type return def yield source names self - Iterator str : yield self.name def yield source links self - Iterator str : file uri = convert to https url if cloud storage uri self.value if file uri.startswith HTTPS PREFIX : yield file uri return if file uri.startswith FILE PREFIX : yield file uri.removeprefix FILE PREFIX return def convert to https url if cloud storage uri uri: str - str: return f"{HTTPS PREFIX}storage.googleapis.com/{uri.removeprefix GOOGLE CLOUD STORAGE PREFIX }" if uri.startswith GOOGLE CLOUD STORAGE PREFIX else uri def local file filename: str - str: return f"{FILE PREFIX}{LOCAL FOLDER}{filename}" You can find public domain books on Project Gutenberg: https://gutenberg.org/ebooks def project gutenberg txt url id: int - str: return f"{HTTPS PREFIX}gutenberg.org/cache/epub/{id}/pg{id}.txt" class Classic Source : en hugo les misérables = project gutenberg txt url 135 en dumas count of monte cristo = project gutenberg txt url 1184 fr zola thérèse raquin = project gutenberg txt url 7461 fr dumas trois mousquetaires = project gutenberg txt url 13951 fr dumas vingt ans après = project gutenberg txt url 13952 fr dumas comte de monte cristo 1 = project gutenberg txt url 17989 fr dumas comte de monte cristo 2 = project gutenberg txt url 17990 fr dumas comte de monte cristo 3 = project gutenberg txt url 17991 fr dumas comte de monte cristo 4 = project gutenberg txt url 17992 class Document Source : en pharma dev agreement = "gs://cloud-samples-data/documentai/ContractDocAI/CUAD v1/Part I/Development/PhasebioPharmaceuticalsInc 20200330 10-K EX-10.21 12086810 EX-10.21 Development Agreement.pdf" class Collection Source : fr dumas comte de monte cristo = Classic.fr dumas comte de monte cristo 1, Classic.fr dumas comte de monte cristo 2, Classic.fr dumas comte de monte cristo 3, Classic.fr dumas comte de monte cristo 4, fr dumas trois mousquetaires vingt ans après = Classic.fr dumas trois mousquetaires, Classic.fr dumas vingt ans après, def yield contents self - Iterator Part : for source in self.value: yield from source.yield contents def yield source names self - Iterator str : for source in self.value: yield from source.yield source names def yield source links self - Iterator str : for source in self.value: yield from source.yield source links def display input data caption source: Source - None: names = list source.yield source names links = list source.yield source links links = ", ".join f" {name} <{link} " for name, link in zip names, links, strict=True md = f" Input data {links} " display markdown md print "✅ Data helpers defined" ✅ Data helpers defined Gemini comes in different versions https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models gemini-models and sizes Flash-Lite, Flash, and Pro . Let's get started with Gemini 3.1 Flash-Lite, as it offers high performance, low latency, and very high output speed: GEMINI 3 1 FLASH LITE = "gemini-3.1-flash-lite" Gemini can be used in different ways, ranging from factual to creative modes. We're essentially dealing with a data-extraction use case . We want the results to be as factual and deterministic as possible. To achieve this, we can adjust the content generation parameters https://docs.cloud.google.com/vertex-ai/generative-ai/docs/multimodal/content-generation-parameters . We'll set the temperature , top p , and seed parameters to minimize randomness: temperature=0.0 top p=0.0 seed=42 arbitrary fixed value python from enum import StrEnum, auto import IPython.display import tenacity from google.genai.errors import ClientError from google.genai.types import FinishReason, GenerateContentConfig, GenerateContentResponse, ThinkingConfig, ThinkingLevel, class Model Enum : GEMINI 3 1 FLASH LITE = "gemini-3.1-flash-lite" GEMINI 3 5 FLASH = "gemini-3.5-flash" GEMINI 2 5 FLASH = "gemini-2.5-flash" GEMINI 2 5 PRO = "gemini-2.5-pro" Preview GEMINI 3 1 PRO = "gemini-3.1-pro-preview" Default model DEFAULT = GEMINI 3 1 FLASH LITE Default configuration for more deterministic outputs DEFAULT CONFIG = GenerateContentConfig temperature=0.0, top p=0.0, seed=42, Arbitrary fixed value class ShowAs StrEnum : DONT SHOW = auto TEXT = auto MARKDOWN = auto def generate content prompt: str, source: Source | str | None = None, , model: Model | None = None, config: GenerateContentConfig | None = None, system instruction: str | None = None, show prompt: ShowAs = ShowAs.DONT SHOW, show response: ShowAs = ShowAs.MARKDOWN, only show prompt: bool = False, return response: bool = False, - GenerateContentResponse | None: disable colab cell scrollbar model = model or Model.DEFAULT model id = model.value prompt contents = get prompt contents prompt, source, show prompt, only show prompt if only show prompt: return None config = config or get generate content config model, system instruction client = check client for model model response = None display request header model id, source for attempt in get retrier : with attempt: response = client.models.generate content model=model id, contents=prompt contents, type: ignore config=config, display response info response display response response, show response return response if return response else None def get prompt contents prompt: str, source: Source | str | None, show prompt: ShowAs, only show prompt: bool, - list str | Part : def yield prompt contents - Iterator str | Part : if not source: yield prompt.strip return yield "==Start of input data==\n" if isinstance source, str : yield f"{source.strip }\n" else: yield from source.yield contents yield "==End of input data==\n" yield f"==Start of user prompt==\n{prompt.strip }\n==End of user prompt==" prompt contents = list yield prompt contents display prompt prompt contents, show prompt, only show prompt return prompt contents def get generate content config model: Model, system instruction: str | None = None, - GenerateContentConfig: thinking config = get thinking config for model model return GenerateContentConfig system instruction=system instruction, temperature=DEFAULT CONFIG.temperature, top p=DEFAULT CONFIG.top p, seed=DEFAULT CONFIG.seed, thinking config=thinking config, def get thinking config for model model: Model - ThinkingConfig | None: Use minimal thinking configurations since our prompt will directly provide a chain of thought match model: case Model.GEMINI 2 5 FLASH: return ThinkingConfig thinking budget=0 case Model.GEMINI 2 5 PRO: return ThinkingConfig thinking budget=128, include thoughts=False case Model.GEMINI 3 1 FLASH LITE | Model.GEMINI 3 5 FLASH: return ThinkingConfig thinking level=ThinkingLevel.MINIMAL case Model.GEMINI 3 1 PRO: return ThinkingConfig thinking level=ThinkingLevel.LOW case : return None Default dynamic thinking is generally enabled def check client for model model: Model - genai.Client: if model.value.endswith "-preview" and client.vertexai and client. api client.location = "global" : Preview models are only available on the "global" location return genai.Client enterprise=client.vertexai, project=client. api client.project, location="global", return client def get retrier - tenacity.Retrying: return tenacity.Retrying stop=tenacity.stop after attempt 7 , wait=tenacity.wait incrementing start=10, increment=1 , retry=tenacity.retry if exception should retry request , reraise=True, def should retry request err: BaseException - bool: if not isinstance err, ClientError : return False print f"❌ ClientError {err.code}: {err.message}" retry = False match err.code: case 400 if err.message is not None and " try again " in err.message: Workshop: project accessing Cloud Storage for the first time service agent provisioning retry = True case 429: Workshop: temporary project with 1 QPM quota retry = True print f"🔄 Retry: {retry}" return retry PRINT COLUMNS = 80 PRINT SEPARATOR CHAR = "-" PRINT SEPARATOR = PRINT COLUMNS PRINT SEPARATOR CHAR def print caption caption: str - None: print f" {caption} ".center PRINT COLUMNS, PRINT SEPARATOR CHAR def print separator - None: print PRINT SEPARATOR def display response info response: GenerateContentResponse - None: if usage metadata := response.usage metadata: if usage metadata.prompt token count: print f"Input tokens : {usage metadata.prompt token count:9,d}" if usage metadata.cached content token count: print f"Cached tokens : {usage metadata.cached content token count:9,d}" if usage metadata.candidates token count: print f"Output tokens : {usage metadata.candidates token count:9,d}" if usage metadata.thoughts token count: print f"Thoughts tokens : {usage metadata.thoughts token count:9,d}" if not response.candidates: print "❌ No response.candidates " return if finish reason := response.candidates 0 .finish reason = FinishReason.STOP: print f"❌ {finish reason = }" if not response.text: print "❌ No response.text " return def display prompt contents: list str | Part , show as: ShowAs, only show prompt: bool, - None: def yield prompt strings - Iterator str : for content in contents: if isinstance content, Part : yield f"{content r}\n" continue yield content if only show prompt and show as == ShowAs.DONT SHOW: show as = ShowAs.TEXT if show as == ShowAs.DONT SHOW: return separator = "\n" if show as == ShowAs.MARKDOWN else "" prompt = separator.join yield prompt strings print caption "Prompt" match show as: case ShowAs.TEXT: print prompt case ShowAs.MARKDOWN: display markdown prompt if only show prompt: print separator def display request header model id: str, source: Source | str | None = None - None: print caption f"Request / {model id}" def display response response: GenerateContentResponse, show as: ShowAs - None: if show as == ShowAs.DONT SHOW or not response text := response.text : return print caption "Start of Response" response text = response text.strip match show as: case ShowAs.TEXT: print response text case ShowAs.MARKDOWN: display markdown response text print caption "End of Response" def display markdown markdown: str - None: IPython.display.display IPython.display.Markdown markdown print "✅ Helpers defined" ✅ Helpers defined Before diving into a solution, it helps to start by prototyping to build some intuition about the natural behavior of the model. Let's define a short text of a few sentences: text = """ - Henry Jones is a famous archaeologist. He is actually a "Junior" because he is named after his father. - Sophie is Henry's daughter, shares his last name, and works as a software engineer. - William Smith is an aerospace engineer and Sophie's lifelong friend. Everybody calls him Bill and Beau is his dog. - Short Round met Henry as a child. They first became close friends, and Henry officially adopted him a few years later. - Sophie and Bill both work at Acme Aerospace. """ 🧪 Let's see if Gemini can spot our characters… prompt = """ Using only the input data, list all people and animals mentioned. """ generate content prompt, text ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 148 Output tokens : 67 ------------------------------ Start of Response ------------------------------- Based on the input data provided, here are the people and animals mentioned: People: Henry Jones also known as Henry Jones Junior Sophie Jones William Smith also known as Bill Short Round Animals: Beau Bill's dog ------------------------------- End of Response -------------------------------- 💡 All people and animals are detected as expected. 🧪 Now, let's see if it can connect the dots and figure out who's who… prompt = """ Using only the input data, list all people and animals mentioned, and how they relate to each other. """ generate content prompt, text ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 156 Output tokens : 168 ------------------------------ Start of Response ------------------------------- Based on the input data provided, here are the people and animals mentioned and their relationships: People: Henry Jones Junior : A famous archaeologist. He is the father of Sophie, the adoptive father of Short Round, and is named after his own father. Sophie Jones: A software engineer at Acme Aerospace. She is the daughter of Henry Jones and a lifelong friend of Bill William Smith . William Bill Smith: An aerospace engineer at Acme Aerospace. He is a lifelong friend of Sophie and the owner of Beau. Short Round: The adopted son of Henry Jones. He met Henry as a child and they became close friends before the adoption. Animals: Beau: A dog owned by William Bill Smith. ------------------------------- End of Response -------------------------------- 💡 Notes We're not domain experts in the field we're exploring yet . An LLM processes instructions based on the given prompt and its training knowledge. This knowledge is part of its long-term memory, and we can learn a lot directly from the model itself. 🧪 Let's ask Gemini: prompt = """ What is the terminology used when building a knowledge graph? Please provide a simple data example in JSON. """ generate content prompt ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 21 Output tokens : 581 ------------------------------ Start of Response ------------------------------- Building a knowledge graph involves representing information as a network of interconnected entities. Here is the core terminology and a simple data example. Core Terminology 1. Entity Node : The "things" in your graph e.g., a person, a place, a product . 2. Relationship Edge/Link : The connection between two entities. It describes how they interact e.g., "works at," "lives in," "is a friend of" . 3. Property Attribute : Key-value pairs that provide more detail about an entity or a relationship e.g., a person's "age" or a relationship's "start date" . 4. Label: A category assigned to a node or edge to define its type e.g., a node might have the label "Person" . 5. Schema Ontology : The formal structure or "blueprint" that defines the types of entities allowed and the rules for how they can be connected. 6. Triple: The fundamental unit of a knowledge graph, consisting of a Subject → Predicate → Object e.g., Alice → works at → Google . --- Simple Data Example JSON In a knowledge graph, data is often represented as a collection of Nodes and Edges . json { "nodes": { "id": "1", "label": "Person", "properties": { "name": "Alice", "age": 30 } }, { "id": "2", "label": "Company", "properties": { "name": "Google", "industry": "Technology" } } , "edges": { "id": "e1", "source": "1", "target": "2", "label": "WORKS AT", "properties": { "since": 2020 } } } Breakdown of the Example: Nodes: We have two entities: "Alice" a Person and "Google" a Company . Edge: We have one relationship: "WORKS AT" connecting Alice to Google. Properties: We stored specific details like Alice's age and the year she started working at Google. Triple representation: This JSON effectively encodes the triple: Alice — WORKS AT — Google . ------------------------------- End of Response -------------------------------- 💡 We learn that knowledge graphs are made of entities and relationships , also called nodes and edges , and we get a nice introduction to the field. Using domain terminology will help make our prompts explicit and precise. To extract knowledge graphs, we'll reason in terms of entities and relationships, adopting domain terminology. If we think of the final result as a database, our goal is to generate two linked tables, allowing us to reason in terms of data and fields. Here is a conceptual view of what we want to achieve: Entities id | name | label | |---|---|---| | 0 | Henry Jones Jr. | person | | 1 | Henry Jones Sr. | person | Relationships source id | link | target id | |---|---|---| | 0 | child of | 1 | Let's call this approach "tabular extraction" and split our instructions to output two successive tables, while still using a single request… In our prototype text, the entities we want to extract are characters like people or animals . We can define an entity data schema with the fields id 0, 1, 2… , name full name of the entity , and label person | animal . 🧪 Let's extract the entities: prompt = """ Data Schema Entity: - id : Unique integer identifier 0, 1, 2… . - name : Full name of the entity. - label : person | animal . Instructions 1. Entity Extraction: - Extract every distinct entity from the input data that matches an allowed label . - Include entities that are explicitly named as well as implied entities whose names can be determined from the context. 2. Output the results as a JSON array inside a fenced code block. """ generate content prompt, text ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 249 Output tokens : 195 ------------------------------ Start of Response ------------------------------- { "id": 0, "name": "Henry Jones Jr.", "label": "person" }, { "id": 1, "name": "Henry Jones Sr.", "label": "person" }, { "id": 2, "name": "Sophie Jones", "label": "person" }, { "id": 3, "name": "William Smith", "label": "person" }, { "id": 4, "name": "Beau", "label": "animal" }, { "id": 5, "name": "Short Round", "label": "person" } ------------------------------- End of Response -------------------------------- 💡 Remarks 🧪 Now, let's extract both the entities and their relationships: prompt = """ Data Schema Entity: - id : Unique integer identifier 0, 1, 2… . - name : Full name of the entity. - label : person | animal . Relationship: - source id : id of the subject entity. - link : snake case predicate describing the relationship. - target id : id of the object entity. Instructions 1. Entity Extraction: - Extract every distinct entity from the input data that matches an allowed label . - Include entities that are explicitly named as well as implied entities whose names can be determined from the context. 2. Relationship Extraction: - Extract every distinct relationship between them. - If a relationship changes over time, make sure to include every distinct stage of the relationship. 3. Output a JSON object with keys entities and relationships inside a fenced code block. """ response = generate content prompt, text, return response=True ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 340 Output tokens : 456 ------------------------------ Start of Response ------------------------------- { "entities": { "id": 0, "name": "Henry Jones Jr.", "label": "person" }, { "id": 1, "name": "Henry Jones Sr.", "label": "person" }, { "id": 2, "name": "Sophie Jones", "label": "person" }, { "id": 3, "name": "William Smith", "label": "person" }, { "id": 4, "name": "Beau", "label": "animal" }, { "id": 5, "name": "Short Round", "label": "person" } , "relationships": { "source id": 0, "link": "child of", "target id": 1 }, { "source id": 2, "link": "child of", "target id": 0 }, { "source id": 3, "link": "friend of", "target id": 2 }, { "source id": 4, "link": "pet of", "target id": 3 }, { "source id": 5, "link": "friend of", "target id": 0 }, { "source id": 0, "link": "adopted", "target id": 5 }, { "source id": 5, "link": "child of", "target id": 0 } } ------------------------------- End of Response -------------------------------- 💡 Remarks relationships array. link predicates are completely dynamic a level of flexibility we left in the prompt . While it's interesting to see this natural behavior, we'll want to make it more deterministic for production since our prompt has too many degrees of freedom. pet of Person " is an asymmetric relationship that could also be extracted as " Person owner of Animal ". This is another area where the prompt is too open-ended. In the finalization section, we'll see an example that asks the model to extract symmetric and asymmetric relationships in both directions.We've concluded our prototyping stage with promising results using a data schema. To move to production, the next step is to control the generation with a specific structured output. The JSON format has industry-wide support and serves as a core or intermediate format in many use cases. For the next step, we would typically define classes using the Pydantic library and request a pure JSON output with a response schema in the config parameters: response mime type="application/json" response schema="CLASS DERIVED FROM PYDANTIC BASE MODEL" ⚠️ However, JSON is a pretty verbose format, designed for interoperability but not optimized for size. Even if we generate compact JSON also called minified JSON , it still has inherent verbosity due to: ℹ️ When using LLMs, once the first token is generated, the remaining generation time is roughly proportional to the number of output tokens. Similarly, the cost of a request is based on token usage input + output , with output tokens being significantly more expensive than input tokens. 💡 A better output structure will positively impact both generation speed and request cost. Let's explore an alternative… Our tabular-extraction problem clearly calls for table outputs. An interesting possibility is to ask for Tab-Separated Values TSV outputs. For example, we can define our output to be formatted as two consecutive TSV tables. Example output format tsv filename="entities.tsv" id{TAB}name{TAB}label rows tsv filename="relationships.tsv" source id{TAB}link{TAB}target id rows Will this work? Generating structured outputs like TSV will work seamlessly, as Gemini excels at patterns. We just need to be explicit about what's expected. Will this be efficient? 💡 For our use case, this structure looks optimal: ℹ️ CSV could be another alternative, but common separators like commas are everywhere in natural language and frequently appear in names and descriptions e.g., if we decide to extend entity fields . If you're interested in this topic, check out the TOON format https://github.com/toon-format/toon , which proposes a JSON alternative using a YAML+CSV mix. To make an informed decision, we should actually compare the number of tokens needed to represent the same data… python import csv import io import json import re from typing import Literal def get data from response response: GenerateContentResponse - dict: response text = response.text or "" pattern = r" json\s . ? \s " match = re.search pattern, response text, re.DOTALL json str = match.group 1 if match else response text try: data = json.loads json str if not isinstance data, dict : print "❌ Returning empty dict could not parse response as dict " data = {} except json.JSONDecodeError, TypeError : print "❌ Returning empty dict failed parsing the JSON string " data = {} return data def get tsv string from data data: dict - str: output = "" for key, items in data.items : rows = "" if items: with io.StringIO as out: headers = list items 0 .keys writer = csv.DictWriter out, fieldnames=headers, delimiter="\t", lineterminator="\n", writer.writeheader writer.writerows items rows = out.getvalue if output: output += "\n" output += f' {% endraw %} tsv filename="{key}.tsv"\n{rows} {% raw %} \n' return output def print text excerpt title: str, text: str, max chars: int = 400 - None: assert max chars 0 chars = len text if chars <= 0: print caption "❌ Empty text" return if chars <= max chars: print caption f"{title} {chars} chars " print text return print caption f"{title} - First {max chars}/{chars} chars {max chars / chars:.0%} " print f"{text :max chars }…" def compare json vs tsv response: GenerateContentResponse | None, only show excerpts: bool = False, - None: def get gain rows: list tuple int, int , col: Literal 0, 1 - str: val 0, val 1 = rows 0 col , rows 1 col return f" {1 - val 1 / val 0:.1%} " if val 0 0 else "?" def yield table string rows source title: str, target title: str, rows: list tuple int, int , - Iterator list str : yield "", "Chars", "Tokens" yield "-", "-:", "-:" for caption, values in zip source title, target title , rows : yield caption, str value for value in values yield " Gain ", get gain rows, 0 , get gain rows, 1 def display gain table source title: str, target title: str, source text: str, target text: str, - None: print caption f"{source title} → {target title}" model = Model.DEFAULT model id = model.value client = check client for model model rows: list tuple int, int = for s in source text, target text : chars = len s tokens = client.models.count tokens model=model id, contents=s .total tokens rows.append chars, tokens or 0 markdown = "\n".join "| " + " | ".join row + " |" for row in yield table string rows source title, target title, rows display markdown markdown if response is None: print "❌ response is None" return data = get data from response response formatted json = f" {% endraw %} json\n{json.dumps data, indent=2 }\n {% raw %} " compact json = f" {% endraw %} json\n{json.dumps data, separators= ',', ':' }\n {% raw %} " tsv = get tsv string from data data if only show excerpts: max chars = len tsv print text excerpt "Formatted JSON", formatted json, max chars print text excerpt "Compact JSON", compact json, max chars print text excerpt "TSV", tsv, max chars return display gain table "Formatted JSON", "Compact JSON", formatted json, compact json display gain table "Compact JSON", "TSV", compact json, tsv display gain table "Formatted JSON", "TSV", formatted json, tsv print separator print "✅ JSON vs TSV helpers defined" ✅ JSON vs TSV helpers defined 🧪 First, let's compare how much data we can represent for the same number of characters based on our latest API response: compare json vs tsv response, only show excerpts=True ----------------- Formatted JSON - First 335/1122 chars 30% ------------------ json { "entities": { "id": 0, "name": "Henry Jones Jr.", "label": "person" }, { "id": 1, "name": "Henry Jones Sr.", "label": "person" }, { "id": 2, "name": "Sophie Jones", "label": "person" }, { "id": 3, "name": "William Smith", … ------------------- Compact JSON - First 335/665 chars 50% ------------------- json {"entities": {"id":0,"name":"Henry Jones Jr.","label":"person"},{"id":1,"name":"Henry Jones Sr.","label":"person"},{"id":2,"name":"Sophie Jones","label":"person"},{"id":3,"name":"William Smith","label":"person"},{"id":4,"name":"Beau","label":"animal"},{"id":5,"name":"Short Round","label":"person"} ,"relationships": {"source i… ------------------------------- TSV 335 chars -------------------------------- tsv filename="entities.tsv" id name label 0 Henry Jones Jr. person 1 Henry Jones Sr. person 2 Sophie Jones person 3 William Smith person 4 Beau animal 5 Short Round person tsv filename="relationships.tsv" source id link target id 0 child of 1 2 child of 0 3 friend of 2 4 pet of 3 5 friend of 0 0 adopted 5 5 child of 0 💡 Notice how much more information can be represented in the same number of text characters. This will apply similarly to token counts. 🧪 And now, let's compare the gains, especially for token counts: compare json vs tsv response ------------------------ Formatted JSON → Compact JSON ------------------------- | Chars | Tokens | | |---|---|---| | Formatted JSON | 1122 | 456 | | Compact JSON | 665 | 216 | Gain | 40.7% | 52.6% | ------------------------------ Compact JSON → TSV ------------------------------ | Chars | Tokens | | |---|---|---| | Compact JSON | 665 | 216 | | TSV | 335 | 137 | Gain | 49.6% | 36.6% | ----------------------------- Formatted JSON → TSV ----------------------------- | Chars | Tokens | | |---|---|---| | Formatted JSON | 1122 | 456 | | TSV | 335 | 137 | Gain | 70.1% | 70.0% | -------------------------------------------------------------------------------- 💡 Savings in output tokens: With a double-digit percentage reduction in output tokens, building knowledge graphs with TSV outputs is significantly faster and cheaper Now, let's finalize our code with optimized structures… First, it helps to define a structured prompt template, so we can focus on specific parts of our solution using a divide-and-conquer approach. Here's a possible prompt template: KNOWLEDGE GRAPH PROMPT TEMPLATE = """ Data Schema {data schema} Instructions {instructions} Output Format {output format} """ Then, here are some possible Entity , Relationship , and KnowledgeGraph data classes with the matching output format: python from dataclasses import dataclass, field @dataclass class Entity: id: int name: str label: str @dataclass class Relationship: source id: int link: str target id: int @dataclass class KnowledgeGraph: entities: list Entity = field default factory=list relationships: list Relationship = field default factory=list TAB = "\t" KNOWLEDGE GRAPH OUTPUT FORMAT = f""" Format the output strictly as two TSV code blocks including the header row : tsv filename="entities.tsv" id{TAB}name{TAB}label data rows tsv filename="relationships.tsv" source id{TAB}link{TAB}target id data rows """ 💡 While the Gen AI SDK natively supports Pydantic models for JSON structured outputs, we're using standard Python data classes here and TSV outputs to maximize our token efficiency. ℹ️ If you use multiple entity or relationship data classes in your solution, you can dynamically generate the output format specification using features of the dataclasses package like class docstrings and field descriptions . python from dataclasses import fields, is dataclass from typing import get args, get origin, get type hints def generate knowledge graph data schema: str, instructions: str, source: Source | str, , model: Model | None = None, config: GenerateContentConfig | None = None, system instruction: str | None = None, show prompt: ShowAs = ShowAs.DONT SHOW, show response: ShowAs = ShowAs.DONT SHOW, - KnowledgeGraph: prompt = get prompt for data schema and instructions data schema, instructions response = generate content prompt, source, model=model, config=config, system instruction=system instruction, show prompt=show prompt, show response=show response, return response=True, knowledge graph = parse list dataclass KnowledgeGraph, response if isinstance response, GenerateContentResponse else KnowledgeGraph display knowledge graph info knowledge graph return knowledge graph def show knowledge graph prompt data schema: str, instructions: str, source: Source | str, , model: Model | None = None, config: GenerateContentConfig | None = None, system instruction: str | None = None, show as: ShowAs = ShowAs.TEXT, - None: prompt = get prompt for data schema and instructions data schema, instructions generate content prompt, source, model=model, config=config, system instruction=system instruction, show prompt=show as, only show prompt=True, def get prompt for data schema and instructions data schema: str, instructions: str, - str: return KNOWLEDGE GRAPH PROMPT TEMPLATE.format data schema=data schema.strip , instructions=instructions.strip , output format=KNOWLEDGE GRAPH OUTPUT FORMAT.strip , .strip def parse list dataclass T cls: type T , response: GenerateContentResponse - T: assert is dataclass cls if not response text := response.text : return cls data = {} for f in fields cls : origin, list types = get origin f.type , get args f.type assert origin is list , f"❌ Field {f.name} must be a list dataclass parameterized list" assert len list types == 1, f"❌ Expected 1 single type: {len list types =}" data f.name = parse tsv block list types 0 , response text, f.name return cls data def parse tsv block T cls: type T , data: str, tsv filestem: str - list T : rows = valid fields = get type hints cls tsv string = extract tsv block data, tsv filestem for row in csv.DictReader io.StringIO tsv string , delimiter="\t" : casted data = {} for key, value in row.items : if key not in valid fields or value is None: continue field type = valid fields key try: Note: Works only for directly castable types such as int, float, str, enum e.g., not bool casted data key = field type value except ValueError, TypeError : print f'❌ Could not cast "{value}" to {field type} → Skipping {row=}' break else: try: rows.append cls casted data except TypeError as e: print f"❌ Could not instantiate {cls. name }: {e} → Skipping {row=}" return rows def extract tsv block data: str, filestem: str - str: pattern = rf' {% endraw %} tsv filename="{re.escape filestem }.tsv"\s \n . ? {% raw %} ' if not match := re.search pattern, data, flags=re.DOTALL : print f'❌ Could not find a TSV block for "{filestem=}"' return "" return match.group 1 def display knowledge graph info kg: KnowledgeGraph - None: print caption "Knowledge Graph Info" print f"Entities : {len kg.entities :3,d}" print f"Relationships : {len kg.relationships :3,d}" print separator print "✅ Knowledge graph generation helpers defined" ✅ Knowledge graph generation helpers defined And here is a possible data schema with some instructions to generate a knowledge graph for our book analysis use case: python from enum import StrEnum, auto class BookAnalysisEntityLabel StrEnum : PERSON = auto ANIMAL = auto ORGANIZATION = auto def pipe delimited union enum: type StrEnum - str: return "|".join f" {e.value} " for e in enum BOOK ANALYSIS DATA SCHEMA = f""" Entity: - id : Unique integer identifier 0, 1, 2… . - name : Most complete name as exclusively determined from the input data. - label : {pipe delimited union BookAnalysisEntityLabel }. Relationship: - source id : id of the subject entity. - link : snake case predicate. - target id : id of the object entity. """ BOOK ANALYSIS INSTRUCTIONS = """ - Extract every distinct entity: - Treat distinct pseudonyms/identities as separate entities. - Include implied entities whose names can be exclusively determined from the input data. - Extract every distinct relationship between them: - Use specific link predicates in snake case as needed e.g., alias of , son of , fiancée of , friend of , murderer of , employer of , in love with , rival of . - If a relationship changes over time, make sure to include every distinct stage of the relationship. - For every asymmetric relationship extracted, make sure to include the logical inverse relationship e.g., A husband of B AND B wife of A , A employer of B AND B employee of A . - For every symmetric relationship extracted, make sure to include both directions e.g., A friend of B AND B friend of A . """ Verify the structured prompt: show knowledge graph prompt BOOK ANALYSIS DATA SCHEMA, BOOK ANALYSIS INSTRUCTIONS, text, show as=ShowAs.TEXT, ------------------------------------ Prompt ------------------------------------ ==Start of input data== - Henry Jones is a famous archaeologist. He is actually a "Junior" because he is named after his father. - Sophie is Henry's daughter, shares his last name, and works as a software engineer. - William Smith is an aerospace engineer and Sophie's lifelong friend. Everybody calls him Bill and Beau is his dog. - Short Round met Henry as a child. They first became close friends, and Henry officially adopted him a few years later. - Sophie and Bill both work at Acme Aerospace. ==End of input data== ==Start of user prompt== Data Schema Entity: - id : Unique integer identifier 0, 1, 2… . - name : Most complete name as exclusively determined from the input data. - label : person | animal | organization . Relationship: - source id : id of the subject entity. - link : snake case predicate. - target id : id of the object entity. Instructions - Extract every distinct entity: - Treat distinct pseudonyms/identities as separate entities. - Include implied entities whose names can be exclusively determined from the input data. - Extract every distinct relationship between them: - Use specific link predicates in snake case as needed e.g., alias of , son of , fiancée of , friend of , murderer of , employer of , in love with , rival of . - If a relationship changes over time, make sure to include every distinct stage of the relationship. - For every asymmetric relationship extracted, make sure to include the logical inverse relationship e.g., A husband of B AND B wife of A , A employer of B AND B employee of A . - For every symmetric relationship extracted, make sure to include both directions e.g., A friend of B AND B friend of A . Output Format Format the output strictly as two TSV code blocks including the header row : tsv filename="entities.tsv" id name label data rows tsv filename="relationships.tsv" source id link target id data rows ==End of user prompt== -------------------------------------------------------------------------------- 🧪 Let's generate a knowledge graph: knowledge graph = generate knowledge graph BOOK ANALYSIS DATA SCHEMA, BOOK ANALYSIS INSTRUCTIONS, text, show response=ShowAs.TEXT, print knowledge graph.entities print knowledge graph.relationships ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 534 Output tokens : 244 ------------------------------ Start of Response ------------------------------- tsv filename="entities.tsv" id name label 0 Henry Jones Jr. person 1 Henry Jones Sr. person 2 Sophie Jones person 3 William Smith person 4 Bill person 5 Beau animal 6 Short Round person 7 Acme Aerospace organization tsv filename="relationships.tsv" source id link target id 0 son of 1 1 father of 0 0 father of 2 2 daughter of 0 2 friend of 3 3 friend of 2 3 alias of 4 4 alias of 3 3 employee of 7 7 employer of 3 2 employee of 7 7 employer of 2 3 owner of 5 5 pet of 3 6 friend of 0 0 friend of 6 0 adopted father of 6 6 adopted son of 0 ------------------------------- End of Response -------------------------------- ----------------------------- Knowledge Graph Info ----------------------------- Entities : 8 Relationships : 18 -------------------------------------------------------------------------------- Entity id=0, name='Henry Jones Jr.', label='person' , Entity id=1, name='Henry Jones Sr.', label='person' , Entity id=2, name='Sophie Jones', label='person' , Entity id=3, name='William Smith', label='person' , Entity id=4, name='Bill', label='person' , Entity id=5, name='Beau', label='animal' , Entity id=6, name='Short Round', label='person' , Entity id=7, name='Acme Aerospace', label='organization' Relationship source id=0, link='son of', target id=1 , Relationship source id=1, link='father of', target id=0 , Relationship source id=0, link='father of', target id=2 , Relationship source id=2, link='daughter of', target id=0 , Relationship source id=2, link='friend of', target id=3 , Relationship source id=3, link='friend of', target id=2 , Relationship source id=3, link='alias of', target id=4 , Relationship source id=4, link='alias of', target id=3 , Relationship source id=3, link='employee of', target id=7 , Relationship source id=7, link='employer of', target id=3 , Relationship source id=2, link='employee of', target id=7 , Relationship source id=7, link='employer of', target id=2 , Relationship source id=3, link='owner of', target id=5 , Relationship source id=5, link='pet of', target id=3 , Relationship source id=6, link='friend of', target id=0 , Relationship source id=0, link='friend of', target id=6 , Relationship source id=0, link='adopted father of', target id=6 , Relationship source id=6, link='adopted son of', target id=0 💡 This is looking good Now, let's go to the next stage and build a network graph from our data… Now that we have our entities and relationships neatly packed into data classes, let's bring them to life. We'll use networkx to build a network graph. Using domain terminology, entities become nodes and relationships become directed edges. We'll also calculate node centralities to identify key nodes and use the Louvain method to detect communities clusters of closely related nodes … python import textwrap from typing import cast import networkx as nx import numpy as np from networkx.algorithms.community.louvain import louvain communities NODE CENTRALITY = "node centrality" NODE COMMUNITY INDEX = "node community index" NODE COLOR = "node color" EDGE COLOR = "edge color" Wrap names over multiple lines avoids long node labels MULTILINE NAMES = True MULTILINE CHARS = 12 def build graph kg: KnowledgeGraph, remove orphan nodes: bool - nx.DiGraph: For simplicity, we build a DiGraph adapt to MultiDiGraph if needed graph = nx.DiGraph Nodes ← Entities node name from id: dict int, str = {} for entity in kg.entities: node name, display name = get node and display names for entity entity node name from id entity.id = node name graph.add node node name, name=display name Edges ← Relationships for relationship in kg.relationships: source node = node name from id.get relationship.source id, "" target node = node name from id.get relationship.target id, "" if not source node or not target node: print f"❌ Skipping relationship due to empty node:\n{relationship}" continue For simplicity, each link has the same weight For better community detection, define your own weight mappings e.g., family members have stronger bonds weight = 1 edge label = relationship.link if graph.has edge source node, target node : existing data = graph source node target node existing data "link" += f"\n{edge label}" existing data "weight" += weight else: graph.add edge source node, target node, link=edge label, weight=weight if remove orphan nodes: graph.remove nodes from list nx.isolates graph return graph def get node and display names for entity entity: Entity - tuple str, str : snake case name = " ".join map str.lower, entity.name.split node name = f"{entity.id} {snake case name}" display name = entity.name if MULTILINE NAMES: display name = "\n".join textwrap.wrap display name, width=MULTILINE CHARS return node name, display name def color gen color count: int - Iterator str : B50, R50, Y50, G50 = " 4285F4", " EA4335", " FBBC04", " 34A853" B20, R20, Y20, G20 = " AECBFA", " F6AEA9", " FDE293", " A8DAB5" B05, R05, Y05, G05 = " E8F0FE", " FCE8E6", " FEF7E0", " E6F4EA" COLORS = B50, R50, Y50, G50, B20, R20, Y20, G20, B05, R05, Y05, G05 for i in range color count : yield COLORS i % len COLORS Node = str Positions = dict Node, np.ndarray Community = set Node Communities = list Community Nodes = list Node INTER COMMUNITY EDGE COLOR = " 8888" def init graph data graph: nx.Graph - Nodes: def node centrality node: Node - float: return graph.nodes node NODE CENTRALITY def community max centrality community: Community - float: return max node centrality node for node in community def nodes sorted by community communities: Communities - Nodes: entities = for community in communities: sorted entities = sorted community, key=node centrality, reverse=True entities.extend sorted entities return entities Node centralities centralities = nx.betweenness centrality graph, endpoints=True for node key in graph.nodes: graph.nodes node key NODE CENTRALITY = centralities node key Communities communities = cast Communities, louvain communities graph, seed=42 sorted communities = sorted communities, key=community max centrality, reverse=True Community colors community count = len sorted communities community colors = list color gen community count Node community indices and colors for community index, community in enumerate sorted communities : for node key in community: node = graph.nodes node key node NODE COMMUNITY INDEX = community index node NODE COLOR = community colors community index Edge colors for node key i, node key j, edge data in graph.edges data=True : node i = graph.nodes node key i node j = graph.nodes node key j same community = node i NODE COMMUNITY INDEX == node j NODE COMMUNITY INDEX edge data EDGE COLOR = node i NODE COLOR if same community else INTER COMMUNITY EDGE COLOR return nodes sorted by community sorted communities def compute node positions graph: nx.DiGraph, entities: Nodes - Positions: Derive an undirected graph preserving the entity order undirected graph = nx.Graph for entity in entities: undirected graph.add node entity undirected graph.add edges from graph.edges if len entities < 10: positions = nx.circular layout undirected graph else: positions = nx.kamada kawai layout undirected graph positions = nx.arf layout undirected graph, positions, seed=42 return positions @dataclass class GraphData: knowledge graph: KnowledgeGraph remove orphan nodes: bool = True graph: nx.DiGraph = field init=False nodes: Nodes = field init=False positions: Positions = field init=False def post init self - None: self.graph = build graph self.knowledge graph, self.remove orphan nodes self.nodes = init graph data self.graph self.positions = compute node positions self.graph, self.nodes print "✅ Network graph helpers defined" ✅ Network graph helpers defined Let's test this: graph data = GraphData knowledge graph print f"{graph data.graph = s}" print f"{graph data.nodes = }" graph data.graph = DiGraph with 8 nodes and 16 edges graph data.nodes = '2 sophie jones', '3 william smith', '7 acme aerospace', '5 beau', '4 bill', '0 henry jones jr.', '1 henry jones sr.', '6 short round' The extracted data is much easier to understand when you can actually see it We can use matplotlib to draw our network graphs. We'll size the nodes based on their centrality and color-code them by community. To make it even easier to digest, we'll generate an animated sequence highlighting each character's connections one by one… python import typing from io import BytesIO import matplotlib.pyplot as plt from IPython import display from matplotlib.axes import Axes from matplotlib.backends.backend agg import FigureCanvasAgg from matplotlib.figure import Figure from PIL import Image as PilImage class AnimationFormat StrEnum : Matches PIL.Image.Image.format WEBP = auto PNG = auto GIF = auto FIGURE DPI = 200 FIGURE FACTOR = 1.0 ANIMATION INTRO DURATION = 2500 ANIMATION FRAME DURATION = 250 EDGE STYLE = "arc3,rad=0.2" NodeSizes = dict Node, int if running in colab env : ANIMATION FORMAT = AnimationFormat.WEBP else: ANIMATION FORMAT = AnimationFormat.GIF def init figure title: str, subtitle: str - tuple Figure, Axes : figsize = 16 FIGURE FACTOR, 9 FIGURE FACTOR fig, ax = plt.subplots figsize=figsize, dpi=FIGURE DPI ax.set title title, loc="left" ax.set title subtitle, loc="right" ax.axis "off" fig.tight layout pad=2 return fig, ax def draw nodes graph: nx.Graph, positions: Positions, ax: Axes - NodeSizes: node view = graph.nodes data=True node sizes = { node: max 500, int 10000 data NODE CENTRALITY for node, data in node view } node colors = str data NODE COLOR for , data in node view border width = 3.0 nx.draw networkx nodes graph, pos=positions, node size=list node sizes.values , node color=node colors, alpha=0.95, ax=ax, linewidths=border width, labels = {node: data.get "name", str node for node, data in node view} nx.draw networkx labels graph, positions, labels=labels, ax=ax if len node sizes = 10: Extend axis limits default rounding changes them too much otherwise x min, x max , y min, y max = ax.get xlim , ax.get ylim x min, x max = int x min - 1.0 , int x max + 1.0 y min, y max = int y min - 1.0 , int y max + 1.0 ax.set xlim x min, x max ax.set ylim y min, y max return node sizes def draw edges graph: nx.DiGraph, positions: Positions, node sizes: NodeSizes, ax: Axes, , focused node: Node | None = None, - None: Draw specific out-edges or all edges if focused node: out edges = graph.edges focused node , data=True else: out edges = graph.edges data=True edge colors = data EDGE COLOR for , , data in out edges edge list = u, v for u, v, in out edges ordered sizes = node sizes n for n in graph.nodes nx.draw networkx edges graph, positions, edge list, edge color=edge colors, style=":", alpha=0.9, arrowstyle="-| ", arrowsize=20, ax=ax, node size=ordered sizes, connectionstyle=EDGE STYLE, edge labels = { u, v : data "link" for u, v, data in out edges} nx.draw networkx edge labels graph, positions, edge labels, font size=8, font family="monospace", bbox=dict ec=" FFF8", fc=" FFF8" , ax=ax, node size=ordered sizes, type: ignore connectionstyle=EDGE STYLE, type: ignore def display graph title: str, subtitle: str, graph data: GraphData - None: , ax = init figure title, subtitle node graph = graph data.graph edge graph = graph data.graph positions = graph data.positions node sizes = draw nodes node graph, positions, ax draw edges edge graph, positions, node sizes, ax plt.show def yield images title: str, subtitle: str, graph data: GraphData, - Iterator PilImage.Image : fig, ax = init figure title, subtitle canvas = FigureCanvasAgg fig positions = graph data.positions node graph = graph data.graph node sizes = draw nodes node graph, positions, ax edge graph = graph data.graph for focused node in None, graph data.nodes : if focused node is not None: draw edges edge graph, positions, node sizes, ax, focused node=focused node canvas.draw image size = canvas.get width height image bytes = canvas.buffer rgba yield PilImage.frombytes "RGBA", image size, image bytes .convert "RGB" plt.close fig def generate animation title: str, subtitle: str, graph data: GraphData, format: AnimationFormat, - BytesIO: frames = list yield images title, subtitle, graph data assert len frames = 1 if format == AnimationFormat.GIF: Dither all frames with the same palette to optimize the animation The animation is cumulative, so most colors are in the last frame method = PilImage.Quantize.MEDIANCUT palettized = frames -1 .quantize method=method frames = frame.quantize method=method, palette=palettized for frame in frames The animation will be played in a loop: start cycling with the most complete frame first frame = frames -1 next frames = frames :-1 durations = ANIMATION INTRO DURATION durations += ANIMATION FRAME DURATION len next frames params: dict str, typing.Any = dict save all=True, append images=next frames, duration=durations, loop=0, match format: case AnimationFormat.GIF: params.update optimize=False case AnimationFormat.PNG: params.update optimize=True case AnimationFormat.WEBP: params.update lossless=True image io = BytesIO first frame.save image io, str format .upper , params return image io def display graph animation title: str, subtitle: str, graph data: GraphData - None: image io = generate animation title, subtitle, graph data, ANIMATION FORMAT ipython image = display.Image data=image io.getvalue display.display ipython image def get graph title and subtitle source: Source | str, model: Model | None = None, domain: str = "", - tuple str, str : if isinstance source, str : source name = f'"{source.strip :20 }…"' else: source name = source.name model id = model.value if model else Model.DEFAULT.value model id = model id.removesuffix "-preview" separator = " • " title parts = "Knowledge Graph" if domain: title parts.append domain subtitle parts = source name, model id return separator.join title parts , separator.join subtitle parts def display knowledge graph knowledge graph: KnowledgeGraph, source: Source | str, model: Model | None = None, animated: bool = False, domain: str = "", remove orphan nodes: bool = True, - None: if not knowledge graph.entities: return title, subtitle = get graph title and subtitle source, model, domain graph data = GraphData knowledge graph, remove orphan nodes if animated: display graph animation title, subtitle, graph data else: display graph title, subtitle, graph data def extract knowledge graph data schema: str, instructions: str, source: Source | str, model: Model | None = None, , domain: str = "", animated: bool = False, - None: if isinstance source, Source : display input data caption source knowledge graph = generate knowledge graph data schema, instructions, source, model=model, display knowledge graph knowledge graph, source, model, animated, domain print "✅ Data visualization helpers defined" ✅ Data visualization helpers defined Let's test this: display knowledge graph knowledge graph, text 💡 We can now quickly visualize and understand our knowledge graphs. This helps us iterate faster when refining prompts. ℹ️ While this simple approach is great for a quick overview, you might want to swap it out for more specialized libraries if you want to explore the graph interactively. Let's define a book analysis function: python def analyze book source: Source, model: Model | None = None, animated: bool = False, - None: extract knowledge graph BOOK ANALYSIS DATA SCHEMA, BOOK ANALYSIS INSTRUCTIONS, source, model, domain="Character Connections", animated=animated, 🧪 First, let's build a knowledge graph based on Zola's Thérèse Raquin : analyze book Classic.fr zola thérèse raquin Input data fr zola thérèse raquin https://gutenberg.org/cache/epub/7461/pg7461.txt ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 119,518 Cached tokens : 118,753 Output tokens : 357 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 11 Relationships : 28 -------------------------------------------------------------------------------- 💡 We can extract and understand the book's narrative in seconds: The love triangle between Thérèse, Camille, and Laurent clearly stands out. Despite being over 200 pages long and 100k+ tokens , this novel is incredibly minimalistic, revolving around a limited set of characters, which is reflected in the network graph. Adding locations to the extracted entities would also reveal the claustrophobic atmosphere of the book in the resulting knowledge graph. ℹ️ We used the original French version with English instructions, which works seamlessly. If you translate the instructions, you can also generate dynamic relationship links in different languages see the 100+ supported languages https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models expandable-1 . 🧪 Let's see how the model handles the interconnected cast of Hugo's Les Misérables : analyze book Classic.en hugo les misérables, model=Model.GEMINI 3 5 FLASH, Input data en hugo les misérables https://gutenberg.org/cache/epub/135/pg135.txt -------------------------- Request / gemini-3.5-flash -------------------------- Input tokens : 803,328 Output tokens : 2,908 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 65 Relationships : 225 -------------------------------------------------------------------------------- 💡 We quickly get the gist of the novel. ℹ️ Donald Knuth famously used Les Misérables as an example back in 1994 see the Stanford GraphBase https://www-cs-faculty.stanford.edu/~knuth/sgb.html . It's been consistently used in Natural Language Processing NLP because of its massive scale, linguistic complexity, and multiple high-quality translations. At 800k+ tokens, this book is still a true stress test. Note that we used Gemini 3.5 Flash instead of 3.1 Flash-Lite . Larger models can infer more and perform deeper consolidation. 🧪 Now, let's process just the first volume of Le Comte de Monte-Cristo in French : analyze book Classic.fr dumas comte de monte cristo 1 Input data fr dumas comte de monte cristo 1 https://gutenberg.org/cache/epub/17989/pg17989.txt ----------------------- Request / gemini-3.1-flash-lite ------------------------ Input tokens : 215,582 Output tokens : 713 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 36 Relationships : 36 -------------------------------------------------------------------------------- 💡 The graph shows the initial setup of Edmond Dantès' world, his friends, and his betrayers. The love triangle between Edmond, Mercédès, and Fernand also comes to light. ℹ️ Note that the antagonist is referred to only as Fernand , which is expected. We learn his full name, Fernand Mondego , only in volume three. In the prompt, we asked the model to "extract" entities and determine names "from the input data" to avoid generating memorized knowledge. To detect regressions when updating the prompt, we can use this as a unit test to ensure that our response actually provides extracted data and not memorized information. 🧪 And finally, let's analyze the entire saga of Le Comte de Monte-Cristo , all four volumes at once 800k+ tokens : analyze book Collection.fr dumas comte de monte cristo, model=Model.GEMINI 3 5 FLASH, animated=True, Input data fr dumas comte de monte cristo 1 https://gutenberg.org/cache/epub/17989/pg17989.txt , fr dumas comte de monte cristo 2 https://gutenberg.org/cache/epub/17990/pg17990.txt , fr dumas comte de monte cristo 3 https://gutenberg.org/cache/epub/17991/pg17991.txt , fr dumas comte de monte cristo 4 https://gutenberg.org/cache/epub/17992/pg17992.txt -------------------------- Request / gemini-3.5-flash -------------------------- Input tokens : 840,385 Cached tokens : 835,551 Output tokens : 2,659 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 50 Relationships : 212 -------------------------------------------------------------------------------- 💡 The multiple aliases of Edmond Dantès and other characters are nicely extracted. The complex plot of this book is built on characters juggling multiple identities in a psychological chess match. 🎉 This is an example of how we can complete the challenge in a single request, using the fewest tokens possible and zero thinking tokens, while providing multiple levels of consolidation. ℹ️ Note that this is a proof of concept. For an exhaustive extraction in a fully professional solution, we would probably set up a multi-step workflow and consider the following: location , date … , and then save them to a database. evidence or excerpt to serve as direct proof, or chapter or page for source attribution . In this step, we could also dig deeper and differentiate literal/figurative relationships e.g., biological vs. figurative parents .With our current setup, generalizing to other types of content is as simple as adapting our data schema and instructions. For example, legal agreements are among the densest types of documents. They're usually made up of many articles and clauses, with every sentence carrying specific legal weight, outlining obligations, or providing definitions. But what kind of knowledge graph do we want to build from a legal agreement? 🧪 Let's test a minimalistic, open-ended prompt: source = Document.en pharma dev agreement model = Model.GEMINI 3 5 FLASH AGREEMENT OPEN DATA SCHEMA = """ Entity: - id : Unique integer identifier 0, 1, 2… . - name : Name of the entity. - label : Entity type. Relationship: - source id : id of the subject entity. - link : snake case predicate. - target id : id of the object entity. """ AGREEMENT OPEN INSTRUCTIONS = """ Perform a comprehensive, highly-granular entity/relationship extraction. """ extract knowledge graph AGREEMENT OPEN DATA SCHEMA, AGREEMENT OPEN INSTRUCTIONS, source, model, domain="Agreement High-Level Extraction", Input data en pharma dev agreement https://storage.googleapis.com/cloud-samples-data/documentai/ContractDocAI/CUAD v1/Part I/Development/PhasebioPharmaceuticalsInc 20200330 10-K EX-10.21 12086810 EX-10.21 Development%20Agreement.pdf -------------------------- Request / gemini-3.5-flash -------------------------- Input tokens : 51,276 Output tokens : 221 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 10 Relationships : 9 -------------------------------------------------------------------------------- 💡 Remarks 🧪 Then, let's try semi-open instructions focusing on legal obligations: class AgreementEntityLabel StrEnum : PARTY = auto PERSON = auto ROLE = auto LOCATION = auto JURISDICTION = auto FINANCIAL AMOUNT = auto DATE = auto EVENT = auto ASSET = auto PRODUCT = auto INTELLECTUAL PROPERTY = auto OBLIGATION TYPE = auto AGREEMENT SEMI OPEN DATA SCHEMA = f""" Entity: - id : Unique integer identifier 0, 1, 2… . - name : Name of the entity. - label : {pipe delimited union AgreementEntityLabel }. Relationship: Obligation, right, or transfer from a source entity to a target entity. - source id : The id of the subject entity. - link : Specific snake case predicate. - target id : The id of the object entity. """ AGREEMENT SEMI OPEN INSTRUCTIONS = """ - Analyze the input data for every covenant e.g., "shall", "will", "must", "is obligated to" and perform an exhaustive extraction. - Make sure to deconstruct complex obligations: For complex clauses e.g., "A shall pay B $X Million within Y days of the Effective Date" , extract: - The primary obligation: A is obligated to pay B - The value: A's obligation subject to $X Million - The trigger: A's obligation triggered by Effective Date """ extract knowledge graph AGREEMENT SEMI OPEN DATA SCHEMA, AGREEMENT SEMI OPEN INSTRUCTIONS, source, model, domain="Agreement Obligations", Input data en pharma dev agreement https://storage.googleapis.com/cloud-samples-data/documentai/ContractDocAI/CUAD v1/Part I/Development/PhasebioPharmaceuticalsInc 20200330 10-K EX-10.21 12086810 EX-10.21 Development%20Agreement.pdf -------------------------- Request / gemini-3.5-flash -------------------------- Input tokens : 51,466 Output tokens : 2,002 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 85 Relationships : 93 -------------------------------------------------------------------------------- 💡 Remarks What if we don't care about legal obligations, but rather the document's architecture? Let's shift our focus to the structure itself and extract how sections, clauses, and defined terms are hierarchically organized… 🧪 And now, let's test closed instructions focusing on the document structure: class AgreementStructureEntityLabel StrEnum : DEFINED TERM = auto DOCUMENT SECTION = auto DOCUMENT = auto class AgreementStructureRelationshipType StrEnum : DEFINED IN = auto CONTAINS = auto AGREEMENT STRUCTURAL DATA SCHEMA = f""" Entity: - id : Unique integer identifier 0, 1, 2… . - name : Name of the entity. - label : {pipe delimited union AgreementStructureEntityLabel }. Relationship: Connection from a source entity to a target entity. - source id : The id of the subject entity. - link : {pipe delimited union AgreementStructureRelationshipType }. - target id : The id of the object entity. """ AGREEMENT STRUCTURAL OPEN INSTRUCTIONS = """ - Extract every distinct entity that matches an allowed label . - Extract every distinct relationship representing a structural connection hierarchical organization between these entities: - You must be comprehensive and highly granular. If multiple distinct relationships exist between the same pair of entities, create a separate entry for each. """ extract knowledge graph AGREEMENT STRUCTURAL DATA SCHEMA, AGREEMENT STRUCTURAL OPEN INSTRUCTIONS, source, model, domain="Agreement Structure", Input data en pharma dev agreement https://storage.googleapis.com/cloud-samples-data/documentai/ContractDocAI/CUAD v1/Part I/Development/PhasebioPharmaceuticalsInc 20200330 10-K EX-10.21 12086810 EX-10.21 Development%20Agreement.pdf -------------------------- Request / gemini-3.5-flash -------------------------- Input tokens : 51,351 Output tokens : 9,108 ----------------------------- Knowledge Graph Info ----------------------------- Entities : 314 Relationships : 513 -------------------------------------------------------------------------------- 💡 If you extract hundreds of entities from a massive document, your graph will quickly turn into an unreadable hairball. For larger datasets, you'll want to export your nodes and edges to a dedicated graph database, which typically comes with its own visualization and exploration tools. We successfully extracted data and built knowledge graphs from documents by following these steps: These principles apply to many other data-extraction domains and will allow you to solve your own complex problems. Have fun and happy solving ➕ More