from __future__ import annotations import importlib.util import shutil import sys from dataclasses import dataclass from typing import Any from ocr_pipeline.config import settings from ocr_pipeline.onboarding import tesseract_install_hint from ocr_pipeline.utils.logging import get_logger logger = get_logger(__name__) @dataclass(frozen=True) class CheckResult: name: str status: str # "ok" | "warn" | "fail" detail: str hint: str | None = None def _has_module(name: str) -> bool: return importlib.util.find_spec(name) is not None def check_tesseract() -> CheckResult: path = shutil.which("tesseract") if path: return CheckResult("Tesseract", "ok", f"found at {path}") return CheckResult("Tesseract", "fail", "not on PATH", f"Install: {tesseract_install_hint()}") def check_paddle() -> CheckResult: if not _has_module("paddleocr"): return CheckResult("PaddleOCR", "warn", "not installed", "uv sync --extra paddle") try: from paddleocr import PaddleOCR # type: ignore[import-not-found] PaddleOCR(use_angle_cls=True, lang="en", use_gpu=False, show_log=False) return CheckResult("PaddleOCR", "ok", "initialized (en)") except Exception as exc: return CheckResult("PaddleOCR", "warn", f"init failed: {exc}") def check_torch() -> CheckResult: if not _has_module("torch"): return CheckResult("PyTorch", "warn", "not installed", "uv sync --extra tables") import torch # type: ignore[import-not-found] cuda = torch.cuda.is_available() mps = getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available() backend = "cuda" if cuda else "mps" if mps else "cpu" return CheckResult("PyTorch", "ok", f"backend={backend}") def check_tables() -> CheckResult: if not _has_module("transformers") or not _has_module("torch"): return CheckResult( "Table detector", "warn", "missing torch/transformers", "uv sync --extra tables" ) return CheckResult("Table detector", "ok", "transformers + torch present") def check_figures() -> CheckResult: if not _has_module("layoutparser"): return CheckResult( "Figure detector", "warn", "layoutparser not installed", "uv sync --extra figures" ) try: import layoutparser as lp # type: ignore[import-not-found] has_detectron = hasattr(lp, "Detectron2LayoutModel") has_auto = hasattr(lp, "AutoLayoutModel") if not (has_detectron or has_auto): return CheckResult( "Figure detector", "warn", "no layout model backend", "Install Detectron2 for your platform", ) return CheckResult("Figure detector", "ok", "layoutparser + backend present") except Exception as exc: return CheckResult("Figure detector", "warn", f"import failed: {exc}") def check_spacy() -> CheckResult: if not _has_module("spacy"): return CheckResult( "scispaCy NER", "warn", "spacy not installed", "uv sync --extra scientific" ) import spacy # type: ignore[import-not-found] try: spacy.load(settings.entities.model) return CheckResult("scispaCy NER", "ok", f"loaded {settings.entities.model}") except OSError: return CheckResult( "scispaCy NER", "warn", f"model {settings.entities.model!r} not downloaded", "uv pip install https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.4/en_core_sci_lg-0.5.4.tar.gz", ) def check_config() -> CheckResult: from ocr_pipeline.config import find_config_file path = find_config_file() if path is None: return CheckResult("Config", "warn", "no config.yaml found", "Run: ocr-pipeline init") return CheckResult("Config", "ok", f"loaded from {path}") def check_db() -> CheckResult: db_path = settings.watch_db_path try: db_path.parent.mkdir(parents=True, exist_ok=True) with db_path.open("a"): pass return CheckResult("Database", "ok", f"writable at {db_path}") except OSError as exc: return CheckResult("Database", "fail", f"not writable: {exc}") def check_output_dir() -> CheckResult: out = settings.output_dir try: out.mkdir(parents=True, exist_ok=True) return CheckResult("Output dir", "ok", str(out)) except OSError as exc: return CheckResult("Output dir", "fail", f"cannot create: {exc}") def check_python() -> CheckResult: return CheckResult( "Python", "ok", f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}", ) CHECKS = [ check_python, check_tesseract, check_paddle, check_torch, check_tables, check_figures, check_spacy, check_config, check_db, check_output_dir, ] def run_checks() -> list[CheckResult]: results: list[CheckResult] = [] for check in CHECKS: try: results.append(check()) except Exception as exc: # noqa: BLE001 - surfaced as a fail result, never propagated results.append(CheckResult(check.__name__, "fail", f"unexpected error: {exc}")) return results def render(results: list[CheckResult], console: Any) -> None: from rich.table import Table status_style = {"ok": "green", "warn": "yellow", "fail": "red"} icon = {"ok": "[green]OK[/green]", "warn": "[yellow]WARN[/yellow]", "fail": "[red]FAIL[/red]"} table = Table(title="ocr-pipeline health check") table.add_column("Status") table.add_column("Check", style="cyan") table.add_column("Detail") for r in results: table.add_row(icon[r.status], r.name, r.detail) console.print(table) failing = [r for r in results if r.status == "fail"] warning = [r for r in results if r.status == "warn"] if failing or warning: console.print() for r in failing + warning: if r.hint: console.print(f" [{status_style[r.status]}]{r.name}:[/] {r.hint}") def exit_code_for(results: list[CheckResult]) -> int: return 1 if any(r.status == "fail" for r in results) else 0