stack/tests/llm/test_api.py

"""llm.api — FastAPI chat app."""

import re
from unittest.mock import patch

from fastapi.testclient import TestClient

from llm.api import _COLLECTION_CHOICES, _KIND_CHOICES, app

client = TestClient(app)


class TestHealth:
    def test_ok(self):
        r = client.get("/health")
        assert r.status_code == 200
        assert r.json() == {"status": "ok"}


class TestIndex:
    def test_serves_chat_page(self):
        r = client.get("/")
        assert r.status_code == 200
        assert "text/html" in r.headers["content-type"]
        assert "Library Chat" in r.text
        assert 'id="form"' in r.text

    def test_serves_chat_page_with_valuation_renderer(self):
        r = client.get("/")
        assert r.status_code == 200
        html = r.text
        assert "function renderValuation" in html
        assert "ev.type === 'valuation'" in html
        assert "table.valuation" in html
        # the money on screen is national and unadjusted — say so
        assert "no geographic (GPCI) adjustment" in html
        assert "createCaption" in html
        # notes explain an unpaid status and which CF is shown
        assert "r.status_note" in html and "r.cf_note" in html
        assert "'cf', 'CF (non-APM)'" in html
        # a citation without a URL is not an empty <a href="">
        assert "if (p.url)" in html

    def test_serves_chat_page_with_lineage_renderer(self):
        r = client.get("/")
        assert r.status_code == 200
        html = r.text
        assert "function renderLineage" in html
        assert "ev.type === 'lineage'" in html
        assert "table.lineage" in html
        assert 'id="mode"' in html
        assert "tr.unanchored" in html
        assert "ul.className = 'elements'" in html
        assert "ul.className = 'guidance'" in html
        # mode is sent to the API and echoed back on the sources meta line
        assert "mode" in html and "ev.mode === 'timeline'" in html

    def test_serves_chat_page_with_markdown_and_cite_links(self):
        """The answer is painted as markdown (headings/lists/bold, never
        raw "## …"), and a cited [label] whose URL arrived on the
        lineage/valuation/sources events is an <a class="cite"> jump
        link to the FR paragraph / eCFR section; the answer is repainted
        when the sources event completes the label→url map."""
        html = client.get("/").text
        assert "function paintInline" in html
        assert "function paint(el, text, links)" in html
        assert "function collectLinks" in html
        assert "a.cite" in html
        assert (
            "collectLinks(links, ev); renderSources(wrap, ev); paint(b, answer, links)"
            in html
        )
        assert "document.createElement('h' + level)" in html
        # model text never reaches innerHTML
        assert "innerHTML = ''" in html and "innerHTML = text" not in html


class TestSearchPage:
    """``GET /ui/search`` — the framework-free search UI (refs #578)."""

    def test_serves_search_page(self):
        r = client.get("/ui/search")
        assert r.status_code == 200
        assert "text/html" in r.headers["content-type"]
        html = r.text
        assert "Library Search" in html
        # the form: query box + every filter the API accepts
        assert 'id="search"' in html
        assert 'id="q"' in html
        for field in ("collection", "kind", "docket", "year"):
            assert 'id="%s"' % field in html
        for value in sorted(_COLLECTION_CHOICES):
            assert 'value="%s"' % value in html
        for value in sorted(_KIND_CHOICES):
            assert 'value="%s"' % value in html
        # llm: tags aren't indexed yet — a placeholder, not a dead filter
        assert "#575" in html and 'id="tag"' not in html

    def test_trailing_slash_also_serves_the_page(self):
        """The headless probe's readiness check appends "/"; a 307 there
        would read as "service not responding"."""
        r = client.get("/ui/search/")
        assert r.status_code == 200
        assert "Library Search" in r.text

    def test_renders_results_similar_and_pagination(self):
        html = client.get("/ui/search").text
        # results: label link, snippet, docket/comment id, similar
        assert "function hitEl(r)" in html
        assert "r.snippet" in html and "r.docket" in html and "r.comment_id" in html
        assert "if (r.url)" in html  # no empty href for an unlinked label
        # /similar is keyed by the bib item key the result carries
        assert "function showSimilar" in html
        assert "r.item_key" in html
        assert "'/similar/' + encodeURIComponent(key)" in html
        # pagination + shareable URL
        assert 'id="prev"' in html and 'id="next"' in html
        assert "offset += LIMIT" in html
        assert "window.history.pushState" in html
        assert "new URLSearchParams(window.location.search)" in html
        # loading + error states, with the API's own 400 detail
        assert "'Searching…'" in html
        assert "body.detail" in html
        assert "setStatus('Error: ' + e.message, true)" in html

    def test_no_external_script_or_style_urls(self):
        """Framework-free and CDN-free: behind oauth2-proxy the page has
        no outbound access, so every byte of CSS/JS is inline."""
        html = client.get("/ui/search").text
        assert re.search(r"<script[^>]*\ssrc=", html) is None
        assert re.search(r"<link[^>]*stylesheet", html) is None
        assert "fonts.googleapis.com" not in html
        # no absolute off-site fetch at all (the favicon is the platform's
        # own, protocol-relative, host)
        assert "http://" not in html and "https://" not in html
        # API data never reaches innerHTML — nothing is ever assigned to it
        assert re.search(r"innerHTML\s*=", html) is None

    def test_nav_links_between_the_two_pages_resolve(self):
        search_html = client.get("/ui/search").text
        chat_html = client.get("/").text
        assert '<a class="nav" href="/">Chat</a>' in search_html
        assert '<a class="nav" href="/ui/search">Search</a>' in chat_html
        assert client.get("/").status_code == 200
        assert client.get("/ui/search").status_code == 200


