kg-scr/kg_ocr/export/neo4j_exporter.py

"""Serialize a knowledge graph to JSON / GraphML, or push it to Neo4j."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

NEO4J_LABELS = {
    "document": "Document",
    "chunk": "Chunk",
    "entity": "Entity",
    "citation": "Citation",
}


def export_json(graph: Any, path: Path) -> Path:
    """Node-link JSON (lossless round-trip with `load_json`)."""
    import networkx as nx

    path.parent.mkdir(parents=True, exist_ok=True)
    data = nx.node_link_data(graph, edges="edges")
    path.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
    return path


def load_json(path: Path) -> Any:
    """Load a graph previously written by `export_json`."""
    import networkx as nx

    data = json.loads(path.read_text(encoding="utf-8"))
    return nx.node_link_graph(data, edges="edges")


def export_graphml(graph: Any, path: Path) -> Path:
    """GraphML for Gephi/Cytoscape. Chunk text is dropped (GraphML attr limits)."""
    import networkx as nx

    path.parent.mkdir(parents=True, exist_ok=True)
    slim = nx.DiGraph()
    for node, data in graph.nodes(data=True):
        slim.add_node(node, **{k: v for k, v in data.items() if k != "text"})
    slim.add_edges_from(
        (u, v, {k: val for k, val in data.items() if isinstance(val, (int, float, str, bool))})
        for u, v, data in graph.edges(data=True)
    )
    nx.write_graphml(slim, str(path))
    return path


class Neo4jExporter:
    """MERGE the graph into a Neo4j database.

    Credentials come from constructor args or NEO4J_URI / NEO4J_USER /
    NEO4J_PASSWORD environment variables.
    """

    def __init__(
        self,
        uri: str | None = None,
        user: str | None = None,
        password: str | None = None,
    ) -> None:
        import os

        self.uri = uri or os.environ.get("NEO4J_URI", "bolt://localhost:7687")
        self.user = user or os.environ.get("NEO4J_USER", "neo4j")
        self.password = password or os.environ.get("NEO4J_PASSWORD", "")
        try:
            from neo4j import GraphDatabase
        except ImportError as exc:
            raise ImportError("Neo4j export needs the driver: uv pip install neo4j") from exc
        self._driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))

    def close(self) -> None:
        self._driver.close()

    def __enter__(self) -> Neo4jExporter:
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.close()

    def push(self, graph: Any, batch_size: int = 500) -> dict[str, int]:
        """MERGE all nodes and edges. Returns counts pushed."""
        nodes = list(graph.nodes(data=True))
        edges = list(graph.edges(data=True))
        with self._driver.session() as session:
            for start in range(0, len(nodes), batch_size):
                session.execute_write(self._push_nodes, nodes[start : start + batch_size])
            for start in range(0, len(edges), batch_size):
                session.execute_write(self._push_edges, edges[start : start + batch_size])
        return {"nodes": len(nodes), "edges": len(edges)}

    @staticmethod
    def _push_nodes(tx: Any, batch: list[tuple[str, dict]]) -> None:
        for node_id, data in batch:
            label = NEO4J_LABELS.get(data.get("kind", ""), "Node")
            props = {k: v for k, v in data.items() if k != "kind"}
            tx.run(
                f"MERGE (n:{label} {{id: $id}}) SET n += $props",
                id=node_id,
                props=props,
            )

    @staticmethod
    def _push_edges(tx: Any, batch: list[tuple[str, str, dict]]) -> None:
        for source, target, data in batch:
            rel = data.get("kind", "RELATED")
            props = {k: v for k, v in data.items() if k != "kind"}
            tx.run(
                f"MATCH (a {{id: $source}}), (b {{id: $target}}) "
                f"MERGE (a)-[r:{rel}]->(b) SET r += $props",
                source=source,
                target=target,
                props=props,
            )