homelab-codex-ws/jobs/documents-ingest/tests/test_summarize.py

"""Unit tests for the summarize+tag job — no DB, no real HTTP, no real Ollama/Anthropic."""
from __future__ import annotations

import json

import pytest

from documents_ingest.summarize import (
    DEFAULT_TAGS_VOCAB_PATH,
    OLLAMA_MAX_NUM_CTX,
    OLLAMA_MIN_NUM_CTX,
    build_partial_prompt,
    build_user_prompt,
    compute_num_ctx,
    get_summary_and_tags,
    insert_summary,
    is_duplicate,
    load_tags_vocab,
    normalize_and_validate_tags,
    normalize_tag,
    run_embed_summaries,
    run_summarize,
    summarize_mapreduce,
)

VOCAB = ["ubezpieczenie", "bank", "faktura"]


def _envelope_row(envelope_id, entities=None):
    return {"id": envelope_id, "entities": json.dumps(entities or [])}


class TestComputeNumCtx:
    def test_short_prompt_at_least_floor(self):
        assert compute_num_ctx(100) >= OLLAMA_MIN_NUM_CTX

    def test_zero_length_prompt_uses_floor(self):
        assert compute_num_ctx(0) == OLLAMA_MIN_NUM_CTX

    def test_scales_with_prompt_length(self):
        # Regression case from the pilot: a 93k-char document was silently truncated to
        # prompt_eval_count=2051 tokens by Ollama's runtime default -- num_ctx must scale
        # well past that for a prompt this size.
        ctx = compute_num_ctx(93_471)
        assert ctx > 30_000

    def test_huge_prompt_clamped_to_model_max(self):
        assert compute_num_ctx(10_000_000) == OLLAMA_MAX_NUM_CTX

    def test_rounds_up_to_1024_step(self):
        assert compute_num_ctx(1000) % 1024 == 0


class TestTagsVocab:
    def test_default_vocab_file_loads_and_is_nonempty(self):
        vocab = load_tags_vocab(DEFAULT_TAGS_VOCAB_PATH)
        assert "ubezpieczenie" in vocab
        assert len(vocab) >= 5


class TestIsDuplicate:
    def test_no_entities_not_duplicate(self):
        assert is_duplicate([]) is False
        assert is_duplicate(None) is False

    def test_duplicate_of_entity_detected(self):
        entities = [{"type": "content", "text": "x"}, {"type": "duplicate_of", "envelope_id": "paperless:14"}]
        assert is_duplicate(entities) is True

    def test_unrelated_entities_not_duplicate(self):
        entities = [{"type": "content", "text": "x"}, {"type": "tag", "name": "faktura"}]
        assert is_duplicate(entities) is False


class TestNormalizeTag:
    def test_lowercases_and_kebab_cases(self):
        assert normalize_tag("  Ubezpieczenie Auto ") == "ubezpieczenie-auto"

    def test_strips_punctuation(self):
        assert normalize_tag("PZU!!") == "pzu"

    def test_collapses_repeated_dashes(self):
        assert normalize_tag("a   b") == "a-b"

    def test_nfc_normalizes_diacritics(self):
        # decomposed "s" + combining cedilla vs precomposed - should normalize the same way
        import unicodedata
        decomposed = unicodedata.normalize("NFD", "ś") + "rodowisko"  # "ś" decomposed
        assert normalize_tag(decomposed) == normalize_tag("środowisko")


class TestNormalizeAndValidateTags:
    def test_keeps_all_vocab_tags(self):
        tags, truncated = normalize_and_validate_tags(["bank", "faktura"], VOCAB)
        assert tags == ["bank", "faktura"]
        assert truncated == 0

    def test_caps_freeform_at_three(self):
        raw = ["bank", "polisa", "pzu", "auto-osobowe", "extra-one"]
        tags, truncated = normalize_and_validate_tags(raw, VOCAB)
        assert tags[0] == "bank"
        assert len(tags) == 1 + 3  # 1 vocab + 3 freeform kept
        assert truncated == 1  # "extra-one" dropped

    def test_dedupes_after_normalization(self):
        tags, truncated = normalize_and_validate_tags(["Bank", "bank", " bank "], VOCAB)
        assert tags == ["bank"]
        assert truncated == 0

    def test_ignores_non_string_entries(self):
        tags, truncated = normalize_and_validate_tags(["bank", 123, None], VOCAB)
        assert tags == ["bank"]