class TestWhoami:
    def test_reads_forwarded_header(self):
        r = client.get("/whoami", headers={"X-Auth-Request-User": "kert"})
        assert r.json() == {"user": "kert"}

    def test_empty_when_absent(self):
        r = client.get("/whoami")
        assert r.json() == {"user": ""}


class TestChat:
    def test_empty_question_400(self):
        r = client.post("/chat", json={"question": "   "})
        assert r.status_code == 400

    @patch("llm.rag.stream_answer")
    def test_streams_sse_events(self, mock_stream):
        mock_stream.return_value = iter(
            [
                {"type": "token", "text": "Hi"},
                {"type": "sources", "sources": []},
                {"type": "done"},
            ]
        )
        r = client.post("/chat", json={"question": "hello"})
        assert r.status_code == 200
        assert "text/event-stream" in r.headers["content-type"]
        body = r.text
        assert 'data: {"type": "token", "text": "Hi"}' in body
        assert '"type": "done"' in body

    @patch("llm.rag.stream_answer")
    def test_error_becomes_event_not_500(self, mock_stream):
        mock_stream.side_effect = RuntimeError("ollama down")
        r = client.post("/chat", json={"question": "hello"})
        assert r.status_code == 200
        assert '"type": "error"' in r.text
        assert "ollama down" in r.text


