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

kg_ocr/__init__.py

@@ -1,17 +1,50 @@
-"""Experimental txtai/litellm RAG interface built on the OCR pipeline.
+"""Knowledge graph + RAG interface built on the OCR pipeline.
 
-This package is a prototype. The supported interface is the `ocr-pipeline` CLI
-(src/ocr_pipeline). Extra dependencies are required: `uv sync --extra kg`.
+Lightweight submodules (graph, export) import without optional deps.
+The txtai/litellm RAG pieces need: uv sync --extra kg.
 """
 
-try:
-    from .ocr import get_screenshots, extract_text
-    from .embeddings import create_and_index
-    from .rag import retrieve, ask_wllm
-except ImportError as exc:
-    raise ImportError(
-        "kg_ocr is experimental and needs extra dependencies: "
-        "uv sync --extra kg  (or: uv pip install txtai litellm python-dotenv)"
-    ) from exc
+from __future__ import annotations
 
-__all__ = ["get_screenshots", "extract_text", "create_and_index", "retrieve", "ask_wllm"]
+from typing import Any
+
+_LAZY: dict[str, tuple[str, str]] = {
+    "get_screenshots": (".ocr", "get_screenshots"),
+    "extract_text": (".ocr", "extract_text"),
+    "create_and_index": (".embeddings", "create_and_index"),
+    "retrieve": (".rag", "retrieve"),
+    "ask_wllm": (".rag", "ask_wllm"),
+    "build_graph": (".graph", "build_graph"),
+    "build_from_directory": (".graph", "build_from_directory"),
+    "summary": (".graph", "summary"),
+    "top_entities": (".graph", "top_entities"),
+    "detect_anomalies": (".graph", "detect_anomalies"),
+    "chunks_for_entity": (".graph", "chunks_for_entity"),
+    "chunks_for_citation": (".graph", "chunks_for_citation"),
+    "related_entities": (".graph", "related_entities"),
+    "expand_context": (".graph", "expand_context"),
+    "export_json": (".export", "export_json"),
+    "export_graphml": (".export", "export_graphml"),
+    "load_json": (".export", "load_json"),
+    "Neo4jExporter": (".export", "Neo4jExporter"),
+}
+
+__all__ = list(_LAZY)
+
+
+def __getattr__(name: str) -> Any:
+    if name not in _LAZY:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = _LAZY[name]
+    try:
+        from importlib import import_module
+
+        module = import_module(module_name, __name__)
+    except ImportError as exc:
+        raise ImportError(
+            f"kg_ocr.{name} needs optional dependencies: uv sync --extra kg "
+            f"(graph/export only need networkx; embeddings/rag also need txtai, litellm)"
+        ) from exc
+    value = getattr(module, attr)
+    globals()[name] = value
+    return value

kg_ocr/cli/build_graph.py

@@ -1,3 +0,0 @@
-# to do
-
-# add entrypoint to setuppy

kg_ocr/cli/export_graph.py

@@ -1 +0,0 @@
-# to do

kg_ocr/cli/process_screenshots.py

@@ -1 +0,0 @@
-# to do

kg_ocr/embeddings/config_loader.py

@@ -1 +1,31 @@
-# in future prefer a config, especially for graph traversal
+"""Config for kg_ocr: graph traversal defaults and Neo4j credentials.
+
+Environment variables: NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class KgConfig:
+    """Tunables for graph building/traversal plus Neo4j connection."""
+
+    # Traversal
+    expand_limit: int = 5
+    related_limit: int = 10
+
+    # Neo4j
+    neo4j_uri: str = "bolt://localhost:7687"
+    neo4j_user: str = "neo4j"
+    neo4j_password: str = field(default="", repr=False)
+
+    @classmethod
+    def from_env(cls) -> KgConfig:
+        return cls(
+            neo4j_uri=os.environ.get("NEO4J_URI", cls.neo4j_uri),
+            neo4j_user=os.environ.get("NEO4J_USER", cls.neo4j_user),
+            neo4j_password=os.environ.get("NEO4J_PASSWORD", ""),
+        )

kg_ocr/embeddings/indexer.py

@@ -1,18 +1,27 @@
-from txtai.embeddings import Embeddings
+from __future__ import annotations
+
+from typing import Any
+
+DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
 
