diff --git a/.env.example b/.env.example index 823ff74..ae9687f 100644 --- a/.env.example +++ b/.env.example @@ -23,4 +23,12 @@ WATCH_DB_PATH=./data/processed_files.db # Optional: GROBID for citation parsing GROBID_ENABLED=false -GROBID_URL=http://localhost:8070 \ No newline at end of file +GROBID_URL=http://localhost:8070 + +# Optional: Neo4j for `ocr-pipeline kg export -f neo4j` +NEO4J_URI=bolt://localhost:7687 +NEO4J_USER=neo4j +NEO4J_PASSWORD= + +# Optional: LLM provider keys for kg_ocr.rag (ask_wllm) +# OPENROUTER_API_KEY= \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index d6686a0..cced17d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Knowledge graph: `kg_ocr.graph` builds a NetworkX graph (documents, chunks, entities, citations, entity co-occurrence) from pipeline chunk markdown; `kg_ocr.graph.analyzer` provides summaries, top entities/citations, and anomaly detection (low OCR confidence, empty chunks/documents, entity hubs). +- Graph traversal: `chunks_for_entity`, `chunks_for_citation`, `related_entities`, `expand_context` (one-hop CO_OCCURS expansion for graph-RAG), exposed via `ocr-pipeline kg query`. +- `kg_ocr.embeddings.config_loader.KgConfig` for traversal tunables and Neo4j credentials from env. +- `kg_ocr.export`: node-link JSON (lossless round-trip), GraphML for Gephi/Cytoscape, and batched MERGE export to Neo4j via `Neo4jExporter`. +- CLI: `ocr-pipeline kg build|stats|query|export` subcommands with lazy kg_ocr imports. +- Optional extras: `kg` (networkx + txtai + litellm + python-dotenv), `neo4j` (driver). - `init` command: writes a minimal `config.yaml` with platform-correct defaults and checks for Tesseract. - `demo` command: runs the pipeline on a bundled sample screenshot, prints the output Markdown, and verifies the install in <30s. - `setup [basic|full]` command: executes install steps with confirmation (`--yes`/`--dry-run`). @@ -25,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - OCR engine default is now `tesseract` (was `paddleocr`). - `run` results table now shows Skipped count. - Run directories are created lazily on first write (no empty `run_NNN`). -- `kg_ocr` package marked experimental; deps under `--extra kg`, helpful ImportError on bare import. +- `kg_ocr` is now a functional knowledge-graph package (graph/export need only networkx); txtai/litellm imports are lazy so the package imports in slim environments. ### Fixed - ISBN regex no longer false-positives on years, URLs, job IDs, or dilutions. Now requires an explicit `ISBN` prefix and matches strict 13- or 10-digit structures. @@ -33,6 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `import ocr_pipeline` no longer depends on the caller's cwd. ### Removed +- Dead `src/ocr_pipeline/watch.py` shim (shadowed by the `watch/` package). - Unused core dependencies: `tqdm`, `scikit-image`, `sqlite-utils`, `xxhash`, `python-slugify`, `python-magic`, `legacy-cgi`. - Core dependency on `spacy` (moved to `scientific` extra). - Core dependency on `paddleocr`/`paddlepaddle` (moved to `paddle` extra). diff --git a/README.md b/README.md index 90cd324..7ef5f6a 100644 --- a/README.md +++ b/README.md @@ -1,90 +1,118 @@ # OCR Pipeline for Life Science Screenshots -[![Tests](https://img.shields.io/badge/tests-passing-brightgreen)](./.github/workflows/tests.yml) - Turns scientific screenshots into RAG-ready Markdown, with metadata attached. - OCR via PaddleOCR, falls back to Tesseract - Preprocessing: deskew, denoise, CLAHE, line removal - Figure/table detection - Entity extraction with scispaCy -- Citation matching +- Citation matching (DOI, PMID, arXiv, ISBN) - Chunking for embedding +- Knowledge graph over the output (`kg` commands) ## Quick start -The default install is small: Tesseract OCR only, no multi-GB ML downloads. - ```bash -uv sync # install the core pipeline -uv run ocr-pipeline setup # installs Tesseract if missing (asks first) -uv run ocr-pipeline init # write a config for this machine -uv run ocr-pipeline demo # verify everything works on a sample screenshot -uv run ocr-pipeline run # process your screenshots +uv sync +uv run ocr-pipeline setup # installs Tesseract if missing +uv run ocr-pipeline init # write a config +uv run ocr-pipeline demo # check it works +uv run ocr-pipeline run # process your screenshots ``` -Config is found in this order: `$OCR_PIPELINE_CONFIG`, `./config.yaml`, then the -per-user config directory (`~/Library/Application Support/ocr-pipeline` on macOS, -`~/.config/ocr-pipeline` on Linux). Running `run` with neither a config nor -`--input-dir` stops with guidance instead of scanning a default directory. +Default install is small: Tesseract only, no multi-GB ML downloads. -Use `--force` to reprocess files already recorded in the processing database. The -`--exclude` option is repeatable and prevents recursive scans from entering bundles -such as macOS `.photoslibrary` directories. +## Commands -### Optional capabilities +- `run` / `watch` — process a directory, or watch for new screenshots +- `status` / `stats` — what's been processed +- `doctor` — check engines, models, config +- `kg build|stats|query|export` — knowledge graph over the output +- `setup [basic|full]` — install dependencies step by step -Heavy ML features are opt-in extras and degrade gracefully when absent (a warning -is logged and that step is skipped): +Use `--force` to reprocess. `--exclude "*.photoslibrary/*"` keeps recursive +scans out of macOS photo bundles. + +
+Optional extras + +Heavy ML features are opt-in and degrade gracefully when missing (logged, step +skipped, run continues): | Extra | Provides | |-------|----------| -| `paddle` | PaddleOCR engine (`ocr.engine: paddleocr` or `auto`) | -| `scientific` | spaCy/scispaCy NER (`entities.enabled: true`) | +| `paddle` | PaddleOCR engine | +| `scientific` | spaCy/scispaCy NER | | `tables` | Table Transformer detection | -| `figures` | LayoutParser figure/caption detection (needs a Detectron2 build) | +| `figures` | LayoutParser figure/caption detection | | `full` | All of the above | +| `kg` | Knowledge graph (networkx, txtai, litellm) | +| `neo4j` | Neo4j export driver | ```bash -uv run ocr-pipeline setup full # runs: uv sync --extra full + downloads en_core_sci_lg +uv sync --extra full +# or guided: +uv run ocr-pipeline setup full ``` -`en_core_sci_lg` is distributed separately from scispaCy. The compatible scispaCy -0.5.4 release requires spaCy 3.7.x; if the model install conflicts with your spaCy -version, use a dedicated Python 3.11 environment as described in -`ocr-pipeline setup full --dry-run` output, and enable `entities` only there. +`en_core_sci_lg` is distributed separately from scispaCy and needs spaCy 3.7.x. +If the model install conflicts, use a dedicated Python 3.11 venv and enable +`entities` only there. -## kg_ocr (experimental) +
-The top-level `kg_ocr` package is a prototype txtai/litellm RAG interface over the -pipeline's output. It is not the supported interface (that is the `ocr-pipeline` -CLI) and its dependencies are not installed by default: +
+Knowledge graph + +Build a graph from pipeline output: documents, chunks, entities, citations, +entity co-occurrence. ```bash uv sync --extra kg + +uv run ocr-pipeline kg build -d data/ocr_output +uv run ocr-pipeline kg stats data/ocr_output/kg_graph.json + +# retrieval +uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --entity BRCA1 +uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --entity BRCA1 --expand +uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --citation 10.1038/nature12345 + +# export +uv run ocr-pipeline kg export data/ocr_output/kg_graph.json -f graphml +uv run ocr-pipeline kg export data/ocr_output/kg_graph.json -f neo4j ``` -## Docker +Anomaly detection flags low-confidence documents, empty chunks, and entity +hubs. Neo4j export needs `uv sync --extra neo4j` plus `NEO4J_URI`, +`NEO4J_USER`, `NEO4J_PASSWORD` env vars (or `--uri/--user/--password`). -The image ships the Tesseract-only core with a ready-made config that reads from -`/data/screenshots` and writes to `/data/output`: +
+ +
+Docker + +The image ships the Tesseract-only core. It reads `/data/screenshots` and +writes `/data/output`: ```bash docker build -t ocr-pipeline . docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline run docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline watch -# ML extras baked in: +# with ML extras: docker build --build-arg EXTRAS="--extra full" -t ocr-pipeline:full . ``` -## Config - -`ocr-pipeline init` writes a minimal working config; the example below shows every -knob for reference. +
-config.yaml +Config + +`ocr-pipeline init` writes a minimal config. Lookup order: +`$OCR_PIPELINE_CONFIG` → `./config.yaml` → per-user config dir. + +Every knob, for reference: ```yaml input: @@ -134,20 +162,6 @@ output: organize_by: "date_run" # date_run | source_dir | flat write_consolidated: true consolidated_filename: "all_ocr.md" - frontmatter: - - source_path - - source_hash - - timestamp - - ocr_engine - - ocr_confidence_mean - - language - - detected_entities - - entity_extraction_backend - - has_figures - - has_tables - - citations_found - - chunk_index - - total_chunks watch: enabled: true @@ -157,7 +171,8 @@ watch:
-## Output +
+Output example Each chunk is a Markdown file with YAML frontmatter: @@ -172,71 +187,61 @@ language: "en" detected_entities: ["GENE", "PROTEIN", "CHEMICAL"] has_figures: true has_tables: false -citations_found: ["DOI:10.1038/nature12345", "PMID:12345678"] +citations_found: ["doi:10.1038/nature12345", "pmid:PMID:12345678"] chunk_index: 0 total_chunks: 2 --- # Screenshot: SCR-20250115-gel.png -## Figures - -### Figure 1 -- BBox: [100, 200, 800, 600] -- Confidence: 0.92 -- Caption: "Western blot showing BRCA1 expression..." - -## Detected Entities - -- BRCA1 -- CRISPR -- β-actin - -## Citations - -- DOI: 10.1038/nature12345 -- PMID: 12345678 - ## Extracted Text **Western Blot Analysis of BRCA1 Expression** Lane 1: WT control Lane 2: BRCA1 KO (CRISPR) -Lane 3: BRCA1 KO + pBRCA1-WT rescue -Lane 4: BRCA1 KO + pBRCA1-C61G mutant - -Anti-BRCA1 (1:1000), Anti-β-actin (1:5000) +... ``` -Files land in `data/ocr_output//run_NNN/`. Individual chunk files are retained for RAG indexing. Set `output.write_consolidated: true` to also write `all_ocr.md` containing one frontmatter block and all source text grouped by image. +Files land in `//run_NNN/`. Chunk files are kept for RAG +indexing. `output.write_consolidated: true` also writes `all_ocr.md` with one +frontmatter block and all source text grouped by image. -## Life science specifics +
-Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas and units (µM, ng/mL, kb/Mb/Gb, °C, ×g) get normalized, scientific notation gets cleaned up (`1.5×10⁻³` → `1.5×10^-3`), and gel/blot figures get their captions pulled out separately. Citations are matched by regex for DOI, PMID, arXiv, PMC, and ISBN. +
+Models and dependencies -## Models - -Model weights are only fetched when the matching extra is installed and enabled. -First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`: +Weights download on first use, cached in `~/.cache/ocr_pipeline/`: - PaddleOCR models (~200MB) - scispaCy `en_core_sci_lg` (~800MB) - Table Transformer (~500MB) - LayoutParser PubLayNet (~300MB) -## Runtime requirements and health +Notes per capability: -The core pipeline can run with Tesseract alone. PaddleOCR, scientific NER, figure detection, and table detection are optional capabilities with heavyweight, platform-specific dependencies. The pipeline records the OCR engine actually used and the entity-extraction backend in generated frontmatter; inspect them after each run rather than assuming configured models loaded. +- **PaddleOCR:** pinned to the legacy 2.x API. Falls back to Tesseract when it + can't initialize; the fallback is logged and recorded in frontmatter. +- **Scientific NER:** without scispaCy + `en_core_sci_lg`, a conservative + regex extractor runs instead and the backend is marked accordingly. +- **Figures:** needs a Detectron2 build matching your Torch/Python. No + universal wheel exists, so it's never installed automatically. +- **Tables:** downloaded by Transformers on first use; needs `timm`. -- **PaddleOCR:** the project pins the legacy 2.x API used by the pipeline. Install the locked environment with `uv sync`; a startup fallback to Tesseract is logged when Paddle cannot initialize. -- **Scientific NER:** install a compatible `scispacy` distribution and the separately distributed `en_core_sci_lg` model before enabling production scientific NER. Without it, the pipeline uses its conservative regex fallback and marks the backend accordingly. -- **Figures:** `layoutparser`'s `Detectron2LayoutModel` requires a Detectron2 build matching your Torch/Python platform. It is intentionally not forced as a universal dependency because no single wheel supports every platform. -- **Tables:** the Table Transformer model is downloaded by Transformers on first use; ensure the selected model's optional dependencies (including `timm`, when required by that model revision) are installed in the runtime image. +`ocr-pipeline doctor` checks all of the above and prints remediation hints. -Use `ocr-pipeline config` to verify effective settings. A missing optional model is logged and produces empty results for that detector instead of failing a complete batch. +
+ +## Life science specifics + +Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas +and units (µM, ng/mL, kb/Mb/Gb, °C, ×g) get normalized, scientific notation +gets cleaned up (`1.5×10⁻³` → `1.5×10^-3`), and gel/blot figures get their +captions pulled out separately. ## TODO -- [ ] Build a knowledge graph from extracted entities/citations - +- [x] Build a knowledge graph from extracted entities/citations (see `kg` commands) +- [ ] Entity linking to MeSH/UniProt identifiers +- [ ] GROBID-based structured citation parsing (server optional, regex today) diff --git a/kg_ocr/__init__.py b/kg_ocr/__init__.py index 4cbaba6..28d03dc 100644 --- a/kg_ocr/__init__.py +++ b/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 diff --git a/kg_ocr/cli/__init__.py b/kg_ocr/cli/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/kg_ocr/cli/build_graph.py b/kg_ocr/cli/build_graph.py deleted file mode 100644 index 70db6b2..0000000 --- a/kg_ocr/cli/build_graph.py +++ /dev/null @@ -1,3 +0,0 @@ -# to do - -# add entrypoint to setuppy \ No newline at end of file diff --git a/kg_ocr/cli/export_graph.py b/kg_ocr/cli/export_graph.py deleted file mode 100644 index d5364a9..0000000 --- a/kg_ocr/cli/export_graph.py +++ /dev/null @@ -1 +0,0 @@ -# to do \ No newline at end of file diff --git a/kg_ocr/cli/process_screenshots.py b/kg_ocr/cli/process_screenshots.py deleted file mode 100644 index d5364a9..0000000 --- a/kg_ocr/cli/process_screenshots.py +++ /dev/null @@ -1 +0,0 @@ -# to do \ No newline at end of file diff --git a/kg_ocr/embeddings/config_loader.py b/kg_ocr/embeddings/config_loader.py index 0ce7cef..0009d39 100644 --- a/kg_ocr/embeddings/config_loader.py +++ b/kg_ocr/embeddings/config_loader.py @@ -1 +1,31 @@ -# in future prefer a config, especially for graph traversal \ No newline at end of file +"""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", ""), + ) diff --git a/kg_ocr/embeddings/indexer.py b/kg_ocr/embeddings/indexer.py index 8a20037..d1c897e 100644 --- a/kg_ocr/embeddings/indexer.py +++ b/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 - - - diff --git a/kg_ocr/export/__init__.py b/kg_ocr/export/__init__.py index e69de29..4f0f980 100644 --- a/kg_ocr/export/__init__.py +++ b/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"] diff --git a/kg_ocr/export/neo4j_exporter.py b/kg_ocr/export/neo4j_exporter.py index e69de29..0995381 100644 --- a/kg_ocr/export/neo4j_exporter.py +++ b/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, + ) diff --git a/kg_ocr/graph/__init__.py b/kg_ocr/graph/__init__.py index e69de29..3ece7fb 100644 --- a/kg_ocr/graph/__init__.py +++ b/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", +] diff --git a/kg_ocr/graph/analyzer.py b/kg_ocr/graph/analyzer.py index f74afe7..f425b29 100644 --- a/kg_ocr/graph/analyzer.py +++ b/kg_ocr/graph/analyzer.py @@ -1 +1,117 @@ -# also anomaly detection here \ No newline at end of file +"""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 diff --git a/kg_ocr/graph/builder.py b/kg_ocr/graph/builder.py index e69de29..954f3f3 100644 --- a/kg_ocr/graph/builder.py +++ b/kg_ocr/graph/builder.py @@ -0,0 +1,193 @@ +"""Build a knowledge graph from OCR pipeline markdown output. + +Schema +------ +Nodes + doc: Document (source image + OCR metadata) + chunk:: Chunk (text segment of a document) + entity: Entity (gene/protein/chemical mention) + citation: 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) diff --git a/kg_ocr/graph/traversal.py b/kg_ocr/graph/traversal.py new file mode 100644 index 0000000..c6e59ea --- /dev/null +++ b/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 diff --git a/kg_ocr/ocr/__init__.py b/kg_ocr/ocr/__init__.py index f20aca4..9afc276 100644 --- a/kg_ocr/ocr/__init__.py +++ b/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"] diff --git a/kg_ocr/ocr/batch_processor.py b/kg_ocr/ocr/batch_processor.py index 88934f9..60da7d8 100644 --- a/kg_ocr/ocr/batch_processor.py +++ b/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]: diff --git a/kg_ocr/ocr/constants.py b/kg_ocr/ocr/constants.py index ab6ee02..a7b6991 100644 --- a/kg_ocr/ocr/constants.py +++ b/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", diff --git a/kg_ocr/ocr/extractor.py b/kg_ocr/ocr/extractor.py index 0d55512..465cda6 100644 --- a/kg_ocr/ocr/extractor.py +++ b/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()) diff --git a/kg_ocr/rag/__init__.py b/kg_ocr/rag/__init__.py index 645b12d..417c5b1 100644 --- a/kg_ocr/rag/__init__.py +++ b/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"] diff --git a/kg_ocr/rag/query.py b/kg_ocr/rag/query.py index dd09627..26fa545 100644 --- a/kg_ocr/rag/query.py +++ b/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 \ No newline at end of file + return response.choices[0].message.content diff --git a/pyproject.toml b/pyproject.toml index 5e65f39..bebdc3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,10 +63,14 @@ scientific = [ "spacy>=3.7.0; python_version >= '3.12'", ] kg = [ + "networkx>=3.2", "txtai>=7.4.0", "litellm>=1.40.0", "python-dotenv>=1.0.0", ] +neo4j = [ + "neo4j>=5.0", +] grobid = [] full = [ "ocr-pipeline[paddle,tables,figures,scientific]", diff --git a/src/ocr_pipeline/__pycache__/cli.cpython-313.pyc b/src/ocr_pipeline/__pycache__/cli.cpython-313.pyc index e31b9fd..3190a7a 100644 Binary files a/src/ocr_pipeline/__pycache__/cli.cpython-313.pyc and b/src/ocr_pipeline/__pycache__/cli.cpython-313.pyc differ diff --git a/src/ocr_pipeline/cli.py b/src/ocr_pipeline/cli.py index d4653ff..bbe532b 100644 --- a/src/ocr_pipeline/cli.py +++ b/src/ocr_pipeline/cli.py @@ -28,10 +28,27 @@ app = typer.Typer( help="OCR pipeline for screenshots -> RAG-ready Markdown", add_completion=False, ) +kg_app = typer.Typer( + name="kg", + help="Knowledge graph over OCR output (needs: uv sync --extra kg)", + add_completion=False, +) +app.add_typer(kg_app, name="kg") console = Console() logger = get_logger(__name__) +def _load_kg(): + """Import kg_ocr lazily so the base CLI works without the kg extra.""" + try: + from kg_ocr import export as kg_export + from kg_ocr import graph as kg_graph + except ImportError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + return kg_graph, kg_export + + def _ensure_configured(input_dir: Path | None) -> None: """Stop early with guidance instead of scanning a default directory.""" if input_dir is not None or find_config_file() is not None: @@ -296,6 +313,169 @@ def config(): console.print(settings.model_dump_json(indent=2)) +@kg_app.command("build") +def kg_build( + output_dir: Path = typer.Option( + ..., "--output-dir", "-d", help="Pipeline output directory with chunk markdown files" + ), + save: Path | None = typer.Option( + None, "--save", "-s", help="Graph JSON path (default: /kg_graph.json)" + ), +): + """Build a knowledge graph from pipeline output and save it as JSON.""" + kg_graph, kg_export = _load_kg() + if not output_dir.is_dir(): + console.print(f"[red]Not a directory:[/red] {output_dir}") + raise typer.Exit(code=2) + graph = kg_graph.build_from_directory(output_dir) + save = save or (output_dir / "kg_graph.json") + kg_export.export_json(graph, save) + + counts = kg_graph.summary(graph) + table = Table(title=f"Knowledge graph -> {save}") + table.add_column("Kind", style="cyan") + table.add_column("Count", style="green") + for kind, count in sorted(counts.items()): + table.add_row(kind, str(count)) + console.print(table) + + +@kg_app.command("stats") +def kg_stats( + graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"), + top: int = typer.Option(10, "--top", "-n", help="Rows per top-list"), +): + """Summarize a graph: counts, top entities/citations, anomalies.""" + kg_graph, kg_export = _load_kg() + if not graph_path.is_file(): + console.print(f"[red]Graph file not found:[/red] {graph_path}") + raise typer.Exit(code=2) + graph = kg_export.load_json(graph_path) + + counts = kg_graph.summary(graph) + table = Table(title="Graph summary") + table.add_column("Kind", style="cyan") + table.add_column("Count", style="green") + for kind, count in sorted(counts.items()): + table.add_row(kind, str(count)) + console.print(table) + + entities = kg_graph.top_entities(graph, limit=top) + if entities: + table = Table(title=f"Top {top} entities") + table.add_column("Entity", style="cyan") + table.add_column("Mentions", style="green") + for text, mentions in entities: + table.add_row(text, str(mentions)) + console.print(table) + + citations = kg_graph.top_citations(graph, limit=top) + if citations: + table = Table(title=f"Top {top} citations") + table.add_column("Type", style="cyan") + table.add_column("Identifier") + table.add_column("Citing chunks", style="green") + for ctype, identifier, citing in citations: + table.add_row(ctype, identifier, str(citing)) + console.print(table) + + anomalies = kg_graph.detect_anomalies(graph) + if anomalies: + console.print(f"\n[yellow]{len(anomalies)} anomalies:[/yellow]") + for anomaly in anomalies[:top]: + console.print(f" [dim]{anomaly.kind}[/dim] {anomaly.node}: {anomaly.detail}") + + +@kg_app.command("query") +def kg_query( + graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"), + entity: str | None = typer.Option(None, "--entity", "-e", help="Entity to look up"), + citation: str | None = typer.Option( + None, "--citation", "-c", help="Citation identifier (e.g. 10.1038/nature12345)" + ), + expand: bool = typer.Option( + False, "--expand", "-x", help="Include chunks from co-occurring entities" + ), + limit: int = typer.Option(5, "--limit", "-n", help="Max chunks to show"), +): + """Retrieve chunks by entity or citation; --expand adds neighbor chunks.""" + kg_graph, kg_export = _load_kg() + if not graph_path.is_file(): + console.print(f"[red]Graph file not found:[/red] {graph_path}") + raise typer.Exit(code=2) + if not entity and not citation: + console.print("[red]Give --entity or --citation.[/red]") + raise typer.Exit(code=2) + graph = kg_export.load_json(graph_path) + + if entity: + chunks = ( + kg_graph.expand_context(graph, entity, limit=limit) + if expand + else kg_graph.chunks_for_entity(graph, entity) + ) + related = kg_graph.related_entities(graph, entity) + if related: + console.print( + "[dim]Related entities: " + + ", ".join(f"{text} ({weight})" for text, weight in related[:5]) + + "[/dim]" + ) + else: + chunks = kg_graph.chunks_for_citation(graph, citation or "") + + if not chunks: + console.print("[yellow]No matching chunks.[/yellow]") + return + + for payload in chunks[:limit]: + via = f" via {payload['via_entity']}" if payload.get("via_entity") else "" + console.print( + f"\n[bold]{Path(str(payload['source_path'])).name}[/bold] " + f"chunk {payload['chunk_index']}{via} " + f"[dim]({payload['ocr_engine']} {payload['ocr_confidence_mean']:.2f})[/dim]" + ) + excerpt = " ".join(str(payload["text"]).split())[:300] + console.print(f" {excerpt}") + + +@kg_app.command("export") +def kg_export_cmd( + graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"), + fmt: str = typer.Option("graphml", "--format", "-f", help="graphml | neo4j"), + out: Path | None = typer.Option(None, "--out", "-o", help="Output path for graphml"), + uri: str | None = typer.Option(None, "--uri", help="Neo4j bolt URI (or NEO4J_URI)"), + user: str | None = typer.Option(None, "--user", help="Neo4j user (or NEO4J_USER)"), + password: str | None = typer.Option( + None, "--password", help="Neo4j password (or NEO4J_PASSWORD)" + ), +): + """Export a graph JSON to GraphML or push it into Neo4j.""" + _, kg_export = _load_kg() + if not graph_path.is_file(): + console.print(f"[red]Graph file not found:[/red] {graph_path}") + raise typer.Exit(code=2) + graph = kg_export.load_json(graph_path) + + if fmt == "graphml": + out = out or graph_path.with_suffix(".graphml") + kg_export.export_graphml(graph, out) + console.print(f"[green]GraphML written:[/green] {out}") + elif fmt == "neo4j": + try: + with kg_export.Neo4jExporter(uri=uri, user=user, password=password) as exporter: + pushed = exporter.push(graph) + except ImportError as exc: + console.print(f"[red]{exc}[/red]") + raise typer.Exit(code=2) from exc + console.print( + f"[green]Pushed to Neo4j:[/green] {pushed['nodes']} nodes, {pushed['edges']} edges" + ) + else: + console.print(f"[red]Unknown format:[/red] {fmt} (choose graphml or neo4j)") + raise typer.Exit(code=2) + + def main(): app() diff --git a/src/ocr_pipeline/watch.py b/src/ocr_pipeline/watch.py deleted file mode 100644 index d1bb487..0000000 --- a/src/ocr_pipeline/watch.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility module for the canonical watcher package.""" - -from .watch.watcher import ScreenshotHandler, Watcher - -__all__ = ["ScreenshotHandler", "Watcher"] diff --git a/tests/test_graph.py b/tests/test_graph.py index e69de29..0164b4a 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +nx = pytest.importorskip("networkx", reason="kg extra not installed") + +from kg_ocr.export import export_graphml, export_json, load_json # noqa: E402 +from kg_ocr.graph import ( # noqa: E402 + build_from_directory, + build_graph, + detect_anomalies, + find_chunk_files, + parse_chunk_file, + summary, + top_citations, + top_cooccurrences, + top_entities, +) + +CHUNK_MD = """--- +source_path: /shots/SCR-001.png +source_hash: aaaabbbbccccdddd +timestamp: '2026-07-19T10:00:00+00:00' +ocr_engine: paddleocr +ocr_confidence_mean: 0.92 +language: en +detected_entities: +- BRCA1 +- PARP +citations_found: +- doi:10.1038/nature12345 +chunk_index: 0 +total_chunks: 1 +--- + +# Screenshot: SCR-001.png + +## Extracted Text + +BRCA1 and PARP inhibition in ovarian cancer. + +## Citations Found + +| Type | Identifier | +""" + +CHUNK_MD_2 = """--- +source_path: /shots/SCR-002.png +source_hash: eeeeffff00001111 +timestamp: '2026-07-19T10:01:00+00:00' +ocr_engine: tesseract +ocr_confidence_mean: 0.2 +language: en +detected_entities: +- BRCA1 +citations_found: [] +chunk_index: 0 +total_chunks: 1 +--- + +# Screenshot: SCR-002.png + +## Extracted Text + +BRCA1 western blot. +""" + + +@pytest.fixture() +def chunk_dir(tmp_path: Path) -> Path: + run_dir = tmp_path / "out" / "2026-07-19" / "run_001" + run_dir.mkdir(parents=True) + (run_dir / "SCR-001_aaaabbbbcccc.md").write_text(CHUNK_MD, encoding="utf-8") + (run_dir / "SCR-002_eeeeffff0000.md").write_text(CHUNK_MD_2, encoding="utf-8") + return tmp_path / "out" + + +def test_parse_chunk_file(chunk_dir: Path) -> None: + path = next(chunk_dir.rglob("SCR-001_*.md")) + record = parse_chunk_file(path) + assert record is not None + assert record.source_hash == "aaaabbbbccccdddd" + assert record.ocr_confidence_mean == pytest.approx(0.92) + assert record.entities == ["BRCA1", "PARP"] + assert record.citations == ["doi:10.1038/nature12345"] + assert "BRCA1 and PARP" in record.text + + +def test_parse_chunk_file_rejects_plain_markdown(tmp_path: Path) -> None: + plain = tmp_path / "notes.md" + plain.write_text("# just notes\n", encoding="utf-8") + assert parse_chunk_file(plain) is None + + +def test_find_chunk_files_skips_consolidated(chunk_dir: Path) -> None: + (chunk_dir / "all_ocr.md").write_text("# consolidated\n", encoding="utf-8") + names = [p.name for p in find_chunk_files(chunk_dir)] + assert "all_ocr.md" not in names + assert len(names) == 2 + + +def test_build_graph_structure(chunk_dir: Path) -> None: + graph = build_from_directory(chunk_dir) + counts = summary(graph) + + assert counts["nodes:document"] == 2 + assert counts["nodes:chunk"] == 2 + assert counts["nodes:entity"] == 2 # brca1, parp (normalized) + assert counts["nodes:citation"] == 1 + assert counts["edges:CONTAINS"] == 2 + assert counts["edges:MENTIONS"] == 3 + assert counts["edges:CITES"] == 1 + assert counts["edges:CO_OCCURS"] == 1 # brca1 <-> parp + + assert graph.has_edge("doc:aaaabbbbccccdddd", "chunk:aaaabbbbccccdddd:0") + assert graph.has_edge("chunk:aaaabbbbccccdddd:0", "entity:brca1") + + +def test_top_entities_and_cooccurrence(chunk_dir: Path) -> None: + graph = build_from_directory(chunk_dir) + top = top_entities(graph) + assert top[0] == ("BRCA1", 2) + pairs = top_cooccurrences(graph) + assert pairs == [("BRCA1", "PARP", 1)] + + +def test_top_citations(chunk_dir: Path) -> None: + graph = build_from_directory(chunk_dir) + rows = top_citations(graph) + assert rows == [("doi", "10.1038/nature12345", 1)] + + +def test_detect_anomalies(chunk_dir: Path) -> None: + graph = build_from_directory(chunk_dir) + kinds = {a.kind for a in detect_anomalies(graph)} + # SCR-002 has confidence 0.2 (< 0.4) -> low_confidence + assert "low_confidence" in kinds + + +def test_json_round_trip(chunk_dir: Path, tmp_path: Path) -> None: + graph = build_from_directory(chunk_dir) + out = export_json(graph, tmp_path / "g.json") + restored = load_json(out) + assert summary(restored) == summary(graph) + assert restored.nodes["entity:brca1"]["mentions"] == 2 + + +def test_graphml_export(chunk_dir: Path, tmp_path: Path) -> None: + graph = build_from_directory(chunk_dir) + out = export_graphml(graph, tmp_path / "g.graphml") + content = out.read_text(encoding="utf-8") + assert "graphml" in content + + +def test_empty_directory(tmp_path: Path) -> None: + graph = build_from_directory(tmp_path) + assert summary(graph) == {} + + +def test_build_graph_from_records_directly() -> None: + from kg_ocr.graph import ChunkRecord + + record = ChunkRecord( + source_hash="x" * 16, + source_path="/a.png", + chunk_index=0, + total_chunks=1, + text="hello", + entities=["TP53", "tp53"], # dedupe via normalization + citations=["pmid:PMID:12345678"], + ) + graph = build_graph([record]) + counts = summary(graph) + assert counts["nodes:entity"] == 1 + assert graph.nodes["entity:tp53"]["mentions"] == 2 + assert graph.nodes["citation:PMID:12345678"]["citation_type"] == "pmid" + + +def test_kg_cli_build_and_stats(chunk_dir: Path, tmp_path: Path) -> None: + from typer.testing import CliRunner + + from ocr_pipeline.cli import app + + runner = CliRunner() + save = tmp_path / "g.json" + + result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)]) + assert result.exit_code == 0, result.output + assert save.is_file() + + result = runner.invoke(app, ["kg", "stats", str(save)]) + assert result.exit_code == 0, result.output + assert "BRCA1" in result.output + assert "low_confidence" in result.output + + out_graphml = tmp_path / "g.graphml" + result = runner.invoke( + app, ["kg", "export", str(save), "-f", "graphml", "-o", str(out_graphml)] + ) + assert result.exit_code == 0, result.output + assert out_graphml.is_file() + + +def test_chunks_for_entity(chunk_dir: Path) -> None: + from kg_ocr.graph import chunks_for_entity + + graph = build_from_directory(chunk_dir) + chunks = chunks_for_entity(graph, "brca1") # case-insensitive + assert len(chunks) == 2 + assert {c["source_path"] for c in chunks} == {"/shots/SCR-001.png", "/shots/SCR-002.png"} + assert chunks_for_entity(graph, "NONEXISTENT") == [] + + +def test_chunks_for_citation(chunk_dir: Path) -> None: + from kg_ocr.graph import chunks_for_citation + + graph = build_from_directory(chunk_dir) + chunks = chunks_for_citation(graph, "10.1038/nature12345") + assert len(chunks) == 1 + assert "BRCA1 and PARP" in chunks[0]["text"] + assert chunks_for_citation(graph, "nope") == [] + + +def test_related_entities(chunk_dir: Path) -> None: + from kg_ocr.graph import related_entities + + graph = build_from_directory(chunk_dir) + assert related_entities(graph, "BRCA1") == [("PARP", 1)] + assert related_entities(graph, "PARP") == [("BRCA1", 1)] + assert related_entities(graph, "NONEXISTENT") == [] + + +def test_expand_context(chunk_dir: Path) -> None: + from kg_ocr.graph import expand_context + + graph = build_from_directory(chunk_dir) + chunks = expand_context(graph, "PARP") + # PARP's own chunk + BRCA1 neighbor chunks (SCR-001 deduped) + assert len(chunks) == 2 + via = {c["via_entity"] for c in chunks} + assert "PARP" in via + assert "BRCA1" in via + + +def test_kg_cli_query(chunk_dir: Path, tmp_path: Path) -> None: + from typer.testing import CliRunner + + from ocr_pipeline.cli import app + + runner = CliRunner() + save = tmp_path / "g.json" + result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)]) + assert result.exit_code == 0, result.output + + result = runner.invoke(app, ["kg", "query", str(save), "--entity", "PARP"]) + assert result.exit_code == 0, result.output + assert "SCR-001.png" in result.output + assert "Related entities" in result.output + + result = runner.invoke(app, ["kg", "query", str(save), "--citation", "10.1038/nature12345"]) + assert result.exit_code == 0, result.output + assert "SCR-001.png" in result.output + + result = runner.invoke(app, ["kg", "query", str(save)]) + assert result.exit_code == 2 # neither --entity nor --citation diff --git a/tests/test_indexer.py b/tests/test_indexer.py index e69de29..0c4109b 100644 --- a/tests/test_indexer.py +++ b/tests/test_indexer.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import sys +import types + +import pytest + +from kg_ocr.embeddings import create_and_index +from kg_ocr.rag import retrieve + + +class _FakeEmbeddings: + instances: list[_FakeEmbeddings] = [] + + def __init__(self, config): + self.config = config + self.indexed: list[str] | None = None + _FakeEmbeddings.instances.append(self) + + def index(self, data): + self.indexed = data + + +def _install_fake_txtai(monkeypatch: pytest.MonkeyPatch) -> None: + fake_txtai = types.ModuleType("txtai") + fake_embeddings_mod = types.ModuleType("txtai.embeddings") + fake_embeddings_mod.Embeddings = _FakeEmbeddings # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "txtai", fake_txtai) + monkeypatch.setitem(sys.modules, "txtai.embeddings", fake_embeddings_mod) + + +def test_create_and_index_config_and_data(monkeypatch: pytest.MonkeyPatch) -> None: + _install_fake_txtai(monkeypatch) + emb = create_and_index(["alpha", "beta"], model="some-model") + assert emb.config["path"] == "some-model" + assert emb.config["content"] is True + assert emb.config["hybrid"] is True + assert emb.indexed == ["alpha", "beta"] + + +def test_create_and_index_raises_without_txtai(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "txtai", None) + monkeypatch.setitem(sys.modules, "txtai.embeddings", None) + with pytest.raises(ImportError, match="txtai"): + create_and_index(["x"]) + + +def test_retrieve_passes_through() -> None: + class FakeEmb: + def search(self, query, limit): + return [{"text": f"{query}-{i}", "score": 1.0} for i in range(limit)] + + out = retrieve(FakeEmb(), "q", limit=2) + assert out == [{"text": "q-0", "score": 1.0}, {"text": "q-1", "score": 1.0}] diff --git a/uv.lock b/uv.lock index 973ff02..090568c 100644 --- a/uv.lock +++ b/uv.lock @@ -2349,6 +2349,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/82/0340caa499416c78e5d8f5f05947ae4bc3cba53c9f038ab6e9ed964e22f1/nbformat-5.10.4-py3-none-any.whl", hash = "sha256:3b48d6c8fbca4b299bf3982ea7db1af21580e4fec269ad087b9e81588891200b", size = 78454 }, ] +[[package]] +name = "neo4j" +version = "6.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/f4/aaa4ac19adae4b01bc742b63afd2672a77e7351566f02721e713e4b863ee/neo4j-6.2.0.tar.gz", hash = "sha256:e1e246b65b572bd8ea97f9e0e721b7d40a5ce53e53d0007c29aef63e4f9124d9", size = 241459 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/cf/1c3795866cefaac6e648d4e98c373cafd97810f6e317c307371007ab4abb/neo4j-6.2.0-py3-none-any.whl", hash = "sha256:b87abdd13a5cc2e3bd51026926c2f20ac38fa3febe98c340520dce19e97388d0", size = 327824 }, +] + [[package]] name = "networkx" version = "3.6.1" @@ -2692,9 +2704,13 @@ full = [ ] kg = [ { name = "litellm" }, + { name = "networkx" }, { name = "python-dotenv" }, { name = "txtai" }, ] +neo4j = [ + { name = "neo4j" }, +] paddle = [ { name = "paddleocr" }, { name = "paddlepaddle" }, @@ -2728,6 +2744,8 @@ requires-dist = [ { name = "litellm", marker = "extra == 'kg'", specifier = ">=1.40.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" }, { name = "nbformat", specifier = ">=5.9.0" }, + { name = "neo4j", marker = "extra == 'neo4j'", specifier = ">=5.0" }, + { name = "networkx", marker = "extra == 'kg'", specifier = ">=3.2" }, { name = "numpy", marker = "python_full_version < '3.12'", specifier = ">=1.26.0,<2.0" }, { name = "numpy", marker = "python_full_version >= '3.12'", specifier = ">=1.26.0" }, { name = "ocr-pipeline", extras = ["paddle", "tables", "figures", "scientific"], marker = "extra == 'full'" }, @@ -3863,6 +3881,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101 }, ] +[[package]] +name = "pytz" +version = "2026.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141 }, +] + [[package]] name = "pywin32" version = "312"