From 7f8e1273843dbbbcdecf8b3ee973fd8135e65e52 Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 11:51:31 -0600 Subject: [PATCH 1/4] Add fix + test --- compass/utilities/parsing.py | 9 ++++-- .../unit/utilities/test_utilities_parsing.py | 30 +++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index 22635b702..d8a1a3fa4 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -233,14 +233,17 @@ 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", + "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 doc.raw_pages + return raw_pages if text_splitter is None: logger.debug( diff --git a/tests/python/unit/utilities/test_utilities_parsing.py b/tests/python/unit/utilities/test_utilities_parsing.py index 9b6eb9aea..17b2c567c 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 @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 +def test_raw_pages_from_pdf_splits_oversized_raw_page(): + """Test PDF raw pages respect the supplied splitter budget""" + page = "word " * 100 + doc = MDDocument([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) + + +def test_raw_pages_from_pdf_preserves_page_within_splitter_budget(): + """Test PDF raw pages remain unchanged when already within budget""" + doc = MDDocument(["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", [ From a3e773926e13ad4c24bea69fbfe677ce9cda834f Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 11:57:05 -0600 Subject: [PATCH 2/4] Add failsafe --- compass/utilities/parsing.py | 3 +++ .../unit/utilities/test_utilities_parsing.py | 16 ++++++++-------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index d8a1a3fa4..f7bdb69e8 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -237,6 +237,9 @@ def raw_pages_from_doc( # "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 + # failsafe check + if text_splitter is not None and len(raw_pages) == 1: + raw_pages = text_splitter.split_text(raw_pages[0]) logger.debug( "PDF Document from %s has %d raw %s", doc.attrs.get("source", "unknown source"), diff --git a/tests/python/unit/utilities/test_utilities_parsing.py b/tests/python/unit/utilities/test_utilities_parsing.py index 17b2c567c..55a2611ec 100644 --- a/tests/python/unit/utilities/test_utilities_parsing.py +++ b/tests/python/unit/utilities/test_utilities_parsing.py @@ -18,7 +18,7 @@ ordinances_bool_index, raw_pages_from_doc, ) -from elm.web.document import MDDocument +from elm.web.document import MDDocument, PDFDocument @pytest.mark.parametrize( @@ -39,10 +39,11 @@ def test_clean_backticks_from_llm_response(in_str, expected): assert clean_backticks_from_llm_response(in_str) == expected -def test_raw_pages_from_pdf_splits_oversized_raw_page(): +@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 = MDDocument([page], attrs={"doc_type": "pdf"}) + doc = doc_class([page], attrs={"doc_type": "pdf"}) text_splitter = RecursiveCharacterTextSplitter( chunk_size=50, chunk_overlap=0 ) @@ -53,16 +54,15 @@ def test_raw_pages_from_pdf_splits_oversized_raw_page(): assert all(len(raw_page) <= 50 for raw_page in raw_pages) -def test_raw_pages_from_pdf_preserves_page_within_splitter_budget(): +@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 = MDDocument(["Short PDF page"], attrs={"doc_type": "pdf"}) + doc = doc_class(["Short PDF page"], attrs={"doc_type": "pdf"}) text_splitter = RecursiveCharacterTextSplitter( - chunk_size=50, - chunk_overlap=0, + chunk_size=50, chunk_overlap=0 ) raw_pages = raw_pages_from_doc(doc, text_splitter=text_splitter) - assert raw_pages == doc.raw_pages From b537532e49ea1728c19b541acf7263144d57e30d Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 12:44:23 -0600 Subject: [PATCH 3/4] Add new function --- compass/utilities/parsing.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index f7bdb69e8..8783c8a85 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -261,6 +261,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))) @@ -272,12 +289,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 From ddb6fa0e4bf6a58970f1ed82c5df252267baad69 Mon Sep 17 00:00:00 2001 From: Paul Date: Wed, 5 Aug 2026 12:44:32 -0600 Subject: [PATCH 4/4] Use new function --- compass/utilities/parsing.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index 8783c8a85..0b477dead 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -240,12 +240,26 @@ def raw_pages_from_doc( # failsafe check if text_splitter is not None and len(raw_pages) == 1: raw_pages = text_splitter.split_text(raw_pages[0]) - 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", - ) + 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: