SnapOtter/packages/ai/python/ocr_runtime_entrypoint.py

"""Persistent JSON-line entrypoint for the isolated OCR runtime."""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any, BinaryIO, Callable, Mapping, TextIO

try:
    # Direct execution inside the immutable runtime generation.
    from ocr_runtime import OcrRuntime
except ModuleNotFoundError:  # pragma: no cover - exercised by package imports in tooling
    from .ocr_runtime import OcrRuntime


PROTOCOL_VERSION = 1
MAX_REQUEST_BYTES = 64 * 1024
MAX_ERROR_CHARS = 400


def configure_offline_environment() -> None:
    """Set offline switches before importing any optional model libraries."""
    os.environ["HF_HUB_OFFLINE"] = "1"
    os.environ["TRANSFORMERS_OFFLINE"] = "1"
    os.environ["SNAPOTTER_OCR_OFFLINE"] = "1"
    os.environ["NO_PROXY"] = "*"
    os.environ["no_proxy"] = "*"


def _request_id(request: Any) -> str:
    if isinstance(request, Mapping):
        value = request.get("requestId")
        if isinstance(value, str) and 0 < len(value) <= 128:
            return value
    return "unknown"


def _failure(request_id: str, code: str, message: str) -> dict[str, Any]:
    bounded = " ".join(str(message).split())[:MAX_ERROR_CHARS] or "OCR runtime request failed"
    return {
        "protocolVersion": PROTOCOL_VERSION,
        "requestId": request_id,
        "ok": False,
        "error": {"code": code, "message": bounded},
    }


def _parse_settings(value: str) -> dict[str, Any]:
    try:
        settings = json.loads(value)
    except json.JSONDecodeError as error:
        raise ValueError("OCR settings are not valid JSON") from error
    if not isinstance(settings, dict):
        raise ValueError("OCR settings must be a JSON object")
    return settings


def _parse_pages(value: str) -> tuple[tuple[int, Path], ...]:
    try:
        raw_pages = json.loads(value)
    except json.JSONDecodeError as error:
        raise ValueError("Rasterized pages are not valid JSON") from error
    if not isinstance(raw_pages, list) or not raw_pages:
        raise ValueError("Rasterized pages must be a non-empty JSON list")
    pages = []
    for index, item in enumerate(raw_pages):
        if isinstance(item, str):
            pages.append((index + 1, Path(item)))
            continue
        if (
            isinstance(item, dict)
            and isinstance(item.get("page"), int)
            and not isinstance(item.get("page"), bool)
            and isinstance(item.get("path"), str)
        ):
            pages.append((item["page"], Path(item["path"])))
            continue
        raise ValueError("Each rasterized page must be a path or {page, path} object")
    return tuple(pages)


def _validate_request(request: Any) -> tuple[str, str, list[str]]:
    if not isinstance(request, dict):
        raise ValueError("OCR runtime request must be an object")
    request_id = _request_id(request)
    if request_id == "unknown":
        raise ValueError("OCR runtime requestId is missing or invalid")
    protocol_version = request.get("protocolVersion")
    if type(protocol_version) is not int or protocol_version != PROTOCOL_VERSION:
        raise ValueError(f"Unsupported OCR runtime protocol version; expected {PROTOCOL_VERSION}")
    script = request.get("script")
    if script not in ("ocr", "ocr_pdf", "smoke"):
        raise ValueError("OCR runtime script must be ocr, ocr_pdf, or smoke")
    args = request.get("args")
    if not isinstance(args, list) or not all(isinstance(arg, str) for arg in args):
        raise ValueError("OCR runtime args must be a list of strings")
    if script == "smoke" and args:
        raise ValueError("OCR runtime smoke requires no arguments")
    if script != "smoke" and len(args) != 2:
        raise ValueError("OCR runtime scripts require exactly two arguments")
    return request_id, script, args


def _prepare_request(request: Any) -> tuple[str, str, tuple[Any, ...]]:
    """Validate and decode a frame without constructing model sessions."""
    request_id, script, args = _validate_request(request)
    if script == "smoke":
        return request_id, script, ()
    settings = _parse_settings(args[1])
    if script == "ocr":
        return request_id, script, (Path(args[0]), settings)
    return request_id, script, (_parse_pages(args[0]), settings)


def _process_prepared_request(
    prepared: tuple[str, str, tuple[Any, ...]], runtime: OcrRuntime
) -> dict[str, Any]:
    request_id, script, arguments = prepared
    try:
        if script == "smoke":
            result = runtime.smoke()
        elif script == "ocr":
            result = runtime.recognize_image(*arguments)
        else:
            result = runtime.recognize_pages(*arguments)
        return {
            "protocolVersion": PROTOCOL_VERSION,
            "requestId": request_id,
            "ok": True,
            "result": result,
        }
    except ValueError as error:
        return _failure(request_id, "invalid-request", str(error))
    except FileNotFoundError as error:
        return _failure(request_id, "file-not-found", str(error))
    except Exception as error:
        return _failure(request_id, "ocr-runtime-failed", str(error))


def process_request(
    request: Any,
    *,
    runtime: OcrRuntime | None = None,
    runtime_factory: Callable[[], OcrRuntime] = OcrRuntime,
) -> dict[str, Any]:
    request_id = _request_id(request)
    try:
        prepared = _prepare_request(request)
        request_id = prepared[0]
        active_runtime = runtime or runtime_factory()
        return _process_prepared_request(prepared, active_runtime)
    except ValueError as error:
        return _failure(request_id, "invalid-request", str(error))
    except FileNotFoundError as error:
        return _failure(request_id, "file-not-found", str(error))
    except Exception as error:
        return _failure(request_id, "ocr-runtime-failed", str(error))


def main(
    argv: list[str] | None = None,
    *,
    runtime_factory: Callable[[], OcrRuntime] = OcrRuntime,
    stdin: BinaryIO | None = None,
    stdout: TextIO | None = None,
    stderr: TextIO | None = None,
) -> int:
    configure_offline_environment()
    arguments = list(sys.argv[1:] if argv is None else argv)
    input_stream = sys.stdin.buffer if stdin is None else stdin
    output_stream = sys.stdout if stdout is None else stdout
    error_stream = sys.stderr if stderr is None else stderr

    if arguments == ["--smoke"]:
        try:
            smoke_result = runtime_factory().smoke()
        except Exception as error:
            error_stream.write(f"OCR runtime smoke failed: {error}\n")
            error_stream.flush()
            return 1
        output_stream.write(
            json.dumps(
                {"smoke": True, **smoke_result},
                ensure_ascii=False,
                separators=(",", ":"),
            )
            + "\n"
        )
        output_stream.flush()
        return 0
    if arguments:
        error_stream.write("Usage: ocr_runner.py [--smoke]\n")
        error_stream.flush()
        return 2

    runtime: OcrRuntime | None = None
    handled_requests = 0

    while True:
        raw = input_stream.readline(MAX_REQUEST_BYTES + 1)
        if not raw:
            if handled_requests == 0:
                response = _failure(
                    "unknown", "invalid-request", "OCR runtime request is empty"
                )
                output_stream.write(
                    json.dumps(response, ensure_ascii=False, separators=(",", ":"))
                    + "\n"
                )
                output_stream.flush()
                return 1
            return 0

        exit_after_response = False
        if len(raw) > MAX_REQUEST_BYTES:
            response = _failure(
                "unknown", "invalid-request", "OCR runtime request is too large"
            )
            # readline() leaves the rest of an oversized unterminated frame in
            # the stream. Exit after the bounded error instead of interpreting
            # the tail as another request and losing frame synchronization.
            exit_after_response = not raw.endswith(b"\n")
        else:
            try:
                request = json.loads(raw.decode("utf-8"))
            except (UnicodeDecodeError, json.JSONDecodeError) as error:
                response = _failure(
                    "unknown", "invalid-request", f"Malformed JSON request: {error}"
                )
            else:
                try:
                    prepared = _prepare_request(request)
                except ValueError as error:
                    response = _failure(
                        _request_id(request), "invalid-request", str(error)
                    )
                else:
                    if runtime is None:
                        try:
                            runtime = runtime_factory()
                        except Exception as error:
                            response = _failure(
                                _request_id(request), "ocr-runtime-failed", str(error)
                            )
                            exit_after_response = True
                        else:
                            response = _process_prepared_request(prepared, runtime)
                    else:
                        response = _process_prepared_request(prepared, runtime)

        output_stream.write(
            json.dumps(response, ensure_ascii=False, separators=(",", ":")) + "\n"
        )
        output_stream.flush()
        handled_requests += 1
        if exit_after_response:
            return 1


if __name__ == "__main__":
    raise SystemExit(main())