graph_artifacts.py

from __future__ import annotations

from pathlib import Path
from typing import Iterable

import pandas as pd


DEFAULT_INPUT_DIR = Path("./input/wiki-1k")

ARTIFACT_FILES = {
    "documents": "documents.parquet",
    "entities": "entities.parquet",
    "relationships": "relationships.parquet",
    "text_units": "text_units.parquet",
}

REQUIRED_COLUMNS = {
    "documents": {"id", "title", "text"},
    "entities": {"id", "title", "type", "description", "text_unit_ids"},
    "relationships": {
        "id",
        "source",
        "target",
        "description",
        "weight",
        "text_unit_ids",
    },
    "text_units": {"id", "text", "document_id"},
}


def artifact_path(input_dir: str | Path, artifact: str) -> Path:
    """Return the path to one artifact in a new-format GraphRAG dataset."""
    if artifact not in ARTIFACT_FILES:
        raise KeyError(f"Unknown GraphRAG artifact: {artifact}")
    path = Path(input_dir) / ARTIFACT_FILES[artifact]
    if not path.is_file():
        raise FileNotFoundError(f"Missing GraphRAG artifact: {path}")
    return path


def _require_columns(
    dataframe: pd.DataFrame,
    required: Iterable[str],
    artifact: str,
    path: Path,
) -> None:
    missing = sorted(set(required).difference(dataframe.columns))
    if missing:
        raise ValueError(f"{path} ({artifact}) is missing columns: {missing}")


def read_graph_artifacts(
    input_dir: str | Path,
    artifacts: Iterable[str] = ARTIFACT_FILES,
) -> dict[str, pd.DataFrame]:
    """Load and validate new-format GraphRAG parquet artifacts."""
    requested = tuple(dict.fromkeys(artifacts))
    paths = {name: artifact_path(input_dir, name) for name in requested}
    frames = {name: pd.read_parquet(path) for name, path in paths.items()}
    for name, dataframe in frames.items():
        _require_columns(dataframe, REQUIRED_COLUMNS[name], name, paths[name])
    return frames