548 lines
13 KiB
Python
548 lines
13 KiB
Python
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()
|