diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index 22635b702..0b477dead 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -233,14 +233,34 @@ def raw_pages_from_doc( num_end_pages_to_keep=2, ): """[NOT PUBLIC API] Get raw pages from an input doc""" - if is_pdf_doc(doc) and hasattr(doc, "raw_pages"): + # Do NOT use `is_pdf_doc` here because MDDocuments could have + # "doc_type" == "pdf" and be treated as a single page doc + if isinstance(doc, PDFDocument) and hasattr(doc, "raw_pages"): raw_pages = doc.raw_pages - logger.debug( - "PDF Document from %s has %d raw pages", - doc.attrs.get("source", "unknown source"), - len(raw_pages), - ) - return doc.raw_pages + # failsafe check + if text_splitter is not None and len(raw_pages) == 1: + raw_pages = text_splitter.split_text(raw_pages[0]) + raw_pages = _down_select_pages( + raw_pages, + percent_raw_pages_to_keep, + max_raw_pages, + num_end_pages_to_keep, + ) + logger.debug( + "PDF Document from %s had 1 raw page; " + "has %d raw %s after splitting", + doc.attrs.get("source", "unknown source"), + len(raw_pages), + "page" if len(raw_pages) == 1 else "pages", + ) + else: + logger.debug( + "PDF Document from %s has %d raw %s", + doc.attrs.get("source", "unknown source"), + len(raw_pages), + "page" if len(raw_pages) == 1 else "pages", + ) + return raw_pages if text_splitter is None: logger.debug( @@ -255,6 +275,23 @@ def raw_pages_from_doc( return [] pages = text_splitter.split_text(text) + raw_pages = _down_select_pages( + pages, percent_raw_pages_to_keep, max_raw_pages, num_end_pages_to_keep + ) + + logger.debug( + "Document from %s has %d raw %s after splitting and trimming", + doc.attrs.get("source", "unknown source"), + len(raw_pages), + "page" if len(raw_pages) == 1 else "pages", + ) + return raw_pages + + +def _down_select_pages( + pages, percent_raw_pages_to_keep, max_raw_pages, num_end_pages_to_keep +): + """Down-select pages based on percentage and max limits""" num_to_keep = percent_raw_pages_to_keep / 100 * len(pages) num_raw_pages_to_keep = min(max_raw_pages, max(1, int(num_to_keep))) @@ -266,12 +303,6 @@ def raw_pages_from_doc( if last_page_index: raw_pages += pages[last_page_index:] - logger.debug( - "Document from %s has %d raw %s after splitting and trimming", - doc.attrs.get("source", "unknown source"), - len(raw_pages), - "page" if len(raw_pages) == 1 else "pages", - ) return raw_pages diff --git a/tests/python/unit/utilities/test_utilities_parsing.py b/tests/python/unit/utilities/test_utilities_parsing.py index 9b6eb9aea..55a2611ec 100644 --- a/tests/python/unit/utilities/test_utilities_parsing.py +++ b/tests/python/unit/utilities/test_utilities_parsing.py @@ -6,6 +6,7 @@ import numpy as np import pandas as pd import pytest +from langchain_text_splitters.character import RecursiveCharacterTextSplitter from compass.utilities.parsing import ( clean_backticks_from_llm_response, @@ -15,7 +16,9 @@ merge_overlapping_texts, num_ordinances_dataframe, ordinances_bool_index, + raw_pages_from_doc, ) +from elm.web.document import MDDocument, PDFDocument @pytest.mark.parametrize( @@ -36,6 +39,33 @@ def test_clean_backticks_from_llm_response(in_str, expected): assert clean_backticks_from_llm_response(in_str) == expected +@pytest.mark.parametrize("doc_class", [MDDocument, PDFDocument]) +def test_raw_pages_from_pdf_splits_oversized_raw_page(doc_class): + """Test PDF raw pages respect the supplied splitter budget""" + page = "word " * 100 + doc = doc_class([page], attrs={"doc_type": "pdf"}) + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0 + ) + + raw_pages = raw_pages_from_doc(doc, text_splitter=text_splitter) + + assert len(raw_pages) > 1 + assert all(len(raw_page) <= 50 for raw_page in raw_pages) + + +@pytest.mark.parametrize("doc_class", [MDDocument, PDFDocument]) +def test_raw_pages_from_pdf_preserves_page_within_splitter_budget(doc_class): + """Test PDF raw pages remain unchanged when already within budget""" + doc = doc_class(["Short PDF page"], attrs={"doc_type": "pdf"}) + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=50, chunk_overlap=0 + ) + + raw_pages = raw_pages_from_doc(doc, text_splitter=text_splitter) + assert raw_pages == doc.raw_pages + + @pytest.mark.parametrize( "in_str,expected", [