kg-scr/kg_ocr/graph/builder.py

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