-def create_and_index(
-    data: list[str], model: str = "sentence-transformers/all-MiniLM-L6-v2"
-) -> Embeddings:
-    """Create and index embeddings from text."""
-    embeddings = Embeddings({
-        "path": model,
-        "content": True,
-        "hybrid": True,
-        "scoring": "bm25",
-    })
+def create_and_index(data: list[str], model: str = DEFAULT_MODEL) -> Any:
+    """Create and index embeddings from text.
+
+    Requires txtai (uv sync --extra kg). Returns a txtai Embeddings instance.
+    """
+    try:
+        from txtai.embeddings import Embeddings
+    except ImportError as exc:
+        raise ImportError("create_and_index needs txtai: uv sync --extra kg") from exc
+
+    embeddings = Embeddings(
+        {
+            "path": model,
+            "content": True,
+            "hybrid": True,
+            "scoring": "bm25",
+        }
+    )
     embeddings.index(data)
     return embeddings
-
-
-

kg_ocr/export/__init__.py

@@ -0,0 +1,3 @@
+from .neo4j_exporter import Neo4jExporter, export_graphml, export_json, load_json
+
+__all__ = ["Neo4jExporter", "export_graphml", "export_json", "load_json"]

kg_ocr/export/neo4j_exporter.py

@@ -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,
+            )

kg_ocr/graph/__init__.py

@@ -0,0 +1,39 @@
+from .analyzer import (
+    Anomaly,
+    detect_anomalies,
+    summary,
+    top_citations,
+    top_cooccurrences,
+    top_entities,
+)
+from .builder import (
+    ChunkRecord,
+    build_from_directory,
+    build_graph,
+    find_chunk_files,
+    parse_chunk_file,
+)
+from .traversal import (
+    chunks_for_citation,
+    chunks_for_entity,
+    expand_context,
+    related_entities,
+)
+
+__all__ = [
+    "Anomaly",
+    "ChunkRecord",
+    "build_from_directory",
+    "build_graph",
+    "chunks_for_citation",
+    "chunks_for_entity",
+    "detect_anomalies",
+    "expand_context",
+    "find_chunk_files",
+    "parse_chunk_file",
+    "related_entities",
+    "summary",
+    "top_citations",
+    "top_cooccurrences",
+    "top_entities",
+]

kg_ocr/graph/analyzer.py