class TestChatSince:
    @patch("llm.rag.stream_answer")
    def test_since_forwarded(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q", "since": "2025-09-01"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["since"] == "2025-09-01"

    def test_bad_since_400(self):
        r = client.post("/chat", json={"question": "q", "since": "last year"})
        assert r.status_code == 400


class TestChatMode:
    @patch("llm.rag.stream_answer")
    def test_mode_forwarded(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q", "mode": "timeline"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["mode"] == "timeline"

    @patch("llm.rag.stream_answer")
    def test_mode_defaults_to_auto(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["mode"] == "auto"

    def test_bad_mode_400(self):
        r = client.post("/chat", json={"question": "q", "mode": "x"})
        assert r.status_code == 400


def _cfg(**kw):
    from llm.config import LlmConfig

    base = dict(
        ollama_hosts=("http://h1:11434", "http://h2:11434"),
        host_vram={"http://h1:11434": 12, "http://h2:11434": 24},
        embed_model="e",
        instruct_model="chat",
        embed_dim=768,
        build_ann_index=False,
        pg_host="x",
        pg_port=5432,
        pg_db="llm",
        pg_user="llm",
    )
    base.update(kw)
    return LlmConfig(**base)


class TestStartup:
    """#699 ruling B3: the replica is warmed at boot, not on the chat's
    first valuation lookup. ``with TestClient(app) as c:`` is required to
    actually trigger FastAPI's startup event — a bare ``TestClient(app)``
    (as ``client`` above, module-level) never sends the ASGI lifespan
    ``startup`` message."""

    @patch("llm.config.load")
    def test_startup_event_warms_the_replica(self, mock_load):
        mock_load.return_value = _cfg()
        with patch("llm.evidence.warm") as mock_warm:
            with TestClient(app):
                pass
        mock_warm.assert_called_once_with(mock_load.return_value)


class TestHosts:
    @patch("llm.pool.pick_model", return_value="big")
    @patch("llm.pool.HostPool.check", return_value=["http://h2:11434"])
    @patch("llm.pool.HostPool.status")
    @patch("llm.config.load")
    def test_reports_fleet_and_pick(self, mock_load, mock_status, _check, _pm):
        mock_load.return_value = _cfg()
        mock_status.return_value = [
            {"host": "http://h2:11434", "vram_gb": 24.0, "models": ["chat:latest"]}
        ]
        r = client.get("/hosts")
        assert r.status_code == 200
        body = r.json()
        assert body["generation"] == {"host": "http://h2:11434", "model": "big"}
        assert [(h["host"], h["live"]) for h in body["hosts"]] == [
            ("http://h1:11434", False),
            ("http://h2:11434", True),
        ]

    @patch("llm.pool.HostPool.check", side_effect=RuntimeError("no Ollama host"))
    @patch("llm.config.load")
    def test_no_live_hosts(self, mock_load, _check):
        mock_load.return_value = _cfg(ollama_hosts=("http://h1:11434",))
        r = client.get("/hosts")
        assert r.json()["generation"] is None
        assert "no Ollama host" in r.json()["error"]


class TestSearchEndpoint:
    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_empty_q_400(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        r = client.get("/search", params={"q": "   "})
        assert r.status_code == 400
        mock_search.assert_not_called()

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_unknown_collection_400(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        r = client.get("/search", params={"q": "ccm", "collection": "bogus"})
        assert r.status_code == 400
        mock_search.assert_not_called()

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_unknown_kind_400(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        r = client.get("/search", params={"q": "ccm", "kind": "bogus"})
        assert r.status_code == 400
        mock_search.assert_not_called()

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_limit_out_of_range_400(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        assert client.get("/search", params={"q": "ccm", "limit": 0}).status_code == 400
        assert (
            client.get("/search", params={"q": "ccm", "limit": 51}).status_code == 400
        )
        mock_search.assert_not_called()

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_negative_offset_400(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        r = client.get("/search", params={"q": "ccm", "offset": -1})
        assert r.status_code == 400
        mock_search.assert_not_called()

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_filtered_search_scopes_by_docket_and_returns_shape(
        self, mock_load, mock_search
    ):
        mock_load.return_value = _cfg()
        src = {
            "id": "CMS-2023-0121-1",
            "label": "CMS-2023-0121-1",
            "kind": "comment",
            "url": "u",
            "title": "t",
            "date": "2024-01-01",
            "docket": "CMS-2023-0121",
            "comment_id": "CMS-2023-0121-1",
            "snippet": "s",
            "score": 0.9,
            "item_key": "K1",
            "p_id": "",
            "seq": "1",
            "section": "",
            "distance": 0.1,
        }
        mock_search.return_value = [src]
        r = client.get(
            "/search",
            params={
                "q": "chronic care management",
                "docket": "CMS-2023-0121",
                "collection": "comments",
                "limit": 10,
            },
        )
        assert r.status_code == 200
        body = r.json()
        assert body["query"] == "chronic care management"
        assert body["filters"] == {"docket": "CMS-2023-0121"}
        assert body["total"] == 1
        assert body["results"] == [src]
        assert all(s["docket"] == "CMS-2023-0121" for s in body["results"])
        kwargs = mock_search.call_args.kwargs
        assert kwargs["collection"] == "comments"
        assert kwargs["filters"] == {"docket": "CMS-2023-0121"}
        assert kwargs["limit"] == 10
        assert kwargs["offset"] == 0

    @patch("llm.search.search")
    @patch("llm.config.load")
    def test_defaults(self, mock_load, mock_search):
        mock_load.return_value = _cfg()
        mock_search.return_value = []
        r = client.get("/search", params={"q": "q"})
        assert r.status_code == 200
        kwargs = mock_search.call_args.kwargs
        assert kwargs["collection"] == "all"
        assert kwargs["filters"] == {}
        assert kwargs["limit"] == 10
        assert kwargs["offset"] == 0


class TestSimilarEndpoint:
    @patch("llm.search.similar")
    @patch("llm.config.load")
    def test_unknown_collection_400(self, mock_load, mock_similar):
        mock_load.return_value = _cfg()
        r = client.get("/similar/K1", params={"collection": "bogus"})
        assert r.status_code == 400
        mock_similar.assert_not_called()

    @patch("llm.search.similar")
    @patch("llm.config.load")
    def test_limit_out_of_range_400(self, mock_load, mock_similar):
        mock_load.return_value = _cfg()
        assert client.get("/similar/K1", params={"limit": 0}).status_code == 400
        assert client.get("/similar/K1", params={"limit": 51}).status_code == 400
        mock_similar.assert_not_called()

    @patch("llm.search.similar", return_value=None)
    @patch("llm.config.load")
    def test_unknown_key_404(self, mock_load, mock_similar):
        mock_load.return_value = _cfg()
        r = client.get("/similar/NOPE")
        assert r.status_code == 404

    @patch("llm.search.similar")
    @patch("llm.config.load")
    def test_returns_neighbours_shape(self, mock_load, mock_similar):
        mock_load.return_value = _cfg()
        src = {
            "id": "91 FR 43949 ¶1",
            "label": "91 FR 43949 ¶1",
            "kind": "rule",
            "url": "u",
            "title": "t",
            "date": "2024-01-01",
            "docket": "",
            "comment_id": "",
            "snippet": "s",
            "score": 0.8,
            "item_key": "R1",
            "p_id": "1",
            "seq": "",
            "section": "",
            "distance": 0.2,
        }
        mock_similar.return_value = [src]
        r = client.get("/similar/K1", params={"limit": 5})
        assert r.status_code == 200
        body = r.json()
        assert body["key"] == "K1"
        assert body["total"] == 1
        assert body["results"] == [src]
        kwargs = mock_similar.call_args.kwargs
        assert kwargs["collection"] == "all"
        assert kwargs["limit"] == 5


class TestPdfViewer:
    """/ui/pdf/<key>, /pdf/<key>, /pdf/<key>/file and the vendored PDF.js."""

    def test_viewer_page_and_vendor_assets(self):
        html = client.get("/ui/pdf/ABCD1234").text
        assert "import('/ui/vendor/pdf.min.mjs')" in html
        assert "workerSrc = '/ui/vendor/pdf.worker.min.mjs'" in html
        assert "IntersectionObserver" in html and "ResizeObserver" in html
        assert "TextLayer" in html and "runFind" in html
        assert client.get("/ui/pdf/not a key").status_code == 404
        js = client.get("/ui/vendor/pdf.min.mjs")
        assert js.status_code == 200 and js.headers["content-type"].startswith(
            "text/javascript"
        )
        assert client.get("/ui/vendor/pdf.worker.min.mjs").status_code == 200
        assert client.get("/ui/vendor/jszip.min.js").status_code == 200
        assert client.get("/ui/vendor/docx-preview.min.js").status_code == 200
        assert (
            "renderAsync" in html
            and "docx-preview.min.js" in html
            and "renderer === 'text'" in html
        )
        assert client.get("/ui/vendor/../api.py").status_code in (404, 400)
        assert client.get("/ui/vendor/evil.mjs").status_code == 404

    def test_list_and_serve(self, tmp_path, monkeypatch):
        import llm.api as api
        import llm.convert as conv
        from llm.pdfs import PdfFile

        monkeypatch.setattr(
            conv, "available", lambda: False
        )  # the host may have soffice

        root = tmp_path / "root"
        (root / "K").mkdir(parents=True)
        pdf = root / "K" / "a.pdf"
        pdf.write_bytes(b"%PDF-1.4\n" + b"x" * 5000)
        docx = root / "K" / "b.docx"
        docx.write_bytes(b"PK")
        files = [
            PdfFile("a.pdf", pdf, pdf.stat().st_size, "comment"),
            PdfFile("b.docx", docx, 2, "comment"),
        ]
        monkeypatch.setattr(
            api, "_pdf_files", lambda key: files if key == "ABCD1234" else []
        )
        monkeypatch.setattr(
            api, "_pdf_paths", lambda: (tmp_path / "z.sqlite", tmp_path / "zs", root)
        )

        class _Store:
            _storage = tmp_path / "bibstorage"

            def get(self, key):
                class _I:
                    title = "Comment on CMS-2026-2377-1"

                return _I()

            def close(self):
                pass

        monkeypatch.setattr("conf.connect.bib", lambda: _Store())
        r = client.get("/pdf/ABCD1234")
        assert r.status_code == 200
        body = r.json()
        assert body["title"] == "Comment on CMS-2026-2377-1"
        assert [(f["name"], f["renderable"], f["renderer"]) for f in body["files"]] == [
            ("a.pdf", True, "pdf"),
            ("b.docx", True, "docx"),
        ]
        assert client.get("/pdf/NOPE9999").json()["files"] == []

        r = client.get("/pdf/ABCD1234/file")
        assert r.status_code == 200 and r.headers["content-type"] == "application/pdf"
        assert r.headers["content-disposition"].startswith("inline")
        assert r.headers["accept-ranges"] == "bytes" and r.content.startswith(b"%PDF")
        r = client.get("/pdf/ABCD1234/file", headers={"Range": "bytes=0-3"})
        assert r.status_code == 206 and r.content == b"%PDF"
        r = client.get("/pdf/ABCD1234/file?name=b.docx")
        assert r.status_code == 200 and r.headers["content-disposition"].startswith(
            "attachment"
        )
        assert client.get("/pdf/ABCD1234/file?name=nope.pdf").status_code == 404
        assert client.get("/pdf/ABCD1234/file?name=../a.pdf").status_code == 404

    def test_chat_and_search_pages_link_the_viewer(self):
        chat = client.get("/").text
        assert "'/ui/pdf/' + encodeURIComponent(src.item_key)" in chat
        assert "src.kind !== 'rule'" in chat
        search = client.get("/ui/search").text
        assert "'/ui/pdf/' + encodeURIComponent(r.item_key)" in search


class TestWordExactPages:
    """A Word attachment lists as renderer 'pdf' with an as=pdf url when
    LibreOffice is in the image, and /pdf/<key>/file?as=pdf serves the
    cached conversion; without the converter it falls back to docx-preview."""

    def _setup(self, tmp_path, monkeypatch, available):
        import llm.api as api
        import llm.convert as conv
        from llm.pdfs import PdfFile

        root = tmp_path / "root"
        (root / "K").mkdir(parents=True)
        docx = root / "K" / "letter.docx"
        docx.write_bytes(b"PK docx")
        files = [PdfFile("letter.docx", docx, 7, "comment")]
        monkeypatch.setattr(api, "_pdf_files", lambda key: files)
        monkeypatch.setattr(
            api, "_pdf_paths", lambda: (tmp_path / "z.sqlite", tmp_path / "zs", root)
        )
        monkeypatch.setattr(conv, "available", lambda: available)

        class _Store:
            _storage = tmp_path / "bibstorage"

            def get(self, key):
                class _I:
                    title = "t"

                return _I()

            def close(self):
                pass

        monkeypatch.setattr("conf.connect.bib", lambda: _Store())
        return docx

    def test_with_converter(self, tmp_path, monkeypatch):
        import llm.convert as conv

        docx = self._setup(tmp_path, monkeypatch, True)
        converted = tmp_path / "letter.pdf"
        converted.write_bytes(b"%PDF-1.7 word-exact")
        calls = []

        def fake_to_pdf(src, *, cache_dir):
            calls.append((src, cache_dir))
            return converted

        monkeypatch.setattr(conv, "to_pdf", fake_to_pdf)
        f = client.get("/pdf/ABCD1234").json()["files"][0]
        assert (
            f["renderer"] == "pdf"
            and f["converted"] is True
            and f["renderable"] is True
        )
        assert f["url"] == "/pdf/ABCD1234/file?name=letter.docx&as=pdf"
        assert f["download_url"] == "/pdf/ABCD1234/file?name=letter.docx"
        r = client.get(f["url"])
        assert r.status_code == 200 and r.headers["content-type"] == "application/pdf"
        assert r.content == b"%PDF-1.7 word-exact" and r.headers[
            "content-disposition"
        ].startswith("inline")
        assert (
            calls
            and calls[0][0] == docx
            and str(calls[0][1]).endswith(".state/llm/pdfcache")
        )
        # the original is still a download
        r = client.get(f["download_url"])
        assert r.status_code == 200 and r.headers["content-disposition"].startswith(
            "attachment"
        )

    def test_conversion_failure_is_a_502(self, tmp_path, monkeypatch):
        import llm.convert as conv

        self._setup(tmp_path, monkeypatch, True)

        def boom(src, *, cache_dir):
            raise conv.ConversionError("letter.docx: LibreOffice failed (1)")

        monkeypatch.setattr(conv, "to_pdf", boom)
        r = client.get("/pdf/ABCD1234/file?name=letter.docx&as=pdf")
        assert r.status_code == 502 and "LibreOffice failed" in r.json()["detail"]

    def test_without_converter_falls_back_to_docx_preview(self, tmp_path, monkeypatch):
        self._setup(tmp_path, monkeypatch, False)
        f = client.get("/pdf/ABCD1234").json()["files"][0]
        assert f["renderer"] == "docx" and f["converted"] is False
        assert f["url"] == "/pdf/ABCD1234/file?name=letter.docx"
        assert (
            client.get("/pdf/ABCD1234/file?name=letter.docx&as=pdf").status_code == 404
        )