apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-07-24 21:41:04 +02:00
Code Issues Projects Releases Packages Wiki Activity

fix: align merged hardening with dev's robustness and content fidelity

Browse source 
The prompt-injection hardening ported from main did not fit dev's direction
(graceful degradation + content-faithful, reasoning-based index). Adjust it:

- process_toc_no_page_numbers: a count-mismatched or reordered/renamed LLM
  response no longer raises ValueError (which aborted the entire index at the
  uncaught top-level path and defeated the return_exceptions degradation the
  rest of the pipeline uses). Skip the untrusted chunk and continue; the
  accuracy check falls back to another mode when too little gets filled.
- _secure_doc_text: drop the keyword blocklist that replaced phrases like
  "act as"/"disregard" with [REDACTED]. Those occur in legitimate titles and
  prose, so redaction corrupted document content. Keep the <user_document>
  framing + _SYSTEM_HARDENING system instruction as the injection defense.
- _validate_chunk_physical_indices: tolerate non-list input (extract_json
  returns {} on parse failure and may return a JSON object) instead of
  crashing while iterating a dict.
- Remove _validate_physical_indices / _parse_physical_index: range validation
  is already done by validate_and_truncate_physical_indices in meta_processor
  and marker->int by convert_physical_index_to_int; the helpers duplicated both.

Tests updated to assert the skip-not-raise behavior and lock in no-redaction
and non-list tolerance. 314 passed, 2 skipped.
This commit is contained in:
mountain 2026-07-17 18:45:13 +08:00
parent 9f01ba41ac
commit ccc8b43bb7
2 changed files with 77 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

86
pageindex/index/page_index.py
View file

@ -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)", "&lt;", 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 <user_document> framing
plus _SYSTEM_HARDENING carry the injection defense without touching content.
"""
return _wrap_doc_text(text)
_PHYSICAL_INDEX_MARKER_RE = re.compile(r"^<physical_index_(\d+)>$")
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"<physical_index_(\d+)>", 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):

60
tests/test_page_index.py
View file

@ -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=["<physical_index_1> <physical_index_2>"]), \
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": "<physical_index_1>"}]
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=["<physical_index_1> <physical_index_2>"]), \
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("&lt; USER_DOCUMENT>", wrapped)
self.assertIn("<physical_index_1>", 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="<physical_index_1>"), {})
obj = {"table_of_contents": [{"physical_index": "<physical_index_1>"}]}
self.assertEqual(_validate_chunk_physical_indices(toc=obj, content="<physical_index_1>"), obj)
if __name__ == "__main__":
unittest.main()