Replace manual JSON parsing with Instructor for structured output reliability

[pegahmansourian]

Problem

The codebase uses manual JSON prompt instructions across 11 locations in
page_index.py, with a custom extract_json() parser in utils.py that
handles malformed output through multiple fallback strategies.

When the LLM returns slightly malformed JSON, extract_json() silently
returns {}, causing downstream KeyError crashes like:

json_content['completed']  # KeyError if extract_json returned {}

Solution

Replace manual JSON parsing with Instructor,
which is patched directly onto the existing LiteLLM client already used in the codebase.

Changes

utils.py

• Add sync_instructor_client and async_instructor_client patched via instructor.from_litellm()
• Add llm_structured() for sync structured calls
• Add llm_astructured() for async structured calls
• extract_json() and get_json_content() can be removed once all callers are updated

page_index.py

• Add Pydantic models for all 11 structured outputs
• Replace all extract_json() calls with llm_structured() / llm_astructured()
• Remove "Directly return the final JSON structure" from all prompts — Instructor handles this automatically

Benefits

• Removes 40 lines of fragile JSON parsing with multiple fallback hacks
• Adds automatic retry on malformed output — no more silent {} returns
• Adds type safety via Pydantic — Literal["yes", "no"] fields are always valid
• No behavior change — same LiteLLM routing, same models, same logic
• instructor is the only new dependency

New dependency

This PR adds one new dependency: instructor

Add to requirements.txt:

instructor

Or install directly:

pip install instructor

[KylinMountain]

@pegahmansourian Thanks for this — moving to structured output (instructor + Pydantic) is a solid direction and it removes a whole class of manual-parsing bugs. A few things to address before we can take it:

1. Move the schemas out of page_index.py. The ~14 BaseModel definitions don't belong in the core logic file — please put them in a dedicated module (e.g. pageindex/schemas.py) and import them. Keeps page_index.py focused and lets the schemas be reused/tested independently.

2. The per-field prompt guidance is being dropped. For example detect_page_index used to tell the model what to put in thinking:"thinking": <why do you think there are page numbers/indices given within the table of contents>.
   After the change that's a bare thinking: str with no guidance, so the model loses that instruction. instructor serializes Pydantic field descriptions into the schema it sends the model, so please carry the original hints. Without this, the change risks lowering answer quality rather than improving it.

3. Do you have before/after test data? Since this PR is about reliability, it'd help to see concrete numbers — e.g. JSON parse-failure rate or TOC accuracy over a set of PDFs, old vs new, ideally including a weaker/local (Ollama) model since Mode.JSON relies on the provider honoring the schema. Right now the gain is asserted but not measured.

Minor/practical: instructor isn't in requirements.txt (install will break), and the branch now conflicts with main — will need a rebase.

[pegahmansourian]

Thanks again for the detailed review — here's where things stand now:

1. Schemas moved. Done — all 14 BaseModel definitions live in pageindex/schemas.py, imported into page_index.py. Confirmed nothing else references them except via import.

2. Per-field guidance preserved. Confirmed unchanged from the original inline prompts; every field's description matches what was previously embedded directly in the prompt text, e.g.:

class PageIndexDetection(BaseModel):
    thinking: str = Field(
        validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"),
        description="why do you think there are page numbers/indices given within the table of contents"
    )
    page_index_given_in_toc: Literal["yes", "no"] = Field(...)

I also added validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale") on reasoning fields — testing showed some models use a synonym key instead of thinking, so this widens what validates without touching the guidance itself.

3. Test data. I don't have a full multi-PDF, multi-model parse-failure-rate benchmark. I don't have budget for paid API access at the scale that would be needed (GPT-4o, etc.), so I can't produce the complete hosted-vs-local comparison you'd ideally want, and I want to be upfront about that rather than let the gap go unaddressed.

What I do have: tests/test_instructor_migration.py (25 tests, all passing, no API calls required, fully mocked) covering the concrete reliability properties this PR touches:

• The TOCIndexItem.structure bug this PR fixes (Optional[str] without default=None isn't actually omittable in Pydantic v2. Any response leaving the field out failed validation with Field required, regardless of model)
• add_page_offset_to_toc_json's guard against calculate_page_offset returning None
• finish_reason-based truncation detection in llm_structured/llm_astructured, max_tokens threading, max_retries capped at 1
• AliasChoices synonym handling
• End-to-end toc_transformer/toc_index_extractor behavior

I also ran real (non-mocked) testing against one PDF (NVIDIA's FY26 Q3 10-Q, 44 pages) using a local model via Ollama (qwen3) and, separately, Cohere's command-a-03-2025 via LiteLLM. That surfaced real, reproducible reliability issues specific to smaller models: task substitution on dense tabular content, and measurable degradation across Instructor's default retry behavior (by the 4th retry, the model was responding to the Pydantic validation error text as if it were the task, rather than the original prompt). These aren't fixed by this PR; I'm tracking them for a follow-up rather than expanding scope here. Happy to write these findings up in more detail if useful context for review.

Minor/practical:

• instructor>=1.13.0,<2.0.0 is now in requirements.txt, confirmed clean install.
• Rebased onto current main. This surfaced something worth flagging directly: main picked up security-hardening commits since I opened this PR (prompt-injection escaping via _wrap_doc_text/_secure_doc_text, and _validate_chunk_physical_indices to reject a physical_index the model claims that isn't actually present in the chunk sent) that touch the exact functions this PR migrates. My Instructor-based rewrites had neither protection, since they were written before that hardening landed. I've ported both into toc_index_extractor and toc_transformer, wrapping untrusted content before it enters the prompt, and validating returned indices against the real chunk afterward, reusing your existing helper functions unmodified.

While working through that, I found a related gap worth calling out rather than quietly leaving: toc_transformer has no equivalent of check_if_toc_transformation_is_complete, which catches something schema validation alone can't; a response that validates fine but silently skipped sections (finish_reason == "stop", not truncated, just incomplete). I've reused that function as-is, and I'm planning to wire it in as a Pydantic model_validator on TOCTransformation (using Instructor's context= passthrough so the validator can access the raw TOC text), so the completeness check participates in Instructor's own retry loop rather than a separate hand-rolled one. Not implemented yet; wanted to get this update in front of you rather than keep expanding the diff further first.

One deliberate trade-off I'm leaving as digestible-but-unresolved: your original toc_transformer had a chat-history continuation loop for genuinely truncated completions (up to 5 attempts, stitching partial JSON back together). Instructor's structured output doesn't compose cleanly with that approach. You can't easily validate a half-built Pydantic object mid-stream, so this branch fails fast instead with a clear error pointing at max_tokens. Should be a non-issue for the vast majority of documents given the token budget in use (8000), but it's a real behavior change for the case of an extremely long TOC.

Let me know if you'd like the completeness-check validator work done in this PR before another look, or if you'd rather review the current state first.