kg: graph build, traversal queries, neo4j export - kg_ocr.graph builds a networkx graph (docs, chunks, entities, citations, co-occurrence) from chunk markdown - analyzer for summaries, top entities/citations, anomaly checks - traversal: chunks_for_entity/citation, related_entities, expand_context - export: JSON round-trip, GraphML, batched MERGE into neo4j - new CLI: ocr-pipeline kg build|stats|query|export - lazy kg_ocr imports, networkx/neo4j behind extras - dropped dead watch.py shim, added KgConfig stub - trimmed README, updated TODO
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
from .neo4j_exporter import Neo4jExporter, export_graphml, export_json, load_json
|
||||
|
||||
__all__ = ["Neo4jExporter", "export_graphml", "export_json", "load_json"]
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user