class TestPromptBuilders:
    def test_user_prompt_includes_vocab_and_content(self):
        prompt = build_user_prompt(VOCAB, "treść dokumentu")
        assert "ubezpieczenie" in prompt
        assert "treść dokumentu" in prompt
        assert "JSON" in prompt

    def test_partial_prompt_includes_index_and_total(self):
        prompt = build_partial_prompt("fragment", 2, 5)
        assert "2/5" in prompt
        assert "fragment" in prompt


class _FakeChatResponse:
    def __init__(self, payload, status=200):
        self._payload = payload
        self._status = status

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def json(self):
        return self._payload

    def raise_for_status(self):
        if self._status >= 400:
            raise RuntimeError(f"HTTP {self._status}")


class _FakeOllamaChatSession:
    """Serves a queue of /api/chat responses, one per call (in order)."""

    def __init__(self, contents: list[str]):
        self._contents = list(contents)
        self.requests: list[dict] = []

    def post(self, url, json):
        self.requests.append({"url": url, "json": json})
        content = self._contents.pop(0)
        return _FakeChatResponse({"message": {"content": content}})

    async def close(self):
        pass


class _FakeBackend:
    """A scripted LLMBackend: returns queued responses in order, regardless of prompt."""

    def __init__(self, responses: list[str]):
        self._responses = list(responses)
        self.calls: list[tuple] = []

    async def complete(self, system, user, json_mode=True):
        self.calls.append((system, user, json_mode))
        return self._responses.pop(0)


class _FailingBackend:
    async def complete(self, system, user, json_mode=True):
        raise RuntimeError("boom")


class TestGetSummaryAndTags:
    async def test_valid_json_first_try(self):
        backend = _FakeBackend([json.dumps({"summary": "s", "tags": ["bank"]})])
        data = await get_summary_and_tags(backend, "sys", "user")
        assert data == {"summary": "s", "tags": ["bank"]}
        assert len(backend.calls) == 1

    async def test_invalid_json_then_valid_on_retry(self):
        backend = _FakeBackend(["not json", json.dumps({"summary": "s", "tags": []})])
        data = await get_summary_and_tags(backend, "sys", "user")
        assert data == {"summary": "s", "tags": []}
        assert len(backend.calls) == 2

    async def test_missing_keys_counts_as_invalid(self):
        backend = _FakeBackend([json.dumps({"summary": "s"}), json.dumps({"summary": "s", "tags": []})])
        data = await get_summary_and_tags(backend, "sys", "user")
        assert data == {"summary": "s", "tags": []}

    async def test_gives_up_after_retries_exhausted(self):
        backend = _FakeBackend(["not json", "still not json"])
        data = await get_summary_and_tags(backend, "sys", "user", retries=1)
        assert data is None

    async def test_backend_exception_is_treated_as_invalid(self):
        data = await get_summary_and_tags(_FailingBackend(), "sys", "user", retries=0)
        assert data is None


class TestSummarizeMapreduce:
    async def test_combines_partials_in_order(self):
        backend = _FakeBackend(["partial one", "partial two"])
        result = await summarize_mapreduce(backend, ["c1", "c2", "c3"], group_size=2)
        assert result == "partial one\n\npartial two"
        # two groups of size 2 and 1 -> two calls, both plain-text (json_mode=False)
        assert len(backend.calls) == 2
        assert all(call[2] is False for call in backend.calls)

    async def test_partial_failure_aborts_whole_document(self):
        result = await summarize_mapreduce(_FailingBackend(), ["c1", "c2"], group_size=2)
        assert result is None

    async def test_empty_partial_aborts(self):
        backend = _FakeBackend(["   "])
        result = await summarize_mapreduce(backend, ["c1"], group_size=2)
        assert result is None


