fix: improve JSON extraction and TOC fallback handling

[KairosMarco]

Summary

This PR improves PageIndex robustness when LLM calls return JSON in common non-ideal formats or omit optional fields during TOC/page-index extraction.

It keeps the existing indexing flow unchanged, but adds safer parsing and fallback behavior for provider responses that include:

• fenced JSON blocks,
• explanatory text before JSON,
• arrays with trailing text,
• Python-style literal tokens: None, True, False,
• missing JSON keys,
• object-shaped TOC output where list-shaped output is expected,
• missing page-offset or physical_index values.

Why

While running PageIndex over a FinanceBench PDF subset, I saw indexing failures from model response shape issues such as:

KeyError: 'toc_detected'
KeyError: 'page_index_given_in_toc'
AttributeError: 'dict' object has no attribute 'extend'
TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'
KeyError: 'physical_index'

The failures were not specific to one document format; they came from LLM JSON formatting variance.

Changes

• Make extract_json() tolerate fenced JSON, embedded JSON, arrays, trailing text, and Python-style literal tokens.
• Use safe defaults for TOC detector/completeness checks when parsed JSON is missing or not a dict.
• Normalize TOC generation output to list[dict] before list operations.
• Skip offset/page repair when the model output is missing required fields.
• Return a low-confidence no-TOC structure instead of raising Processing failed after fallback attempts.
• Add focused unittest coverage for the JSON parser and TOC fallback helpers.

Validation

python -m unittest discover -s tests
python -m py_compile pageindex\utils.py pageindex\page_index.py tests\test_json_resilience.py

Local result:

Ran 7 tests
OK

I also validated equivalent fixes in a local benchmark workspace:

Expanded PageIndex structures: 24 / 24 source documents
Expanded PageIndex retrieval-only QA: 25 / 25 generated
Expanded LLM QA: 25 / 25 generated

The benchmark artifacts are available here:

https://github.com/KairosMarco/pageindex-benchlab

[KairosMarco]

Hi maintainers, I wanted to check whether this scope is aligned with the project direction.

The PR is focused on JSON response resilience and TOC fallback handling, with unittest coverage. If this is too broad, I am happy to split it into a smaller parser-only PR first, then a separate TOC fallback PR.

[KylinMountain]

@KairosMarco Nice hardening PR — the defensive .get() calls, the graceful meta_processor fallback, and the _as_toc_list normalization all look good, and thanks for adding tests. One regression to fix before merge, plus two minor notes:

1. (please fix) extract_json no longer tolerates unescaped newlines inside string values. The old code did .replace('\n', ' '); the rewrite drops it, and json.loads/raw_decode reject control chars in strings:

extract_json('{"thinking": "line1\nline2", "answer": "yes"}')
# before -> {'thinking': 'line1 line2', 'answer': 'yes'}
# after  -> {}

Many prompts here ask for a multi-line thinking field, so this will silently degrade to {} → .get(..., 'no') (e.g. a real TOC read as "not detected"). Simplest fix is strict=False:

return json.loads(json_content, strict=False)
...
decoder = json.JSONDecoder(strict=False)

2. (minor) _normalize_json_candidate can corrupt string values — re.sub(r"\bNone\b", "null", ...) also rewrites the standalone word in e.g. {"title": "None of the above"}. Pre-existing behavior, not a blocker, but worth tightening since this is the JSON-hardening PR.

3. (minor) _as_toc_list doesn't recognize a table_of_contents wrapper — a {"table_of_contents": [...]} response normalizes to []. Adding that key alongside toc/items is a cheap safety net.

Also a tiny nit: tests/test_json_resilience.py starts with a UTF-8 BOM.

[KairosMarco]

Thanks for the careful review. I pushed a follow-up commit addressing the requested items:

• Fixed the newline regression by using json.JSONDecoder(strict=False) and json.loads(..., strict=False).
• Tightened Python-literal normalization so None / True / False are converted only outside JSON string values.
• Added table_of_contents wrapper support in _as_toc_list.
• Removed the UTF-8 BOM from tests/test_json_resilience.py.
• Added regression tests for unescaped newlines inside strings, preserving None inside string values, and table_of_contents wrappers.

