105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
"""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"
|