class TestInsertSql:
    def test_on_conflict_target_includes_model(self):
        from documents_ingest.summarize import _INSERT_SQL
        assert "ON CONFLICT (envelope_id, model)" in _INSERT_SQL


class _FakeConn:
    """Fakes the three query shapes this job issues: envelope fetch, active-chunk-text
    fetch, and existing-summary-keys fetch; plus execute() for INSERT/UPDATE."""

    def __init__(self, docs=None, chunks_by_envelope=None, existing_summaries=None,
                 existing_model="bge-m3", execute_results=None):
        self._docs = docs or []
        self._chunks_by_envelope = chunks_by_envelope or {}
        self._existing_summaries = list(existing_summaries or [])
        self._existing_model = existing_model
        self._execute_results = list(execute_results) if execute_results is not None else None
        self.execute_calls: list[tuple] = []

    async def fetch(self, query, *params):
        if "FROM document_chunk" in query:
            envelope_id = params[0]
            return [{"text": t} for t in self._chunks_by_envelope.get(envelope_id, [])]
        if "FROM document_summary" in query:
            if params and params[0] != self._existing_model:
                return []
            return [{"envelope_id": eid} for eid in self._existing_summaries]
        return self._docs

    async def execute(self, query, *params):
        self.execute_calls.append(params)
        if self._execute_results is not None:
            return self._execute_results.pop(0)
        return "INSERT 0 1"

    async def close(self):
        pass


