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