dp-zp-agent/tools/graphrag/evaluation/reproducibility.py

"""Content fingerprints used to freeze retrieval before holdout evaluation."""

from __future__ import annotations

import hashlib
import json
import os
import platform
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from tools.retrieval.publication import DEFAULT_PUBLISHED_ONLY, PUBLICATION_POLICY_VERSION

ROOT = Path(__file__).resolve().parents[3]


def sha256_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def source_corpus_fingerprint(root: Path = ROOT) -> dict[str, Any]:
    """Bind independent expectations to the original Markdown, not graph output."""
    wiki = Path(os.getenv("ZPWIKI_ROOT", str(root.parent / "zpwiki")))
    pages = wiki / "pages"
    if not pages.is_dir():
        raise RuntimeError(f"Original ZP Wiki Markdown directory is missing: {pages}")
    documents = {
        str(path.relative_to(wiki)): sha256_file(path)
        for path in sorted(pages.rglob("README.md"))
        if path.is_file()
    }
    serialized = json.dumps(documents, sort_keys=True, separators=(",", ":")).encode()
    return {
        "sha256": hashlib.sha256(serialized).hexdigest(),
        "document_count": len(documents),
        "documents": documents,
    }


def content_fingerprints(questions: Path, root: Path = ROOT) -> dict[str, Any]:
    # Production and offline measurement code have separate scopes. Historical
    # manifests keep their original bytes and deliberately fail on this new code.
    files = sorted((root / "app").rglob("*.py"))
    for package in ("answers", "retrieval", "runtime", "sync", "knowledge", "rag", "graphrag"):
        files.extend(
            path
            for path in (root / "tools" / package).rglob("*.py")
            if not {"tests", "evaluation"} & set(path.relative_to(root / "tools" / package).parts)
        )
    files.extend(
        root / name
        for name in ["deployment/requirements.txt", "tools/answers/graph_answer_prompt.txt"]
    )
    files = sorted(set(files))
    evaluation = sorted((root / "tools/evaluation").rglob("*.py"))
    evaluation.extend((root / "tools/rag/evaluation").rglob("*.py"))
    evaluation.extend((root / "tools/graphrag/evaluation").rglob("*.py"))
    return {
        "implementation": {str(path.relative_to(root)): sha256_file(path) for path in files},
        "evaluation": {
            str(path.relative_to(root)): sha256_file(path) for path in sorted(evaluation)
        },
        "dataset_sha256": sha256_file(questions),
        "inputs": {
            name: sha256_file(root / name) for name in ["data/documents.json", "data/chunks.json"]
        },
        "source_corpus": source_corpus_fingerprint(root),
        "retrieval_policy": {
            "published_only": DEFAULT_PUBLISHED_ONLY,
            "publication_policy": PUBLICATION_POLICY_VERSION,
            "academic_preparation_is_work": False,
        },
    }


def graph_fingerprint(session: Any) -> dict[str, Any]:
    nodes = session.run(
        "MATCH (n) WHERE n.graph_scope = $scope "
        "RETURN labels(n) AS labels, properties(n) AS properties",
        scope="zpwiki",
    ).data()
    edges = session.run(
        "MATCH (a)-[r]->(b) WHERE a.graph_scope = $scope AND b.graph_scope = $scope "
        "RETURN labels(a) AS source_labels, coalesce(a.id, a.path, a.name) AS source, "
        "type(r) AS type, properties(r) AS properties, labels(b) AS target_labels, "
        "coalesce(b.id, b.path, b.name) AS target",
        scope="zpwiki",
    ).data()
    # elementId and row order are deliberately excluded: they change on rebuild.
    canonical = lambda rows: sorted(
        json.dumps(row, sort_keys=True, ensure_ascii=False, default=str) for row in rows
    )
    content = json.dumps(
        {"nodes": canonical(nodes), "relationships": canonical(edges)}, ensure_ascii=False
    )
    return {
        "sha256": hashlib.sha256(content.encode()).hexdigest(),
        "nodes": len(nodes),
        "relationships": len(edges),
    }


def run_metadata(questions: Path) -> dict[str, Any]:
    def revision(path: Path) -> str | None:
        result = subprocess.run(
            ["git", "-C", str(path), "rev-parse", "HEAD"],
            text=True,
            capture_output=True,
            check=False,
        )
        return result.stdout.strip() if result.returncode == 0 else None

    return {
        "started_at_utc": datetime.now(timezone.utc).isoformat(),
        "python": platform.python_version(),
        "git_head": revision(ROOT),
        "wiki_git_head": revision(Path(os.getenv("ZPWIKI_ROOT", str(ROOT.parent / "zpwiki")))),
        "fingerprints": content_fingerprints(questions),
    }


def verify_freeze(
    manifest: dict,
    metadata: dict,
    graph: dict,
    *,
    limit: int,
    max_entities: int,
    audit: Path | None,
) -> None:
    if manifest["fingerprints"] != metadata["fingerprints"]:
        raise RuntimeError("Frozen implementation, benchmark or inputs changed. TEST refused.")
    if manifest["graph"] != graph:
        raise RuntimeError("Neo4j graph changed since DEV freeze. TEST refused.")
    if manifest["parameters"] != {"limit": limit, "max_entities": max_entities}:
        raise RuntimeError("Retrieval parameters differ from DEV freeze. TEST refused.")
    if audit is None or sha256_file(audit) != manifest["audit_sha256"]:
        raise RuntimeError("Source audit differs from DEV freeze. TEST refused.")