class TestRunSummarize:
    def _patch(self, monkeypatch, conn, session=None):
        async def _fake_connect(dsn):
            return conn
        monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect)
        if session is not None:
            def _fake_session_factory(*args, **kwargs):
                return session
            monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", _fake_session_factory)

    async def test_dry_run_counts_without_calling_backend_or_db(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["treść dokumentu"]},
        )
        self._patch(monkeypatch, conn)
        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False)
        assert stats["documents_fetched"] == 1
        assert stats["summarized"] == 1
        assert stats["llm_errors"] == 0
        assert conn.execute_calls == []

    async def test_duplicate_envelope_skipped_whole(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:74", entities=[{"type": "duplicate_of", "envelope_id": "paperless:14"}])],
        )
        self._patch(monkeypatch, conn)
        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False)
        assert stats["duplicates_skipped"] == 1
        assert stats["summarized"] == 0

    async def test_no_active_chunks_counted_separately(self, monkeypatch):
        conn = _FakeConn(docs=[_envelope_row("paperless:2")], chunks_by_envelope={})
        self._patch(monkeypatch, conn)
        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False)
        assert stats["no_active_chunks"] == 1
        assert stats["summarized"] == 0

    async def test_already_summarized_skipped(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["text"]},
            existing_summaries=["paperless:1"],
            existing_model="gemma3:12b",
        )
        self._patch(monkeypatch, conn)
        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False)
        assert stats["already_summarized"] == 1
        assert stats["summarized"] == 0

    async def test_apply_writes_summary_and_tags(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["treść z kwotą 100 zł"]},
        )
        session = _FakeOllamaChatSession([json.dumps({"summary": "Streszczenie.", "tags": ["bank", "extra"]})])
        self._patch(monkeypatch, conn, session)

        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True)

        assert stats["summarized"] == 1
        assert stats["llm_errors"] == 0
        assert len(conn.execute_calls) == 1
        envelope_id, summary, tags_json, model = conn.execute_calls[0]
        assert envelope_id == "paperless:1"
        assert summary == "Streszczenie."
        assert json.loads(tags_json) == ["bank", "extra"]
        assert model == "gemma3:12b"
        # Regression guard: every Ollama call must carry an explicit num_ctx (plan §5.1
        # limitation) -- Ollama's runtime default silently truncates long documents otherwise.
        assert session.requests[0]["json"]["options"]["num_ctx"] >= OLLAMA_MIN_NUM_CTX

    async def test_apply_with_anthropic_backend(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["treść"]},
        )

        class _TextBlock:
            type = "text"
            text = json.dumps({"summary": "Streszczenie API.", "tags": ["faktura"]})

        class _FakeResponse:
            content = [_TextBlock()]

        class _FakeMessages:
            def __init__(self):
                self.create_calls = []

            async def create(self, **kwargs):
                self.create_calls.append(kwargs)
                return _FakeResponse()

        class _FakeAsyncAnthropic:
            def __init__(self, api_key=None):
                self.api_key = api_key
                self.messages = _FakeMessages()

            async def close(self):
                pass

        fake_client_holder = {}

        def _factory(api_key=None):
            client = _FakeAsyncAnthropic(api_key=api_key)
            fake_client_holder["client"] = client
            return client

        async def _fake_connect(dsn):
            return conn
        monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect)
        monkeypatch.setattr("documents_ingest.summarize.AsyncAnthropic", _factory)

        stats = await run_summarize(
            dsn="x", backend_name="anthropic", model="claude-haiku-4-5",
            anthropic_api_key="sk-fake", apply=True,
        )

        assert stats["summarized"] == 1
        client = fake_client_holder["client"]
        assert client.api_key == "sk-fake"
        assert len(client.messages.create_calls) == 1
        call = client.messages.create_calls[0]
        assert call["model"] == "claude-haiku-4-5"
        assert call["temperature"] == 0
        assert call["output_config"]["format"]["type"] == "json_schema"

    async def test_llm_error_isolated_and_counted(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1"), _envelope_row("paperless:2")],
            chunks_by_envelope={"paperless:1": ["a"], "paperless:2": ["b"]},
        )
        session = _FakeOllamaChatSession([
            "not json", "still not json",  # paperless:1 exhausts retry -> llm_error
            json.dumps({"summary": "ok", "tags": []}),  # paperless:2 succeeds
        ])
        self._patch(monkeypatch, conn, session)

        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True)

        assert stats["documents_fetched"] == 2
        assert stats["llm_errors"] == 1
        assert stats["summarized"] == 1
        assert len(conn.execute_calls) == 1  # only the successful one got inserted

    async def test_conflict_skip_counted_as_already_summarized(self, monkeypatch):
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["a"]},
            execute_results=["INSERT 0 0"],  # ON CONFLICT DO NOTHING no-op
        )
        session = _FakeOllamaChatSession([json.dumps({"summary": "s", "tags": []})])
        self._patch(monkeypatch, conn, session)

        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True)

        assert stats["already_summarized"] == 1
        assert stats["summarized"] == 0

    async def test_stats_balance_invariant(self, monkeypatch):
        conn = _FakeConn(
            docs=[
                _envelope_row("paperless:1", entities=[{"type": "duplicate_of", "envelope_id": "paperless:0"}]),
                _envelope_row("paperless:2"),
                _envelope_row("paperless:3"),
            ],
            chunks_by_envelope={"paperless:3": ["content"]},
        )
        self._patch(monkeypatch, conn)
        stats = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=False)
        balance = (
            stats["duplicates_skipped"] + stats["no_active_chunks"]
            + stats["already_summarized"] + stats["summarized"] + stats["llm_errors"]
        )
        assert balance == stats["documents_fetched"] == 3

    async def test_limit_and_offset_passed_through(self, monkeypatch):
        class _Conn(_FakeConn):
            async def fetch(self, query, *params):
                if "FROM envelope" in query:
                    self.captured = (query, params)
                    return []
                return await super().fetch(query, *params)

        conn = _Conn()
        self._patch(monkeypatch, conn)
        await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", limit=5, offset=10, apply=False)
        query, params = conn.captured
        assert "LIMIT" in query and "OFFSET" in query
        assert params == (5, 10)

    async def test_rerun_after_apply_writes_nothing_new(self, monkeypatch):
        conn1 = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["a"]},
        )
        session1 = _FakeOllamaChatSession([json.dumps({"summary": "s", "tags": []})])
        self._patch(monkeypatch, conn1, session1)
        await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True)

        conn2 = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": ["a"]},
            existing_summaries=["paperless:1"],
            existing_model="gemma3:12b",
        )
        session2 = _FakeOllamaChatSession([])
        self._patch(monkeypatch, conn2, session2)
        stats2 = await run_summarize(dsn="x", backend_name="ollama", model="gemma3:12b", apply=True)

        assert stats2["already_summarized"] == 1
        assert stats2["summarized"] == 0