@@ -1 +1,117 @@
-# also anomaly detection here
+"""Analyze a knowledge graph built from OCR pipeline output."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+LOW_CONFIDENCE_THRESHOLD = 0.4
+HUB_MENTIONS_THRESHOLD = 20
+
+
+@dataclass(frozen=True)
+class Anomaly:
+    kind: str
+    node: str
+    detail: str
+
+
+def _nodes_of_kind(graph: Any, kind: str) -> list[tuple[str, dict]]:
+    return [(n, d) for n, d in graph.nodes(data=True) if d.get("kind") == kind]
+
+
+def summary(graph: Any) -> dict[str, int]:
+    """Node and edge counts grouped by kind."""
+    counts: dict[str, int] = {}
+    for _, data in graph.nodes(data=True):
+        kind = data.get("kind", "unknown")
+        counts[f"nodes:{kind}"] = counts.get(f"nodes:{kind}", 0) + 1
+    for _, _, data in graph.edges(data=True):
+        kind = data.get("kind", "unknown")
+        counts[f"edges:{kind}"] = counts.get(f"edges:{kind}", 0) + 1
+    return counts
+
+
+def top_entities(graph: Any, limit: int = 10) -> list[tuple[str, int]]:
+    """Most-mentioned entities as (text, mentions), descending."""
+    entities = [
+        (data.get("text", ""), int(data.get("mentions", 0)))
+        for _, data in _nodes_of_kind(graph, "entity")
+    ]
+    entities.sort(key=lambda item: item[1], reverse=True)
+    return entities[:limit]
+
+
+def top_cooccurrences(graph: Any, limit: int = 10) -> list[tuple[str, str, int]]:
+    """Strongest entity co-occurrence pairs as (left, right, weight)."""
+    pairs = []
+    for left, right, data in graph.edges(data=True):
+        if data.get("kind") != "CO_OCCURS":
+            continue
+        left_text = graph.nodes[left].get("text", left)
+        right_text = graph.nodes[right].get("text", right)
+        pairs.append((left_text, right_text, int(data.get("weight", 1))))
+    pairs.sort(key=lambda item: item[2], reverse=True)
+    return pairs[:limit]
+
+
+def top_citations(graph: Any, limit: int = 10) -> list[tuple[str, str, int]]:
+    """Most-cited identifiers as (type, identifier, citing chunks)."""
+    rows = []
+    for node, data in _nodes_of_kind(graph, "citation"):
+        citing = sum(
+            1 for _, _, edge in graph.in_edges(node, data=True) if edge.get("kind") == "CITES"
+        )
+        rows.append((data.get("citation_type", ""), data.get("identifier", node), citing))
+    rows.sort(key=lambda item: item[2], reverse=True)
+    return rows[:limit]
+
+
+def detect_anomalies(
+    graph: Any,
+    low_confidence: float = LOW_CONFIDENCE_THRESHOLD,
+    hub_mentions: int = HUB_MENTIONS_THRESHOLD,
+) -> list[Anomaly]:
+    """Flag suspicious graph regions worth re-checking.
+
+    - documents with very low OCR confidence (likely OCR failure)
+    - chunks with empty extracted text
+    - documents with zero entities and zero citations (nothing extracted)
+    - entity hubs (mentioned everywhere; often a junk pattern)
+    """
+    anomalies: list[Anomaly] = []
+
+    documents = _nodes_of_kind(graph, "document")
+    for node, data in documents:
+        confidence = float(data.get("ocr_confidence_mean") or 0.0)
+        if confidence < low_confidence:
+            anomalies.append(
+                Anomaly(
+                    "low_confidence",
+                    node,
+                    f"ocr_confidence_mean={confidence:.2f} < {low_confidence}",
+                )
+            )
+
+        meaningful_out = 0
+        for _, target, edge in graph.out_edges(node, data=True):
+            if edge.get("kind") != "CONTAINS":
+                continue
+            for _, _, chunk_edge in graph.out_edges(target, data=True):
+                if chunk_edge.get("kind") in ("MENTIONS", "CITES"):
+                    meaningful_out += 1
+        if meaningful_out == 0:
+            anomalies.append(Anomaly("empty_document", node, "no entities or citations extracted"))
+
+    for node, data in _nodes_of_kind(graph, "chunk"):
+        if not (data.get("text") or "").strip():
+            anomalies.append(Anomaly("empty_chunk", node, "extracted text is empty"))
+
+    for node, data in _nodes_of_kind(graph, "entity"):
+        mentions = int(data.get("mentions", 0))
+        if mentions >= hub_mentions:
+            anomalies.append(
+                Anomaly("entity_hub", node, f"mentioned {mentions} times; possible junk pattern")
+            )
+
+    return anomalies

kg_ocr/graph/builder.py

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

kg_ocr/graph/traversal.py

@@ -0,0 +1,89 @@
+"""Traversal queries over the knowledge graph: entity/citation -> chunks, neighbors."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .builder import _normalize_entity
+
+
+def _entity_key(text: str) -> str:
+    return f"entity:{_normalize_entity(text)}"
+
+
+def _chunk_payload(graph: Any, chunk_node: str) -> dict[str, Any]:
+    data = graph.nodes[chunk_node]
+    doc_node = next(
+        (src for src, _, e in graph.in_edges(chunk_node, data=True) if e.get("kind") == "CONTAINS"),
+        None,
+    )
+    doc = graph.nodes[doc_node] if doc_node is not None else {}
+    return {
+        "chunk": chunk_node,
+        "chunk_index": data.get("chunk_index", 0),
+        "text": data.get("text", ""),
+        "source_path": doc.get("source_path", ""),
+        "ocr_engine": doc.get("ocr_engine", ""),
+        "ocr_confidence_mean": doc.get("ocr_confidence_mean", 0.0),
+        "timestamp": doc.get("timestamp", ""),
+    }
+
+
+def chunks_for_entity(graph: Any, entity_text: str) -> list[dict[str, Any]]:
+    """All chunks mentioning the given entity (case-insensitive)."""
+    key = _entity_key(entity_text)
+    if key not in graph:
+        return []
+    chunks = [
+        src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "MENTIONS"
+    ]
+    return [_chunk_payload(graph, c) for c in sorted(chunks)]
+
+
+def chunks_for_citation(graph: Any, identifier: str) -> list[dict[str, Any]]:
+    """All chunks citing the given identifier (e.g. '10.1038/nature12345')."""
+    key = identifier if identifier.startswith("citation:") else f"citation:{identifier}"
+    if key not in graph:
+        return []
+    chunks = [src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "CITES"]
+    return [_chunk_payload(graph, c) for c in sorted(chunks)]
+
+
+def related_entities(graph: Any, entity_text: str, limit: int = 10) -> list[tuple[str, int]]:
+    """Entities co-occurring with the given one, as (text, shared-chunk weight)."""
+    key = _entity_key(entity_text)
+    if key not in graph:
+        return []
+    neighbors: dict[str, int] = {}
+    for _, target, edge in graph.out_edges(key, data=True):
+        if edge.get("kind") == "CO_OCCURS":
+            neighbors[target] = edge.get("weight", 1)
+    for source, _, edge in graph.in_edges(key, data=True):
+        if edge.get("kind") == "CO_OCCURS":
+            neighbors[source] = max(neighbors.get(source, 0), edge.get("weight", 1))
+    ranked = sorted(
+        ((graph.nodes[n].get("text", n), w) for n, w in neighbors.items()),
+        key=lambda item: item[1],
+        reverse=True,
+    )
+    return ranked[:limit]
+
+
+def expand_context(graph: Any, entity_text: str, limit: int = 5) -> list[dict[str, Any]]:
+    """Chunks mentioning the entity or its strongest co-occurring neighbors.
+
+    This is the graph-RAG primitive: one hop of CO_OCCURS expansion, so a query
+    for 'BRCA1' also surfaces chunks that only mention its interactors.
+    """
+    seen: set[str] = set()
+    results: list[dict[str, Any]] = []
+
+    seeds = [entity_text] + [text for text, _ in related_entities(graph, entity_text, limit=limit)]
+    for seed in seeds:
+        for payload in chunks_for_entity(graph, seed):
+            if payload["chunk"] in seen:
+                continue
+            seen.add(payload["chunk"])
+            payload = {**payload, "via_entity": seed}
+            results.append(payload)
+    return results

