stack/tests/llm/test_index.py

"""llm.index — incremental, resumable embedding indexer."""

from unittest.mock import MagicMock, patch

from sqlalchemy.exc import ProgrammingError

from llm.chunk import Doc, content_hash
from llm.config import LlmConfig
from llm.index import index_docs

CFG = LlmConfig(
    ollama_hosts=("http://h1:11434",),
    embed_model="m",
    instruct_model="g",
    embed_dim=768,
    build_ann_index=False,
    pg_host="x",
    pg_port=5432,
    pg_db="llm",
    pg_user="llm",
)

DOC = Doc(key="K1", text="Some body text.", metadata={"docket": "D"})


def _run(docs, state_rows, force=False, cfg=CFG):
    """Run index_docs with everything external mocked; return mocks."""
    store = MagicMock()
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    conn.execute.return_value.fetchall.return_value = state_rows
    with (
        patch("llm.index._engine", return_value=engine),
        patch("llm.index.vectorstore", return_value=store),
        patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
        patch("llm.index.ensure_hnsw") as mock_ensure_hnsw,
        patch("llm.index.HostPool") as MockPool,
    ):
        MockPool.return_value.check.return_value = ["http://h1:11434"]
        stats = index_docs(
            docs,
            collection="comments",
            cfg=cfg,
            pool=MockPool.return_value,
            force=force,
        )
    return stats, store, conn, mock_ensure_hnsw


class TestIndexDocs:
    def test_new_doc_embedded_and_recorded(self):
        stats, store, conn, _ensure_hnsw = _run([DOC], state_rows=[])
        assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
        store.add_embeddings.assert_called_once()
        kwargs = store.add_embeddings.call_args.kwargs
        assert kwargs["ids"][0].startswith("K1:")

    def test_unchanged_doc_skipped(self):
        h = content_hash(DOC.text)
        stats, store, _, _ensure_hnsw = _run([DOC], state_rows=[("K1", h, "")])
        assert stats["skipped"] == 1
        store.add_embeddings.assert_not_called()

    def test_changed_doc_deletes_old_chunks_first(self):
        stats, store, conn, _ensure_hnsw = _run(
            [DOC], state_rows=[("K1", "stalehash", "")]
        )
        assert stats["indexed"] == 1
        deletes = [
            c
            for c in conn.execute.call_args_list
            if "DELETE FROM langchain_pg_embedding" in str(c.args[0])
        ]
        assert len(deletes) == 1
        sql = str(deletes[0].args[0])
        assert "item_key" in sql  # DELETE ... cmetadata->>'item_key'
        # Scoped to the target collection, not a global delete by item_key.
        assert "collection_id" in sql
        assert "langchain_pg_collection" in sql
        params = deletes[0].args[1]
        assert params["k"] == "K1"
        assert params["c"] == "comments"

    def test_force_reembeds_unchanged(self):
        h = content_hash(DOC.text)
        stats, store, _, _ensure_hnsw = _run(
            [DOC], state_rows=[("K1", h, "")], force=True
        )
        assert stats["indexed"] == 1

    def test_empty_doc_counts_skipped(self):
        empty = Doc(key="K2", text="   ", metadata={})
        stats, store, _, _ensure_hnsw = _run([empty], state_rows=[])
        assert stats == {"indexed": 0, "skipped": 1, "chunks": 0}


class TestPageEnrichmentHook:
    def test_enriches_only_docs_that_embed(self, monkeypatch):
        """index_docs runs llm.pages.enrich_pdf_pages only on docs that reach
        the embed path — a hash match must not open the doc's PDFs."""
        from llm import index as index_mod

        seen = []
        monkeypatch.setattr(
            index_mod,
            "enrich_pdf_pages",
            lambda doc, chunks: seen.append(doc.key) or chunks,
        )
        _run([DOC], [])
        assert seen == ["K1"]
        seen.clear()
        _run([DOC], [("K1", content_hash(DOC.text), "")])
        assert seen == []


class TestDeleteOldChunksFallback:
    def test_programming_error_swallowed_and_indexing_continues(self):
        """A fresh DB where langchain_pg_embedding doesn't exist yet raises
        ProgrammingError on DELETE; _delete_old_chunks must swallow it and
        indexing must still proceed to add_embeddings + record state."""
        store = MagicMock()
        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value

        def fake_execute(clause, *args, **kwargs):
            sql = str(clause)
            if "DELETE FROM langchain_pg_embedding" in sql:
                raise ProgrammingError("stmt", {}, Exception("relation missing"))
            result = MagicMock()
            result.fetchall.return_value = []  # no prior state
            return result

        conn.execute.side_effect = fake_execute

        with (
            patch("llm.index._engine", return_value=engine),
            patch("llm.index.vectorstore", return_value=store),
            patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
            patch("llm.index.ensure_hnsw"),
            patch("llm.index.HostPool") as MockPool,
        ):
            MockPool.return_value.check.return_value = ["http://h1:11434"]
            stats = index_docs(
                [DOC],
                collection="comments",
                cfg=CFG,
                pool=MockPool.return_value,
                force=False,
            )

        assert stats == {"indexed": 1, "skipped": 0, "chunks": 1}
        store.add_embeddings.assert_called_once()


class TestAnnIndexGating:
    def test_build_ann_index_false_skips_ensure_hnsw(self):
        _stats, _store, _conn, mock_ensure_hnsw = _run([DOC], state_rows=[], cfg=CFG)
        mock_ensure_hnsw.assert_not_called()

    def test_build_ann_index_true_calls_ensure_hnsw(self):
        cfg = LlmConfig(
            ollama_hosts=("http://h1:11434",),
            embed_model="m",
            instruct_model="g",
            embed_dim=768,
            build_ann_index=True,
            pg_host="x",
            pg_port=5432,
            pg_db="llm",
            pg_user="llm",
        )
        _stats, _store, _conn, mock_ensure_hnsw = _run([DOC], state_rows=[], cfg=cfg)
        assert mock_ensure_hnsw.called