class TestMapreduceIntegration:
    def _patch(self, monkeypatch, conn, session=None):
        async def _fake_connect(dsn):
            return conn
        monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect)
        if session is not None:
            monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", lambda *a, **k: session)

    async def test_long_document_triggers_mapreduce(self, monkeypatch):
        long_chunks = ["x" * 100_000, "y" * 150_000]  # combined > threshold
        conn = _FakeConn(
            docs=[_envelope_row("paperless:1")],
            chunks_by_envelope={"paperless:1": long_chunks},
        )
        session = _FakeOllamaChatSession([
            "partial summary 1",  # map-reduce partial (plain text, one group covers both chunks)
            json.dumps({"summary": "final", "tags": []}),  # final synthesis (JSON)
        ])
        self._patch(monkeypatch, conn, session)

        stats = await run_summarize(
            dsn="x", backend_name="ollama", model="gemma3:12b", apply=True,
            mapreduce_threshold_chars=200_000, chunks_per_group=20,
        )

        assert stats["documents_mapreduce"] == 1
        assert stats["summarized"] == 1
        # first call has no "format" key (plain text), second has format=json
        assert "format" not in session.requests[0]["json"]
        assert session.requests[1]["json"]["format"] == "json"


class TestRunEmbedSummaries:
    class _EmbedConn:
        def __init__(self, rows):
            self._rows = rows
            self.update_calls: list[tuple] = []

        async def fetch(self, query, *params):
            return self._rows

        async def execute(self, query, *params):
            self.update_calls.append(params)
            return "UPDATE 1"

        async def close(self):
            pass

    class _EmbedSession:
        def __init__(self, dim=1024):
            self._dim = dim

        def post(self, url, json):
            return _FakeEmbedResponseForEmbed({"embedding": [0.1] * self._dim})

        async def close(self):
            pass

    def _patch(self, monkeypatch, conn, session=None):
        async def _fake_connect(dsn):
            return conn
        monkeypatch.setattr("documents_ingest.summarize.asyncpg.connect", _fake_connect)
        if session is not None:
            monkeypatch.setattr("documents_ingest.summarize.aiohttp.ClientSession", lambda *a, **k: session)

    async def test_dry_run_counts_without_writing(self, monkeypatch):
        conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}])
        self._patch(monkeypatch, conn)
        stats = await run_embed_summaries(dsn="x", apply=False)
        assert stats["summaries_fetched"] == 1
        assert stats["embedded"] == 1
        assert conn.update_calls == []

    async def test_apply_embeds_and_updates(self, monkeypatch):
        conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}])
        session = self._EmbedSession()
        self._patch(monkeypatch, conn, session)
        stats = await run_embed_summaries(dsn="x", embed_model="bge-m3", apply=True)
        assert stats["embedded"] == 1
        assert stats["errors"] == 0
        assert len(conn.update_calls) == 1
        vector_literal, embed_model, summary_id = conn.update_calls[0]
        assert embed_model == "bge-m3"
        assert summary_id == 1

    async def test_balance_invariant(self, monkeypatch):
        conn = self._EmbedConn(rows=[{"id": 1, "summary": "s"}, {"id": 2, "summary": "t"}])
        self._patch(monkeypatch, conn)
        stats = await run_embed_summaries(dsn="x", apply=False)
        assert stats["embedded"] + stats["errors"] == stats["summaries_fetched"]


class _FakeEmbedResponseForEmbed:
    def __init__(self, payload):
        self._payload = payload

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    async def json(self):
        return self._payload

    def raise_for_status(self):
        pass