Make sure text splitter is used even for PDF's

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
 
 

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",
     [