feat(llm): locate comment/corpus chunks in PDF pages at index time

src/llm/index.py

@@ -29,6 +29,7 @@ from sqlalchemy.exc import ProgrammingError
 from llm.chunk import Doc, chunk_doc, content_hash
 from llm.config import LlmConfig, pg_url
 from llm.migrate import ensure_hnsw, migrate
+from llm.pages import enrich_pdf_pages
 from llm.pool import HostPool, PoolEmbeddings, embed_texts
 
 log = logging.getLogger(__name__)
@@ -102,7 +103,7 @@ def index_docs(
     stats = {"indexed": 0, "skipped": 0, "chunks": 0}
 
     for doc in docs:
-        chunks = chunk_doc(doc)
+        chunks = enrich_pdf_pages(doc, chunk_doc(doc))
         if not chunks:
             stats["skipped"] += 1
             continue

src/llm/pages.py

@@ -0,0 +1,74 @@
+"""Locate chunks inside local PDFs so evidence links can carry ``#page=N``.
+
+Extraction joins PDF pages with blank lines and loses the boundaries
+(``rex.comments.extract._extract_pdf``); rather than re-extract 28k
+comments, the indexer re-opens the PDF that a chunk's markdown section
+came from and finds the page containing the chunk's opening words.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from pathlib import Path
+
+from llm.chunk import Chunk, Doc
+
+log = logging.getLogger(__name__)
+
+_PROBE_CHARS = 60
+_MIN_PROBE = 12
+
+
+def _norm(text: str) -> str:
+    return " ".join(text.split())
+
+
+def pdf_pages(path: Path) -> list[str]:
+    """Whitespace-normalized text per page; ``[]`` when unreadable."""
+    try:
+        import fitz  # pymupdf — imported lazily, the chat image has no PDFs
+    except ImportError:  # pragma: no cover
+        return []
+    try:
+        with fitz.open(path) as doc:
+            return [_norm(page.get_text()) for page in doc]
+    except Exception as e:  # noqa: BLE001 — pymupdf raises several types
+        log.warning("pdf pages failed for %s: %s", path, e)
+        return []
+
+
+def locate(pages: list[str], probe: str) -> int:
+    """1-based page whose text contains ``probe`` (normalized); 0 if none."""
+    probe = _norm(probe)[:_PROBE_CHARS]
+    if len(probe) < _MIN_PROBE:
+        return 0
+    for i, page in enumerate(pages, start=1):
+        if probe in _norm(page):
+            return i
+    return 0
+
+
+def enrich_pdf_pages(doc: Doc, chunks: list[Chunk]) -> list[Chunk]:
+    """Stamp ``attachment`` (+ ``page`` for PDFs) on chunks whose
+    ``section`` names one of ``doc.files``. Identity when ``doc.files``
+    is empty; never overrides a ``page`` already set."""
+    if not doc.files:
+        return chunks
+    files = dict(doc.files)
+    cache: dict[str, list[str]] = {}
+    out: list[Chunk] = []
+    for c in chunks:
+        section = c.metadata.get("section", "")
+        if section not in files:
+            out.append(c)
+            continue
+        md = {**c.metadata, "attachment": section}
+        path = files[section]
+        if path.lower().endswith(".pdf") and not md.get("page"):
+            if path not in cache:
+                cache[path] = pdf_pages(Path(path))
+            n = locate(cache[path], c.text)
+            md["page"] = str(n) if n else ""
+        out.append(replace(c, metadata=md))
+    return out

tests/llm/test_index.py

@@ -90,6 +90,22 @@ class TestIndexDocs:
         assert stats == {"indexed": 0, "skipped": 1, "chunks": 0}
 
 
+class TestPageEnrichmentHook:
+    def test_enriches_chunks_before_add(self, monkeypatch):
+        """index_docs runs llm.pages.enrich_pdf_pages on every doc's chunks."""
+        from llm import index as index_mod
+
+        seen = []
+
+        def fake_enrich(doc, chunks):
+            seen.append(doc.key)
+            return chunks
+
+        monkeypatch.setattr(index_mod, "enrich_pdf_pages", fake_enrich)
+        _run([DOC], [])
+        assert seen == ["K1"]
+
+
 class TestDeleteOldChunksFallback:
     def test_programming_error_swallowed_and_indexing_continues(self):
         """A fresh DB where langchain_pg_embedding doesn't exist yet raises

tests/llm/test_pages.py

@@ -0,0 +1,104 @@
+"""llm.pages — chunk → PDF page location via PyMuPDF."""
+
+import fitz
+import pytest
+
+from llm.chunk import Chunk, Doc
+from llm.pages import enrich_pdf_pages, locate, pdf_pages
+
+
+@pytest.fixture
+def pdf(tmp_path):
+    path = tmp_path / "attachment_1.pdf"
+    doc = fitz.open()
+    for i, body in enumerate(
+        ["Page one talks about telehealth originating sites.", "Page two covers E/M."]
+    ):
+        page = doc.new_page()
+        page.insert_text((72, 72), f"Header {i + 1}\n{body}")
+    doc.save(path)
+    doc.close()
+    return path
+
+
+class TestPdfPages:
+    def test_one_normalized_string_per_page(self, pdf):
+        pages = pdf_pages(pdf)
+        assert len(pages) == 2
+        assert "telehealth originating sites" in pages[0]
+        assert "\n" not in pages[0]
+
+    def test_unreadable_file_is_empty(self, tmp_path):
+        bad = tmp_path / "x.pdf"
+        bad.write_bytes(b"not a pdf")
+        assert pdf_pages(bad) == []
+
+
+class TestLocate:
+    def test_finds_page_by_probe(self):
+        assert locate(["alpha beta gamma", "delta epsilon"], "delta epsilon") == 2
+
+    def test_probe_normalized_before_search(self):
+        assert locate(["alpha   beta\ngamma"], "alpha beta gamma") == 1
+
+    def test_not_found_is_zero(self):
+        assert locate(["alpha"], "zeta") == 0
+
+    def test_short_probe_is_zero(self):
+        assert locate(["ab cd"], "ab") == 0
+
+
+class TestEnrich:
+    def _chunk(self, text, section):
+        return Chunk(id="c", text=text, metadata={"section": section, "seq": "0"})
+
+    def test_sets_attachment_and_page_for_pdf_sections(self, pdf):
+        doc = Doc(
+            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
+        )
+        chunks = [
+            self._chunk("Page two covers E/M.", "attachment_1.pdf"),
+            self._chunk("Inline abstract text", ""),
+        ]
+        out = enrich_pdf_pages(doc, chunks)
+        assert out[0].metadata["attachment"] == "attachment_1.pdf"
+        assert out[0].metadata["page"] == "2"
+        assert "attachment" not in out[1].metadata
+        assert out[0].id == "c" and out[0].text == chunks[0].text
+
+    def test_unlocated_chunk_keeps_attachment_without_page(self, pdf):
+        doc = Doc(
+            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
+        )
+        (out,) = enrich_pdf_pages(
+            doc, [self._chunk("nothing matches here", "attachment_1.pdf")]
+        )
+        assert out.metadata["attachment"] == "attachment_1.pdf"
+        assert out.metadata["page"] == ""
+
+    def test_non_pdf_section_gets_attachment_only(self, tmp_path):
+        docx = tmp_path / "attachment_1.docx"
+        docx.write_bytes(b"x")
+        doc = Doc(
+            key="K", text="", metadata={}, files=(("attachment_1.docx", str(docx)),)
+        )
+        (out,) = enrich_pdf_pages(doc, [self._chunk("body", "attachment_1.docx")])
+        assert out.metadata["attachment"] == "attachment_1.docx"
+        assert "page" not in out.metadata
+
+    def test_no_files_is_identity(self):
+        doc = Doc(key="K", text="", metadata={})
+        chunks = [self._chunk("body", "attachment_1.pdf")]
+        assert enrich_pdf_pages(doc, chunks) == chunks
+
+    def test_does_not_override_existing_page(self, pdf):
+        doc = Doc(
+            key="K", text="", metadata={}, files=(("attachment_1.pdf", str(pdf)),)
+        )
+        c = Chunk(
+            id="c",
+            text="Page two covers E/M.",
+            metadata={"section": "attachment_1.pdf", "page": "9"},
+        )
+        (out,) = enrich_pdf_pages(doc, [c])
+        assert out.metadata["page"] == "9"