Octop/tests/unit/knowledge/test_ocr.py

"""Tests for optional knowledge-base OCR configuration and routing."""

from __future__ import annotations

import logging
from pathlib import Path
from types import SimpleNamespace

import pytest
from pypdf import PdfWriter

from octop.infra.knowledge import ocr
from octop.infra.knowledge.parse import parse_document


def test_warns_when_enabled_local_ocr_deps_are_missing(
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    monkeypatch.setattr(ocr, "local_ocr_deps_available", lambda: False)
    monkeypatch.setattr(ocr, "pdf_ocr_deps_available", lambda: False)
    values = {
        "knowledge_ocr_enabled": "true",
        "knowledge_ocr_backend": "onnx",
    }

    with caplog.at_level(logging.WARNING):
        ocr.warn_if_local_ocr_deps_missing(values.get)

    assert "knowledge-ocr" in caplog.text


def test_skips_ocr_dep_warning_when_runtime_is_present_or_unused(
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    monkeypatch.setattr(ocr, "local_ocr_deps_available", lambda: True)
    enabled = {"knowledge_ocr_enabled": "true", "knowledge_ocr_backend": "onnx"}
    disabled = {"knowledge_ocr_enabled": "false", "knowledge_ocr_backend": "onnx"}
    remote = {
        "knowledge_ocr_enabled": "true",
        "knowledge_ocr_backend": "remote",
        "knowledge_ocr_model": "vision-1",
        "knowledge_ocr_provider_id": "7",
    }

    with caplog.at_level(logging.WARNING):
        ocr.warn_if_local_ocr_deps_missing(enabled.get)
        ocr.warn_if_local_ocr_deps_missing(disabled.get)
        monkeypatch.setattr(ocr, "local_ocr_deps_available", lambda: False)
        ocr.warn_if_local_ocr_deps_missing(remote.get)

    assert caplog.records == []


def test_ocr_disabled_by_default() -> None:
    capability = ocr.get_ocr_capability(lambda _key: None)

    assert capability["enabled"] is False
    assert capability["backend"] == "onnx"
    assert capability["usable"] is False


def test_set_local_ocr_settings(monkeypatch: pytest.MonkeyPatch) -> None:
    values: dict[str, str] = {}
    monkeypatch.setattr(ocr, "local_ocr_deps_available", lambda: True)

    ocr.set_ocr_settings(
        values.__setitem__,
        enabled=True,
        backend="onnx",
        model=None,
        provider_id=None,
    )

    assert values == {
        "knowledge_ocr_enabled": "true",
        "knowledge_ocr_backend": "onnx",
        "knowledge_ocr_model": "rapidocr",
        "knowledge_ocr_provider_id": "",
    }


def test_remote_ocr_requires_image_capable_model() -> None:
    provider = SimpleNamespace(
        enabled=True,
        api_key="secret",
        base_url="https://example.test/v1",
        name="Remote",
        get_models=lambda: [
            {"id": "text-only", "input": ["text"]},
            {"id": "vision-1", "input": ["text", "image"]},
        ],
    )
    repo = SimpleNamespace(get=lambda provider_id: provider if provider_id == 7 else None)
    values: dict[str, str] = {}

    with pytest.raises(ValueError, match="provider is not ready"):
        ocr.set_ocr_settings(
            values.__setitem__,
            enabled=True,
            backend="remote",
            model="text-only",
            provider_id="7",
            provider_repo=repo,
        )

    ocr.set_ocr_settings(
        values.__setitem__,
        enabled=True,
        backend="remote",
        model="vision-1",
        provider_id="7",
        provider_repo=repo,
    )
    assert ocr.get_ocr_capability(values.get, repo)["usable"] is True


def test_image_and_blank_pdf_use_ocr(tmp_path: Path) -> None:
    image = tmp_path / "scan.png"
    image.write_bytes(b"not-decoded-by-the-fake")
    pdf = tmp_path / "scan.pdf"
    writer = PdfWriter()
    writer.add_blank_page(width=72, height=72)
    with pdf.open("wb") as output:
        writer.write(output)
    calls: list[Path] = []

    def fake_ocr(path: Path) -> str:
        calls.append(path)
        return "recognized text"

    assert parse_document(image, ocr=fake_ocr) == "recognized text"
    assert parse_document(pdf, ocr=fake_ocr) == "recognized text"
    assert calls == [image, pdf]


def test_text_pdf_does_not_require_ocr(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    pdf = tmp_path / "text.pdf"
    pdf.write_bytes(b"placeholder")

    class Page:
        @staticmethod
        def extract_text() -> str:
            return "embedded text"

    monkeypatch.setattr(
        "pypdf.PdfReader",
        lambda _path: SimpleNamespace(pages=[Page()]),
    )

    def unexpected_ocr(_path: Path) -> str:
        raise AssertionError("OCR must not run for a text PDF")

    assert parse_document(pdf, ocr=unexpected_ocr) == "embedded text"


def test_image_requires_enabled_ocr(tmp_path: Path) -> None:
    image = tmp_path / "scan.jpg"
    image.write_bytes(b"image")

    with pytest.raises(RuntimeError, match="OCR is not enabled"):
        parse_document(image)


def test_local_ocr_joins_rapidocr_text_lines(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    image = tmp_path / "scan.png"
    image.write_bytes(b"image")

    def engine(_data: bytes) -> SimpleNamespace:
        return SimpleNamespace(txts=("第一行", "Second line"))

    monkeypatch.setattr(ocr, "_rapidocr_engine", lambda: engine)

    assert ocr._extract_local(image) == "第一行\nSecond line"


def test_remote_ocr_sends_image_block(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    image = tmp_path / "scan.png"
    image.write_bytes(b"image")
    messages: list[object] = []

    class Model:
        @staticmethod
        def invoke(value: list[object]) -> SimpleNamespace:
            messages.extend(value)
            return SimpleNamespace(content="remote text")

    monkeypatch.setattr(ocr, "build_probe_chat_model", lambda *_a, **_k: Model())
    extractor = ocr._RemoteOcr(SimpleNamespace(), "vision")

    assert extractor(image) == "remote text"
    content = messages[0].content
    assert content[1]["type"] == "image_url"
    assert content[1]["image_url"]["url"].startswith("data:image/png;base64,")


def _remote_extractor(monkeypatch: pytest.MonkeyPatch, reply: object) -> ocr._RemoteOcr:
    """Build a ``_RemoteOcr`` whose model returns *reply* (str or a per-call sequence)."""

    class Model:
        @staticmethod
        def invoke(_value: list[object]) -> SimpleNamespace:
            if isinstance(reply, list):
                return SimpleNamespace(content=reply.pop(0))
            return SimpleNamespace(content=reply)

    monkeypatch.setattr(ocr, "build_probe_chat_model", lambda *_a, **_k: Model())
    return ocr._RemoteOcr(SimpleNamespace(), "vision")


@pytest.mark.parametrize(
    "refusal",
    [
        "No image was attached. Please upload an image.",
        "I don't see an image attached to your message. Please upload the image you'd "
        "like me to transcribe, and I'll provide the exact transcription.",
        "未收到图片，请上传图片后重试。",
    ],
)
def test_remote_ocr_ignores_no_image_refusal(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, refusal: str
) -> None:
    """A refusal is not source text: indexing it would mark the document ready with junk."""
    image = tmp_path / "scan.png"
    image.write_bytes(b"image")

    assert _remote_extractor(monkeypatch, refusal)(image) == ""


def test_remote_ocr_keeps_long_page_that_mentions_uploading_an_image(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """The refusal filter must not drop real transcriptions that happen to mention an image."""
    image = tmp_path / "scan.png"
    image.write_bytes(b"image")
    page = (
        "No image was attached. Please upload an image. "
        + "发票明细：办公用品 128.00 元，差旅费 340.00 元。" * 12
    )

    assert _remote_extractor(monkeypatch, page)(image) == page


def test_remote_ocr_skips_refusal_page_but_keeps_transcribed_page(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Refusals are dropped per page, so a partly readable document keeps its good pages."""
    pdf = tmp_path / "scan.pdf"
    pdf.write_bytes(b"%PDF-1.4")
    monkeypatch.setattr(
        ocr,
        "_image_inputs",
        lambda _path: iter([(b"page-1", "image/png"), (b"page-2", "image/png")]),
    )

    extractor = _remote_extractor(
        monkeypatch,
        ["No image was attached. Please upload an image.", "第二页正文"],
    )

    assert extractor(pdf) == "第二页正文"