dp-zp-agent/tools/graphrag/build_graphrag.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

from .graph_db import (
    create_neo4j_driver,
    get_neo4j_settings,
)
from .graphrag import (
    GRAPH_SCOPE,
    build_graph_payload,
)
from tools.sync.index_state import DATA_DIR, PENDING_FILE, publication_lock, read_version

PROJECT_ROOT = Path(__file__).resolve().parents[2]

DOCUMENTS_FILE = PROJECT_ROOT / "data" / "documents.json"

CHUNKS_FILE = PROJECT_ROOT / "data" / "chunks.json"


CONSTRAINTS = (
    ("zpwiki_index_state_unique", "IndexState", "id"),
    (
        "zpwiki_document_path_unique",
        "Document",
        "path",
    ),
    (
        "zpwiki_person_id_unique",
        "Person",
        "id",
    ),
    (
        "zpwiki_document_author_id_unique",
        "DocumentAuthor",
        "id",
    ),
    (
        "zpwiki_category_name_unique",
        "Category",
        "name",
    ),
    (
        "zpwiki_topic_id_unique",
        "Topic",
        "id",
    ),
    (
        "zpwiki_keyword_id_unique",
        "Keyword",
        "id",
    ),
    (
        "zpwiki_thesis_id_unique",
        "Thesis",
        "id",
    ),
    (
        "zpwiki_work_id_unique",
        "Work",
        "id",
    ),
)


LEGACY_CONSTRAINTS = ("zpwiki_author_id_unique",)


def graph_version(session) -> dict | None:
    record = session.run(
        "MATCH (s:IndexState {id:$scope}) RETURN s.metadata_json AS metadata", scope=GRAPH_SCOPE
    ).single()
    return json.loads(record["metadata"]) if record and record["metadata"] else None


def write_graph_transaction(
    tx, payload: dict, *, replace: bool = True, version: dict | None = None
) -> None:
    """All destructive data changes belong to the caller's single transaction."""
    if replace:
        clear_managed_graph(tx)
    for key, query in NODE_QUERIES.items():
        write_rows(tx, query, payload[key])
    for key, query in RELATIONSHIP_QUERIES.items():
        write_rows(tx, query, payload[key])
    if replace:
        nodes, relationships = graph_counts(tx)
        labels = {
            "people": "Person",
            "document_authors": "DocumentAuthor",
            "documents": "Document",
            "categories": "Category",
            "topics": "Topic",
            "keywords": "Keyword",
            "theses": "Thesis",
            "works": "Work",
        }
        types = {
            "person_documents": "HAS_DOCUMENT",
            "document_author_documents": "AUTHORED_DOCUMENT",
            "document_categories": "IN_CATEGORY",
            "document_keywords": "HAS_KEYWORD",
            "document_describes_topics": "DESCRIBES",
            "person_theses": "AUTHORED",
            "person_works": "HAS_WORK",
        }
        expected_nodes = {labels[key]: len(payload[key]) for key in labels if payload[key]}
        expected_relations = {types[key]: len(payload[key]) for key in types if payload[key]}
        evidenced = len(payload["thesis_documents"]) + len(payload["work_documents"])
        if evidenced:
            expected_relations["EVIDENCED_BY"] = evidenced
        if dict(nodes) != expected_nodes or dict(relationships) != expected_relations:
            raise RuntimeError("Neo4j počty nezodpovedajú pripravenému grafu; transakcia sa zruší.")
    if version is None:
        tx.run("MATCH (s:IndexState {id:$scope}) DELETE s", scope=GRAPH_SCOPE).consume()
    else:
        # Operational state is deliberately outside the knowledge graph scope.
        tx.run(
            "MERGE (s:IndexState {id:$scope}) SET s.metadata_json=$metadata",
            scope=GRAPH_SCOPE,
            metadata=json.dumps(version, sort_keys=True),
        ).consume()


NODE_QUERIES = {
    "people": """
        UNWIND $rows AS row
        MERGE (n:Person {id: row.id})
        SET n += row
    """,
    "document_authors": """
        UNWIND $rows AS row
        MERGE (n:DocumentAuthor {id: row.id})
        SET n += row
    """,
    "documents": """
        UNWIND $rows AS row
        MERGE (n:Document {path: row.path})
        SET n += row
    """,
    "categories": """
        UNWIND $rows AS row
        MERGE (n:Category {name: row.name})
        SET n += row
    """,
    "topics": """
        UNWIND $rows AS row
        MERGE (n:Topic {id: row.id})
        SET n += row
    """,
    "keywords": """
        UNWIND $rows AS row
        MERGE (n:Keyword {id: row.id})
        SET n += row
    """,
    "theses": """
        UNWIND $rows AS row
        MERGE (n:Thesis {id: row.id})
        SET n += row
    """,
    "works": """
        UNWIND $rows AS row
        MERGE (n:Work {id: row.id})
        SET n += row
    """,
}