Validation run locally:

py -3.10 -m unittest tests.test_json_resilience
py -3.10 -m py_compile pageindex\utils.py pageindex\page_index.py tests\test_json_resilience.py
py -3.10 -m unittest discover -s tests -p "test_*.py"
git diff --check

All passed.

[KylinMountain]

Thanks for the quick turnaround — I verified the follow-up and all four items are fixed correctly (unescaped newlines parse, None/True/False preserved inside strings, table_of_contents recognized, BOM removed, nice regression tests). 🙏

One same-class regression slipped through, and it points at a bigger simplification.

Bug: the trailing-comma cleanup in _normalize_json_candidate isn't string-aware. You made the Python-literal replacement string-aware, but the line right after — candidate.replace(",]", "]").replace(",}", "}") — runs unconditionally on the whole candidate before every parse, so it silently drops commas inside string values:

extract_json('{"note": "see items,] and more"}')  # -> {'note': 'see items] and more'}   (comma eaten)
extract_json('{"formula": "f(x,} )"}')             # -> {'formula': 'f(x} )'}

(Regression vs the old code, which only ran the comma cleanup in the fallback after the first parse failed, so valid JSON never hit it.)

Suggestion: rather than keep growing a hand-rolled repairer, consider json_repair (pure-Python). It handles everything this file is hand-rolling — correctly — and more. Verified locally:

• None/True/False outside strings, preserved inside strings ✓
• unescaped newlines in strings ✓
• the ,] / ,} in-string case above ✓ (no corruption)
• extra prose before/after the JSON ✓
• structural trailing commas ✓
• truncated / unclosed JSON auto-completed ✓ — e.g. [1,2,3 → [1,2,3]

That last one is a bonus: it recovers truncated responses structurally, which could reduce reliance on the "continue generation" retry loops (they cost extra LLM calls).

extract_json would then collapse to roughly:

from json_repair import repair_json

def extract_json(content):
    if not content:
        return {}
    return repair_json(content, return_objects=True)

and _replace_python_literals_outside_strings / _normalize_json_candidate / _json_candidates can be removed.

Caveat: json_repair is best-effort and never raises — it repairs even badly-broken input into a partial object (same tolerant spirit as returning {} today, just more capable).

If you'd rather keep the manual approach, the minimal fix is to make the trailing-comma removal string-aware too (fold it into the same scan as the literals). WDYT?

[KairosMarco]

Thanks again for catching this. I pushed a minimal follow-up fix in fbbc83c.

I kept the current manual path for this PR to avoid adding a new runtime dependency in the same hardening change, but I agree json_repair is a good direction for a broader simplification PR.

Changes in this follow-up:

• Replaced the global candidate.replace(",]", "]").replace(",}", "}") with _remove_trailing_commas_outside_strings(...).
• Preserved ,] and ,} when they appear inside JSON string values.
• Kept structural trailing-comma repair for arrays/objects.
• Added regression tests for both in-string cases and one structural trailing-comma case.

Local validation:

py -3.10 -m unittest tests.test_json_resilience
py -3.10 -m py_compile pageindex\utils.py pageindex\page_index.py tests\test_json_resilience.py
py -3.10 -m unittest discover -s tests -p "test_*.py"
git diff --check

All passed. The full unittest run emitted only a LiteLLM remote model-cost-map timeout warning and fell back to the local backup.

[KairosMarco]

I also did one more same-class pass over this PR and pushed a small follow-up in 234c54c.

Additional cleanup/hardening:

• Removed a duplicate import re introduced while adding JSON candidate scanning.
• Added dict-shape guards for the remaining extract_json(...) call sites that still immediately used .get(...):
  • check_title_appearance_in_start(...) now returns "no" for non-dict JSON.
  • toc_transformer(...) now returns [] for non-dict JSON in both completion paths.
• Added regression tests for those non-dict JSON cases.

Validation re-run:

py -3.10 -m unittest tests.test_json_resilience
py -3.10 -m py_compile pageindex\utils.py pageindex\page_index.py tests\test_json_resilience.py
py -3.10 -m unittest discover -s tests -p "test_*.py"
git diff --check

All passed.