kg_ocr/ocr/__init__.py

@@ -1,4 +1,4 @@
-from .extractor import get_screenshots
 from .batch_processor import extract_text
+from .extractor import get_screenshots
 
 __all__ = ["get_screenshots", "extract_text"]

kg_ocr/ocr/batch_processor.py

@@ -1,5 +1,5 @@
-from PIL import Image
 import pytesseract
+from PIL import Image
 
 
 def extract_text(images: list[str]) -> list[str]:

kg_ocr/ocr/constants.py

@@ -1,7 +1,5 @@
 from pathlib import Path
 
-import platform
-
 def_paths = {
     "Darwin": Path.home() / "Desktop",
     "Windows": Path.home() / "Pictures" / "Screenshots",

kg_ocr/ocr/extractor.py

@@ -1,11 +1,10 @@
 import platform
 from pathlib import Path
-from typing import Optional
 
 from .constants import def_paths, sc_pathpatterns
 
 
-def get_screenshots(path: Optional[str | Path] = None) -> list[str]:
+def get_screenshots(path: str | Path | None = None) -> list[str]:
     """Find screenshot files for the current OS."""
     if path is None:
         path = def_paths.get(platform.system(), Path.home())

kg_ocr/rag/__init__.py

@@ -1,3 +1,3 @@
-from .query import retrieve, ask_wllm
+from .query import ask_wllm, retrieve
 
 __all__ = ["retrieve", "ask_wllm"]

kg_ocr/rag/query.py

@@ -1,32 +1,47 @@
-from txtai.embeddings import Embeddings
-from txtai import LLM
-import litellm
-from dotenv import load_dotenv
-import os
+from __future__ import annotations
 
-load_dotenv()
+from typing import Any
 
-def retrieve(embeddings: Embeddings, query: str, limit: int = 3) -> list[dict]:
-    """Search embeddings and return results with scores"""
+DEFAULT_CHAT_MODEL = "openrouter/minimax/minimax-m2.5:free"
+
+SYSTEM_PROMPT = (
+    "Answer ONLY using the provided context. Cite which parts you're drawing "
+    "from. If the context doesn't cover something, say 'not in my documents'."
+)
+
+
+def retrieve(embeddings: Any, query: str, limit: int = 3) -> list[dict]:
+    """Search embeddings and return results with scores."""
     return embeddings.search(query, limit)
 
-def ask_wllm(embeddings: Embeddings, question: str, model: str = "openrouter/minimax/minimax-m2.5:free", limit: int = 3) -> str:
+
+def ask_wllm(
+    embeddings: Any,
+    question: str,
+    model: str = DEFAULT_CHAT_MODEL,
+    limit: int = 3,
+) -> str:
     """RAG: retrieve context from embeddings, then answer with an LLM."""
+    try:
+        import litellm
+    except ImportError as exc:
+        raise ImportError("ask_wllm needs litellm: uv sync --extra kg") from exc
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
     results = retrieve(embeddings, question, limit)
     context = "\n\n".join([r["text"] for r in results])
 
     response = litellm.completion(
         model=model,
         messages=[
-            {
-                "role": "system",
-                "content": "Answer ONLY using the provided context. Cite which parts you're drawing from. If the context doesn't cover something, say 'not in my documents'."
-            },
+            {"role": "system", "content": SYSTEM_PROMPT},
             {
                 "role": "user",
-                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}"
-            }
-        ]
+                "content": f"Context from my documents:\n{context}\n\nQuestion: {question}",
+            },
+        ],
     )
 
-    return response.choices[0].message.content
+    return response.choices[0].message.content