diff --git a/pageindex/index/page_index.py b/pageindex/index/page_index.py index d6468d2..a0ea4ff 100644 --- a/pageindex/index/page_index.py +++ b/pageindex/index/page_index.py @@ -10,23 +10,6 @@ from concurrent.futures import ThreadPoolExecutor, as_completed ######################### Hardening for prompt injection patterns #################################################### -_INJECTION_PATTERNS = re.compile( - r"(?i)(" - r"system\s+override|" - r"ignore\s+(all\s+)?(previous|prior|above)\s+instructions?|" - r"forget\s+(all\s+)?(previous|prior|above)\s+instructions?|" - r"you\s+are\s+now|act\s+as|new\s+instructions?|" - r"do\s+not\s+follow|override\s+(the\s+)?(system|previous|prior)|" - r"disregard|jailbreak|ALL\s+sections\s+MUST" - r")" -) - - -def _sanitize_doc_text(text: str) -> str: - """Redact known prompt-injection keywords from PDF-extracted text.""" - return _INJECTION_PATTERNS.sub("[REDACTED]", text) - - def _wrap_doc_text(text: str) -> str: """Wrap untrusted document text in delimiter tags so the LLM treats it as data.""" text = re.sub(r"(?i)<(?=\s*/?\s*user_document\b)", "<", text) @@ -50,40 +33,19 @@ _SYSTEM_HARDENING = ( def _secure_doc_text(text: str) -> str: - """Sanitize + delimiter-frame a PDF text block before LLM injection.""" - return _wrap_doc_text(_sanitize_doc_text(text)) + """Delimiter-frame a document text block so the LLM treats it as data. + + Deliberately no keyword redaction: phrases like "act as" / "disregard" + also appear in legitimate titles and prose, and blanking them corrupts the + content this reasoning-based index relies on. The framing + plus _SYSTEM_HARDENING carry the injection defense without touching content. + """ + return _wrap_doc_text(text) _PHYSICAL_INDEX_MARKER_RE = re.compile(r"^$") -def _parse_physical_index(raw): - if raw is None: - return None - marker_match = _PHYSICAL_INDEX_MARKER_RE.match(str(raw).strip()) - if marker_match: - return int(marker_match.group(1)) - try: - return int(raw) - except (TypeError, ValueError): - return None - - -def _validate_physical_indices(toc: list, total_pages: int, start_index: int = 1) -> list: - """Nullify any physical_index the LLM produced that falls outside the real page range.""" - max_idx = start_index + total_pages - 1 - for entry in toc: - raw = entry.get("physical_index") - if raw is None: - continue - val = _parse_physical_index(raw) - if val is None or not (start_index <= val <= max_idx): - entry["physical_index"] = None - else: - entry["physical_index"] = val - return toc - - def _extract_chunk_marker_set(content: str) -> set: return {int(m) for m in re.findall(r"", content)} @@ -94,9 +56,16 @@ def _validate_chunk_physical_indices(toc: list, content: str) -> list: This prevents the model from referencing markers that exist elsewhere in the document but not in the current prompt. """ + if not isinstance(toc, list): + # extract_json returns {} on parse failure (or an object the LLM wrapped + # the array in); leave non-list payloads untouched instead of crashing. + return toc + valid_indices = _extract_chunk_marker_set(content) for entry in toc: + if not isinstance(entry, dict): + continue raw = entry.get("physical_index") if raw is None: continue @@ -711,11 +680,6 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None): toc=toc_with_page_number, content=group_texts[0] ) - toc_with_page_number = _validate_physical_indices( - toc=toc_with_page_number, - total_pages=len(page_list), - start_index=start_index - ) for group_text in group_texts[1:]: toc_with_page_number_additional = generate_toc_continue( @@ -727,11 +691,6 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None): toc=toc_with_page_number_additional, content=group_text ) - toc_with_page_number_additional = _validate_physical_indices( - toc=toc_with_page_number_additional, - total_pages=len(page_list), - start_index=start_index - ) toc_with_page_number.extend(toc_with_page_number_additional) logger.info(f'generate_toc: {toc_with_page_number}') @@ -756,16 +715,21 @@ def process_toc_no_page_numbers(toc_content, toc_page_list, page_list, start_in toc_with_page_number = copy.deepcopy(toc_content) for group_text in group_texts: llm_result = add_page_number_to_toc(group_text, toc_with_page_number, model) - if len(llm_result) != len(toc_with_page_number): - raise ValueError( - "LLM returned a different number of TOC entries than expected." - ) + # Don't trust a response that changed the entry count or reordered/renamed + # entries: skip filling from this chunk rather than aborting the whole + # document (meta_processor's accuracy check falls back to another mode if + # too little gets filled). Aborting here would defeat the graceful + # degradation the rest of the pipeline is built around. + if not isinstance(llm_result, list) or len(llm_result) != len(toc_with_page_number): + logger.info("Skipping chunk: LLM returned an unexpected number of TOC entries.") + continue if any( (update.get("structure"), update.get("title")) != (current.get("structure"), current.get("title")) for update, current in zip(llm_result, toc_with_page_number) ): - raise ValueError("LLM returned reordered or modified TOC entries.") + logger.info("Skipping chunk: LLM returned reordered or modified TOC entries.") + continue valid_indices = _extract_chunk_marker_set(group_text) for idx, current in enumerate(toc_with_page_number): diff --git a/tests/test_page_index.py b/tests/test_page_index.py index d7c99d4..35b9eba 100644 --- a/tests/test_page_index.py +++ b/tests/test_page_index.py @@ -3,13 +3,17 @@ from unittest.mock import Mock, patch from pageindex.index.page_index import ( _secure_doc_text, + _validate_chunk_physical_indices, process_no_toc, process_toc_no_page_numbers, ) class ProcessTocNoPageNumbersTest(unittest.TestCase): - def test_rejects_same_length_reordered_llm_toc(self): + def test_skips_same_length_reordered_llm_toc(self): + # A reordered/renamed LLM response must not be trusted, but it must also + # not abort the whole document: the chunk is skipped and processing + # continues with no physical_index filled from it. toc = [ {"structure": "1", "title": "First"}, {"structure": "2", "title": "Second"}, @@ -23,13 +27,37 @@ class ProcessTocNoPageNumbersTest(unittest.TestCase): patch("pageindex.index.page_index.count_tokens", return_value=1), \ patch("pageindex.index.page_index.page_list_to_group_text", return_value=[" "]), \ patch("pageindex.index.page_index.add_page_number_to_toc", return_value=reordered): - with self.assertRaises(ValueError): - process_toc_no_page_numbers( - "toc", - [], - [["page one"], ["page two"]], - logger=Mock(), - ) + result = process_toc_no_page_numbers( + "toc", + [], + [["page one"], ["page two"]], + logger=Mock(), + ) + + self.assertEqual(len(result), 2) + self.assertTrue(all(item.get("physical_index") is None for item in result)) + + def test_skips_count_mismatch_llm_toc(self): + # A response with a different entry count is untrusted -> skipped, not raised. + toc = [ + {"structure": "1", "title": "First"}, + {"structure": "2", "title": "Second"}, + ] + short = [{"structure": "1", "title": "First", "physical_index": ""}] + + with patch("pageindex.index.page_index.toc_transformer", return_value=toc), \ + patch("pageindex.index.page_index.count_tokens", return_value=1), \ + patch("pageindex.index.page_index.page_list_to_group_text", return_value=[" "]), \ + patch("pageindex.index.page_index.add_page_number_to_toc", return_value=short): + result = process_toc_no_page_numbers( + "toc", + [], + [["page one"], ["page two"]], + logger=Mock(), + ) + + self.assertEqual(len(result), 2) + self.assertTrue(all(item.get("physical_index") is None for item in result)) def test_process_no_toc_validates_continuation_chunks(self): with patch("pageindex.index.page_index.count_tokens", return_value=1), \ @@ -64,6 +92,22 @@ class ProcessTocNoPageNumbersTest(unittest.TestCase): self.assertIn("< USER_DOCUMENT>", wrapped) self.assertIn("", wrapped) + def test_secure_doc_text_preserves_legitimate_content(self): + # Framing must NOT redact legitimate prose/titles that happen to contain + # phrases a keyword blocklist would flag (this corrupts a reasoning-based + # index). Guards against re-introducing keyword redaction. + title = "Chapter 5: Act as a Servant Leader and Disregard Old Habits" + wrapped = _secure_doc_text(title) + self.assertIn(title, wrapped) + self.assertNotIn("[REDACTED]", wrapped) + + def test_validate_chunk_tolerates_non_list(self): + # extract_json returns {} on parse failure and may return a JSON object; + # the validator must pass it through, not crash iterating a dict. + self.assertEqual(_validate_chunk_physical_indices(toc={}, content=""), {}) + obj = {"table_of_contents": [{"physical_index": ""}]} + self.assertEqual(_validate_chunk_physical_indices(toc=obj, content=""), obj) + if __name__ == "__main__": unittest.main()