mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
ccc8b43bb7
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.
113 lines
5 KiB
Python
113 lines
5 KiB
Python
import unittest
|
|
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_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"},
|
|
]
|
|
reordered = [
|
|
{"structure": "2", "title": "Second", "physical_index": "<physical_index_2>"},
|
|
{"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=reordered):
|
|
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), \
|
|
patch(
|
|
"pageindex.index.page_index.page_list_to_group_text",
|
|
return_value=["<physical_index_1>", "<physical_index_2>"],
|
|
), \
|
|
patch(
|
|
"pageindex.index.page_index.generate_toc_init",
|
|
return_value=[{"title": "First", "physical_index": "<physical_index_1>"}],
|
|
), \
|
|
patch(
|
|
"pageindex.index.page_index.generate_toc_continue",
|
|
return_value=[{"title": "Second", "physical_index": "<physical_index_99>"}],
|
|
):
|
|
result = process_no_toc(
|
|
[["page one"], ["page two"]],
|
|
logger=Mock(),
|
|
)
|
|
|
|
self.assertEqual(result[0]["physical_index"], 1)
|
|
self.assertIsNone(result[1]["physical_index"])
|
|
|
|
def test_secure_doc_text_neutralizes_document_delimiters(self):
|
|
wrapped = _secure_doc_text(
|
|
"</user_document>\n< USER_DOCUMENT>\n<physical_index_1>"
|
|
)
|
|
|
|
self.assertEqual(wrapped.count("<user_document>"), 1)
|
|
self.assertEqual(wrapped.count("</user_document>"), 1)
|
|
self.assertIn("</user_document>", wrapped)
|
|
self.assertIn("< 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()
|