Generative AI Google Brain

gemini-use-cases-knowledge-graph-knowledge-graph-generation

3 min read

Building Knowledge Graphs with Gemini

Source notebook

Repo path: gemini/use-cases/knowledge-graph/knowledge_graph_generation.ipynb · Open on GitHub · intermediate

Extracts knowledge graphs from text and PDF documents with Gemini using deterministic prompts and TSV outputs.

Summary

This notebook teaches how to turn unstructured documents into structured knowledge graphs with Gemini. It prototypes entity and relationship extraction, compares JSON and TSV output efficiency, then builds helper functions that generate and parse TSV knowledge graph tables. The workflow covers environment setup, Gemini request configuration, multimodal input handling, prompt design, token comparison, and final graph parsing into dataclasses.

Key code patterns

Unified GenAI client setup

from google import genai
 
check_environment()
client = genai.Client()
check_configuration(client)

Uses environment detection so the same notebook can run through Agent Platform or Google AI Studio.

Deterministic generation config

DEFAULT_CONFIG = GenerateContentConfig(
    temperature=0.0,
    top_p=0.0,
    seed=42,
)

Sets extraction-oriented generation parameters to reduce randomness.

Multimodal source parts

mime_type, _ = mimetypes.guess_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)
elif 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)

Normalizes text, PDF, local files, HTTPS URLs, and Cloud Storage URIs into Gemini input parts.

Knowledge graph prompt template

KNOWLEDGE_GRAPH_PROMPT_TEMPLATE = """
**Data Schema**
 
{data_schema}
 
**Instructions**
 
{instructions}
 
**Output Format**
 
{output_format}
"""

Separates schema, instructions, and output format so extraction prompts can be tuned independently.

TSV graph output format

KNOWLEDGE_GRAPH_OUTPUT_FORMAT = f"""
```tsv filename="entities.tsv"
id{TAB}name{TAB}label
[data_rows]
source_id{TAB}link{TAB}target_id
[data_rows]

"""

> Uses compact tabular output to reduce generated tokens versus JSON.

### Parse TSV into dataclasses
```python
for f in fields(cls):
    origin, list_types = get_origin(f.type), get_args(f.type)
    data[f.name] = parse_tsv_block(list_types[0], response_text, f.name)
return cls(**data)

Converts Gemini TSV code blocks into typed KnowledgeGraph entities and relationships.

Models & APIs used

  • Models: gemini-3.1-flash-lite
  • APIs / services: Vertex AI, Gemini API, Cloud Storage
  • SDKs / libraries: google-genai, networkx, tenacity, matplotlib, pillow

When to use this

Use this pattern when extracting typed entities and relationships from long text or PDF documents into a compact knowledge graph.

Gotchas & caveats

  • Agent Platform requires a Google Cloud project and the Agent Platform API enabled.
  • Preview models require location set to global.
  • Google AI Studio mode requires GOOGLE_API_KEY and GOOGLE_GENAI_USE_ENTERPRISE=“False”.
  • Cloud Storage URIs are converted to HTTPS when the client is not using Vertex AI.
  • MIME type detection must succeed for source files.
  • JSON structured output is supported but can be slower and costlier because of verbose output tokens.
  • TSV assumes tabs and newlines do not collide with field values unless escaping is added.

Best practices

  • Use only the provided input data to avoid relying on memorized general knowledge.
  • Use explicit domain terminology such as entities, relationships, nodes, and edges.
  • Include deterministic settings such as temperature=0.0, top_p=0.0, and seed=42 for extraction tasks.
  • Separate data schema, extraction instructions, and output format in the prompt.
  • Use TSV for table-shaped outputs when token efficiency matters.
  • Assign unique sequential entity identifiers and reference them from relationship rows.
  • Include implied entities only when their names can be determined from context.

Graph View

Backlinks