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

This commit is contained in:
2026-07-20 19:50:38 +02:00
parent 39655fc35f
commit 7503a441c2
29 changed files with 1343 additions and 160 deletions

View File

@@ -0,0 +1,193 @@
"""Build a knowledge graph from OCR pipeline markdown output.
Schema
------
Nodes
doc:<source_hash> Document (source image + OCR metadata)
chunk:<source_hash>:<index> Chunk (text segment of a document)
entity:<normalized text> Entity (gene/protein/chemical mention)
citation:<identifier> Citation (doi/pmid/arxiv/isbn)
Edges
doc -[:CONTAINS]-> chunk
chunk -[:MENTIONS]-> entity
chunk -[:CITES]-> citation
entity -[:CO_OCCURS]-> entity (weight = shared chunks)
"""
from __future__ import annotations
import re
from collections.abc import Iterable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import yaml
DEFAULT_CONSOLIDATED_NAME = "all_ocr.md"
_TEXT_SECTION = "## Extracted Text"
@dataclass
class ChunkRecord:
"""One parsed pipeline markdown chunk file."""
source_hash: str
source_path: str
chunk_index: int
total_chunks: int
text: str
timestamp: str = ""
ocr_engine: str = ""
ocr_confidence_mean: float = 0.0
language: str = ""
has_figures: bool = False
has_tables: bool = False
entities: list[str] = field(default_factory=list)
citations: list[str] = field(default_factory=list) # "type:identifier" strings
@property
def doc_key(self) -> str:
return f"doc:{self.source_hash}"
@property
def chunk_key(self) -> str:
return f"chunk:{self.source_hash}:{self.chunk_index}"
def _split_frontmatter(raw: str) -> tuple[dict[str, Any], str]:
"""Return (frontmatter dict, body). Empty dict when no frontmatter."""
match = re.match(r"\A---\s*\n(.*?)\n---\s*\n?(.*)\Z", raw, re.DOTALL)
if not match:
return {}, raw
data = yaml.safe_load(match.group(1))
return (data if isinstance(data, dict) else {}), match.group(2)
def _extract_text_section(body: str) -> str:
"""Pull the text between '## Extracted Text' and the next '## ' heading."""
if _TEXT_SECTION not in body:
return ""
after = body.split(_TEXT_SECTION, 1)[1]
after = re.split(r"\n## ", after, maxsplit=1)[0]
return after.strip()
def parse_chunk_file(path: Path) -> ChunkRecord | None:
"""Parse one pipeline chunk markdown file. None when it is not a chunk file."""
try:
raw = path.read_text(encoding="utf-8")
except OSError:
return None
frontmatter, body = _split_frontmatter(raw)
if not frontmatter or "source_hash" not in frontmatter:
return None
citations = frontmatter.get("citations_found") or []
if not isinstance(citations, list):
citations = [str(citations)]
return ChunkRecord(
source_hash=str(frontmatter["source_hash"]),
source_path=str(frontmatter.get("source_path", "")),
chunk_index=int(frontmatter.get("chunk_index", 0)),
total_chunks=int(frontmatter.get("total_chunks", 1)),
text=_extract_text_section(body),
timestamp=str(frontmatter.get("timestamp", "")),
ocr_engine=str(frontmatter.get("ocr_engine", "")),
ocr_confidence_mean=float(frontmatter.get("ocr_confidence_mean") or 0.0),
language=str(frontmatter.get("language", "")),
has_figures=bool(frontmatter.get("has_figures", False)),
has_tables=bool(frontmatter.get("has_tables", False)),
entities=[str(e) for e in (frontmatter.get("detected_entities") or [])],
citations=citations,
)
def find_chunk_files(
output_dir: Path, consolidated_name: str = DEFAULT_CONSOLIDATED_NAME
) -> list[Path]:
"""All chunk markdown files under an output directory, excluding consolidated files."""
return sorted(
p for p in output_dir.rglob("*.md") if p.is_file() and p.name != consolidated_name
)
def _normalize_entity(text: str) -> str:
return " ".join(text.strip().split()).lower()
def _split_citation(raw: str) -> tuple[str, str]:
"""'isbn:ISBN:9780...' -> ('isbn', 'ISBN:9780...')."""
if ":" in raw:
ctype, identifier = raw.split(":", 1)
return ctype.strip().lower(), identifier.strip()
return "unknown", raw.strip()
def build_graph(records: Iterable[ChunkRecord]) -> Any:
"""Assemble the knowledge graph from parsed chunk records."""
import networkx as nx
graph = nx.DiGraph()
for record in records:
graph.add_node(
record.doc_key,
kind="document",
source_path=record.source_path,
timestamp=record.timestamp,
ocr_engine=record.ocr_engine,
ocr_confidence_mean=record.ocr_confidence_mean,
language=record.language,
has_figures=record.has_figures,
has_tables=record.has_tables,
total_chunks=record.total_chunks,
)
graph.add_node(
record.chunk_key,
kind="chunk",
text=record.text,
chunk_index=record.chunk_index,
total_chunks=record.total_chunks,
)
graph.add_edge(record.doc_key, record.chunk_key, kind="CONTAINS")
for entity in record.entities:
key = f"entity:{_normalize_entity(entity)}"
if key not in graph:
graph.add_node(key, kind="entity", text=entity.strip(), mentions=0)
graph.nodes[key]["mentions"] += 1
graph.add_edge(record.chunk_key, key, kind="MENTIONS")
for raw_citation in record.citations:
ctype, identifier = _split_citation(raw_citation)
key = f"citation:{identifier}"
if key not in graph:
graph.add_node(key, kind="citation", citation_type=ctype, identifier=identifier)
graph.add_edge(record.chunk_key, key, kind="CITES")
# Co-occurrence: entities sharing a chunk, one edge per unordered pair.
normalized = sorted({_normalize_entity(e) for e in record.entities if e.strip()})
for i, left in enumerate(normalized):
for right in normalized[i + 1 :]:
a_key, b_key = f"entity:{left}", f"entity:{right}"
if graph.has_edge(a_key, b_key):
graph[a_key][b_key]["weight"] += 1
else:
graph.add_edge(a_key, b_key, kind="CO_OCCURS", weight=1)
return graph
def build_from_directory(
output_dir: Path, consolidated_name: str = DEFAULT_CONSOLIDATED_NAME
) -> Any:
"""Parse every chunk file under output_dir and build the graph."""
records = (
record
for record in (parse_chunk_file(p) for p in find_chunk_files(output_dir, consolidated_name))
if record is not None
)
return build_graph(records)