RELATIONSHIP_QUERIES = {
    "person_documents": """
        UNWIND $rows AS row
        MATCH (a:Person {
            id: row.person_id
        })
        MATCH (b:Document {
            path: row.document_path
        })
        MERGE (a)-[:HAS_DOCUMENT]->(b)
    """,
    "document_author_documents": """
        UNWIND $rows AS row
        MATCH (a:DocumentAuthor {
            id: row.document_author_id
        })
        MATCH (b:Document {
            path: row.document_path
        })
        MERGE (a)-[:AUTHORED_DOCUMENT]->(b)
    """,
    "document_categories": """
        UNWIND $rows AS row
        MATCH (a:Document {
            path: row.document_path
        })
        MATCH (b:Category {
            name: row.category_name
        })
        MERGE (a)-[:IN_CATEGORY]->(b)
    """,
    "document_keywords": """
        UNWIND $rows AS row
        MATCH (a:Document {
            path: row.document_path
        })
        MATCH (b:Keyword {
            id: row.keyword_id
        })
        MERGE (a)-[:HAS_KEYWORD]->(b)
    """,
    "document_describes_topics": """
        UNWIND $rows AS row
        MATCH (a:Document {
            path: row.document_path
        })
        MATCH (b:Topic {
            id: row.topic_id
        })
        MERGE (a)-[:DESCRIBES]->(b)
    """,
    "person_theses": """
        UNWIND $rows AS row
        MATCH (a:Person {
            id: row.person_id
        })
        MATCH (b:Thesis {
            id: row.thesis_id
        })
        MERGE (a)-[:AUTHORED]->(b)
    """,
    "person_works": """
        UNWIND $rows AS row
        MATCH (a:Person {
            id: row.person_id
        })
        MATCH (b:Work {
            id: row.work_id
        })
        MERGE (a)-[:HAS_WORK]->(b)
    """,
    "thesis_documents": """
        UNWIND $rows AS row
        MATCH (a:Thesis {
            id: row.thesis_id
        })
        MATCH (b:Document {
            path: row.document_path
        })
        MERGE (a)-[:EVIDENCED_BY]->(b)
    """,
    "work_documents": """
        UNWIND $rows AS row
        MATCH (a:Work {
            id: row.work_id
        })
        MATCH (b:Document {
            path: row.document_path
        })
        MERGE (a)-[:EVIDENCED_BY]->(b)
    """,
}


def load_json(
    path: Path,
) -> list[dict[str, Any]]:
    if not path.exists():
        raise FileNotFoundError(f"Missing file: {path}")

    data = json.loads(path.read_text(encoding="utf-8"))

    if not isinstance(
        data,
        list,
    ):
        raise ValueError(f"Expected list in {path}")

    return data


def create_constraints(
    session,
) -> None:
    for (
        name,
        label,
        property_name,
    ) in CONSTRAINTS:
        session.run(
            (
                f"CREATE CONSTRAINT "
                f"{name} IF NOT EXISTS "
                f"FOR (n:{label}) "
                f"REQUIRE n.{property_name} "
                f"IS UNIQUE"
            )
        ).consume()


def drop_legacy_constraints(
    session,
) -> None:
    for name in LEGACY_CONSTRAINTS:
        session.run((f"DROP CONSTRAINT " f"{name} IF EXISTS")).consume()


def clear_managed_graph(
    session,
) -> None:
    session.run(
        """
        MATCH (n)
        WHERE n.graph_scope = $scope
        DETACH DELETE n
        """,
        scope=GRAPH_SCOPE,
    ).consume()


def write_rows(
    session,
    query: str,
    rows: list[dict[str, Any]],
) -> None:
    if not rows:
        return

    session.run(
        query,
        rows=rows,
    ).consume()


