Skip to content

fix: enforce max-tokens-per-node for TOC prompts and harden llm returns#353

Open
fooSynaptic wants to merge 1 commit into
VectifyAI:mainfrom
fooSynaptic:fix/15-enforce-token-limit
Open

fix: enforce max-tokens-per-node for TOC prompts and harden llm returns#353
fooSynaptic wants to merge 1 commit into
VectifyAI:mainfrom
fooSynaptic:fix/15-enforce-token-limit

Conversation

@fooSynaptic

Copy link
Copy Markdown

Summary

Fixes #15.

--max-tokens-per-node was only used when recursively splitting large tree nodes. TOC construction still grouped pages with a hard-coded 20k budget and sent the full group text into the LLM prompt, so small-context models (e.g. gpt-4o-mini on a multi-page PDF) hit provider tokens_limit_reached errors and could crash on unexpected return shapes when unpacking response, finish_reason = ....

Root cause

  1. process_no_toc / process_toc_no_page_numbers called page_list_to_group_text(...) without passing opt.max_token_num_each_node, so the CLI flag had no effect on TOC prompt size.
  2. Even after grouping, a single oversized page/group could still exceed the intended budget; there was no pre-call truncation.
  3. Callers assume llm_completion(..., return_finish_reason=True) always returns a 2-tuple; provider/None content edge cases could produce unpack failures.

Fix

  • Wire max_tokens=opt.max_token_num_each_node into TOC page grouping.
  • Add truncate_to_token_limit and apply it to document parts in generate_toc_init / generate_toc_continue / add_page_number_to_toc before building the LLM prompt.
  • Make llm_completion always return (str, finish_reason) when return_finish_reason=True (including None content → "").

No behavior change when the document already fits under the configured budget.

What changed

File Change
pageindex/utils.py truncate_to_token_limit; harden llm_completion 2-tuple returns
pageindex/page_index.py Pass max_token_num_each_node into TOC grouping; truncate parts before TOC LLM calls
tests/test_issue_15.py Mock-based unit tests for grouping, truncation, return shape, and wiring

Test plan

  • python -m pytest tests/test_issue_15.py -v (11 passed; no live API key required)
  • Existing TOC robustness tests still importable alongside this change
python -m pytest tests/test_issue_15.py -v

Wire --max-tokens-per-node into TOC page grouping and truncate oversized
document parts before LLM calls so small context models stop hitting
tokens_limit_reached. Ensure llm_completion always returns a clean
2-tuple when return_finish_reason=True to avoid unpack crashes (VectifyAI#15).
@fooSynaptic

Copy link
Copy Markdown
Author

Hi @KylinMountain — could you take a look at this when you have a moment?

It fixes #15 by wiring --max-tokens-per-node into TOC page grouping, truncating oversized document parts before LLM calls, and hardening llm_completion so return_finish_reason=True always returns a clean 2-tuple. Covered with mock unit tests in tests/test_issue_15.py (no live API needed). Happy to adjust anything in review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Issue: Token limit not enforced → tokens_limit_reached + unpacking crash

1 participant