Building Knowledge Graphs with Gemini

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

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

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

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

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

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

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