def graph_counts(
    session,
) -> tuple[
    list[tuple[str, int]],
    list[tuple[str, int]],
]:
    node_rows = session.run(
        """
        MATCH (n)
        WHERE n.graph_scope = $scope
        UNWIND labels(n) AS label
        RETURN
            label,
            count(*) AS count
        ORDER BY label
        """,
        scope=GRAPH_SCOPE,
    )

    nodes = [
        (
            str(record["label"]),
            int(record["count"]),
        )
        for record in node_rows
    ]

    relationship_rows = session.run(
        """
        MATCH (a)-[r]->(b)
        WHERE
            a.graph_scope = $scope
            AND b.graph_scope = $scope
        RETURN
            type(r) AS relationship,
            count(*) AS count
        ORDER BY relationship
        """,
        scope=GRAPH_SCOPE,
    )

    relationships = [
        (
            str(record["relationship"]),
            int(record["count"]),
        )
        for record in relationship_rows
    ]

    return nodes, relationships


def titled_entity_counts(
    session,
    label: str,
) -> tuple[int, int]:
    if label not in {
        "Thesis",
        "Work",
    }:
        raise ValueError(f"Unsupported label: {label}")

    record = session.run(
        f"""
        MATCH (n:{label})
        WHERE n.graph_scope = $scope
        RETURN
            count(n) AS total,
            count(n.title) AS with_title
        """,
        scope=GRAPH_SCOPE,
    ).single()

    if record is None:
        return 0, 0

    return (
        int(record["total"]),
        int(record["with_title"]),
    )


def print_payload_summary(
    payload: dict[
        str,
        list[dict[str, Any]],
    ],
) -> None:
    print()
    print("Prepared GraphRAG payload")
    print("=" * 60)

    for key in (
        "people",
        "document_authors",
        "documents",
        "categories",
        "topics",
        "keywords",
        "theses",
        "works",
    ):
        print(f"{key:<24}" f"{len(payload[key]):>6}")

    print()
    print("Relationships")

    for key in (
        "person_documents",
        "document_author_documents",
        "document_categories",
        "document_keywords",
        "document_describes_topics",
        "person_theses",
        "person_works",
        "thesis_documents",
        "work_documents",
    ):
        print(f"{key:<30}" f"{len(payload[key]):>6}")

    print("=" * 60)


def print_database_summary(
    session,
) -> None:
    nodes, relationships = graph_counts(session)

    print()
    print("Neo4j GraphRAG knowledge graph")
    print("=" * 60)

    print("Nodes")

    for label, count in nodes:
        print(f"  {label:<22}" f"{count:>6}")

    print()
    print("Relationships")

    for relationship, count in relationships:
        print(f"  {relationship:<22}" f"{count:>6}")

    thesis_total, thesis_with_title = titled_entity_counts(
        session,
        "Thesis",
    )

    work_total, work_with_title = titled_entity_counts(
        session,
        "Work",
    )

    print()
    print("Theses with extracted title: " f"{thesis_with_title}/" f"{thesis_total}")

    print("Works with extracted title: " f"{work_with_title}/" f"{work_total}")

    print("=" * 60)


def main() -> None:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--replace",
        action=argparse.BooleanOptionalAction,
        default=True,
    )

    args = parser.parse_args()

    settings = get_neo4j_settings()

    with create_neo4j_driver(settings) as driver:
        driver.verify_connectivity()

        with driver.session(database=settings.database) as session:
            drop_legacy_constraints(session)

            create_constraints(session)

            import hashlib

            with publication_lock(DATA_DIR, exclusive=True):
                if (DATA_DIR / PENDING_FILE).exists():
                    raise RuntimeError(
                        "Nedokončená spoločná synchronizácia; použite tools.rag.rebuild_index."
                    )
                payload = build_graph_payload(load_json(DOCUMENTS_FILE), load_json(CHUNKS_FILE))
                print_payload_summary(payload)
                version = read_version(DATA_DIR)
                if not args.replace:
                    version = None  # An incremental merge cannot certify removed source records.
                if version and any(
                    hashlib.sha256(path.read_bytes()).hexdigest() != version[key]
                    for path, key in [
                        (DOCUMENTS_FILE, "documents_sha256"),
                        (CHUNKS_FILE, "chunks_sha256"),
                    ]
                ):
                    version = None  # A standalone rebuild must not claim stale provenance.
                with session.begin_transaction() as tx:
                    write_graph_transaction(tx, payload, replace=args.replace, version=version)
                    tx.commit()

            print_database_summary(session)


if __name__ == "__main__":
    main()