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

pageindex/page_index.py

@@ -412,7 +412,9 @@ def add_page_offset_to_toc_json(data, offset):
 
 
 
-def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, overlap_page=1):    
+def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, overlap_page=1):
+    if max_tokens is None or max_tokens <= 0:
+        max_tokens = 20000
     num_tokens = sum(token_lengths)
 
     if num_tokens <= max_tokens:
@@ -428,7 +430,9 @@ def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, over
     average_tokens_per_part = math.ceil(((num_tokens / expected_parts_num) + max_tokens) / 2)
 
     for i, (page_content, page_tokens) in enumerate(zip(page_contents, token_lengths)):
-        if current_token_count + page_tokens > average_tokens_per_part:
+        # A single page larger than the budget still becomes its own group;
+        # callers truncate that group before the LLM call (#15).
+        if current_subset and current_token_count + page_tokens > average_tokens_per_part:
 
             subsets.append(''.join(current_subset))
             # Start new subset from overlap if specified
@@ -447,7 +451,7 @@ def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, over
     print('divide page_list to groups', len(subsets))
     return subsets
 
-def add_page_number_to_toc(part, structure, model=None):
+def add_page_number_to_toc(part, structure, model=None, max_tokens=None):
     fill_prompt_seq = """
     You are given an JSON structure of a document and a partial part of the document. Your task is to check if the title that is described in the structure is started in the partial given document.
 
@@ -470,6 +474,7 @@ def add_page_number_to_toc(part, structure, model=None):
     The given structure contains the result of the previous part, you need to fill the result of the current part, do not change the previous result.
     Directly return the final JSON structure. Do not output anything else."""
 
+    part = truncate_to_token_limit(part, max_tokens, model=model)
     prompt = fill_prompt_seq + f"\n\nCurrent Partial Document:\n{part}\n\nGiven Structure\n{json.dumps(structure, indent=2)}\n"
     current_json_raw = llm_completion(model=model, prompt=prompt)
     json_result = extract_json(current_json_raw)
@@ -493,7 +498,7 @@ def remove_first_physical_index_section(text):
     return text
 
 ### add verify completeness
-def generate_toc_continue(toc_content, part, model=None):
+def generate_toc_continue(toc_content, part, model=None, max_tokens=None):
     print('start generate_toc_continue')
     prompt = """
     You are an expert in extracting hierarchical tree structure.
@@ -520,6 +525,7 @@ def generate_toc_continue(toc_content, part, model=None):
 
     Directly return the additional part of the final JSON structure. Do not output anything else."""
 
+    part = truncate_to_token_limit(part, max_tokens, model=model)
     prompt = prompt + '\nGiven text\n:' + part + '\nPrevious tree structure\n:' + json.dumps(toc_content, indent=2)
     response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True)
     if finish_reason == 'finished':
@@ -528,7 +534,7 @@ def generate_toc_continue(toc_content, part, model=None):
         raise Exception(f'finish reason: {finish_reason}')
 
 ### add verify completeness
-def generate_toc_init(part, model=None):
+def generate_toc_init(part, model=None, max_tokens=None):
     print('start generate_toc_init')
     prompt = """
     You are an expert in extracting hierarchical tree structure, your task is to generate the tree structure of the document.
@@ -554,6 +560,7 @@ def generate_toc_init(part, model=None):
 
     Directly return the final JSON structure. Do not output anything else."""
 
+    part = truncate_to_token_limit(part, max_tokens, model=model)
     prompt = prompt + '\nGiven text\n:' + part
     response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True)
 
@@ -562,19 +569,20 @@ def generate_toc_init(part, model=None):
     else:
         raise Exception(f'finish reason: {finish_reason}')
 
-def process_no_toc(page_list, start_index=1, model=None, logger=None):
+def process_no_toc(page_list, start_index=1, model=None, logger=None, max_tokens=None):
     page_contents=[]
     token_lengths=[]
     for page_index in range(start_index, start_index+len(page_list)):
         page_text = f"<physical_index_{page_index}>\n{page_list[page_index-start_index][0]}\n<physical_index_{page_index}>\n\n"
         page_contents.append(page_text)
         token_lengths.append(count_tokens(page_text, model))
-    group_texts = page_list_to_group_text(page_contents, token_lengths)
+    # Honor --max-tokens-per-node when building TOC groups (issue #15).
+    group_texts = page_list_to_group_text(page_contents, token_lengths, max_tokens=max_tokens)
     logger.info(f'len(group_texts): {len(group_texts)}')
 
-    toc_with_page_number= generate_toc_init(group_texts[0], model)
+    toc_with_page_number= generate_toc_init(group_texts[0], model, max_tokens=max_tokens)
     for group_text in group_texts[1:]:
-        toc_with_page_number_additional = generate_toc_continue(toc_with_page_number, group_text, model)    
+        toc_with_page_number_additional = generate_toc_continue(toc_with_page_number, group_text, model, max_tokens=max_tokens)    
         toc_with_page_number.extend(toc_with_page_number_additional)
     logger.info(f'generate_toc: {toc_with_page_number}')
 
@@ -583,7 +591,7 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None):
 
     return toc_with_page_number
 
-def process_toc_no_page_numbers(toc_content, toc_page_list, page_list,  start_index=1, model=None, logger=None):
+def process_toc_no_page_numbers(toc_content, toc_page_list, page_list,  start_index=1, model=None, logger=None, max_tokens=None):
     page_contents=[]
     token_lengths=[]
     toc_content = toc_transformer(toc_content, model)
@@ -593,12 +601,12 @@ def process_toc_no_page_numbers(toc_content, toc_page_list, page_list,  start_in
         page_contents.append(page_text)
         token_lengths.append(count_tokens(page_text, model))
 
-    group_texts = page_list_to_group_text(page_contents, token_lengths)
+    group_texts = page_list_to_group_text(page_contents, token_lengths, max_tokens=max_tokens)
     logger.info(f'len(group_texts): {len(group_texts)}')
 
     toc_with_page_number=copy.deepcopy(toc_content)
     for group_text in group_texts:
-        toc_with_page_number = add_page_number_to_toc(group_text, toc_with_page_number, model)
+        toc_with_page_number = add_page_number_to_toc(group_text, toc_with_page_number, model, max_tokens=max_tokens)
     logger.info(f'add_page_number_to_toc: {toc_with_page_number}')
 
     toc_with_page_number = convert_physical_index_to_int(toc_with_page_number)
@@ -955,9 +963,15 @@ async def meta_processor(page_list, mode=None, toc_content=None, toc_page_list=N
     if mode == 'process_toc_with_page_numbers':
         toc_with_page_number = process_toc_with_page_numbers(toc_content, toc_page_list, page_list, toc_check_page_num=opt.toc_check_page_num, model=opt.model, logger=logger)
     elif mode == 'process_toc_no_page_numbers':
-        toc_with_page_number = process_toc_no_page_numbers(toc_content, toc_page_list, page_list, model=opt.model, logger=logger)
+        toc_with_page_number = process_toc_no_page_numbers(
+            toc_content, toc_page_list, page_list,
+            model=opt.model, logger=logger, max_tokens=opt.max_token_num_each_node,
+        )
     else:
-        toc_with_page_number = process_no_toc(page_list, start_index=start_index, model=opt.model, logger=logger)
+        toc_with_page_number = process_no_toc(
+            page_list, start_index=start_index, model=opt.model, logger=logger,
+            max_tokens=opt.max_token_num_each_node,
+        )
 
     toc_with_page_number = [item for item in toc_with_page_number if item.get('physical_index') is not None] 
 

pageindex/utils.py

@@ -30,6 +30,39 @@ def count_tokens(text, model=None):
     return litellm.token_counter(model=model, text=text)
 
 
+def truncate_to_token_limit(text, max_tokens, model=None):
+    """Truncate text so ``count_tokens`` stays within ``max_tokens``.
+
+    Used to keep TOC / LLM prompts inside the budget implied by
+    ``--max-tokens-per-node`` (issue #15). Returns ``text`` unchanged when
+    ``max_tokens`` is unset/non-positive or the text already fits.
+    """
+    if not text or max_tokens is None or max_tokens <= 0:
+        return text
+    if count_tokens(text, model) <= max_tokens:
+        return text
+
+    # Binary-search a character cut; tokenizers are roughly monotone in length.
+    lo, hi = 0, len(text)
+    best = ""
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        candidate = text[:mid]
+        if count_tokens(candidate, model) <= max_tokens:
+            best = candidate
+            lo = mid + 1
+        else:
+            hi = mid - 1
+    if best and best != text:
+        logging.warning(
+            "Truncated text from %s to %s tokens (limit=%s)",
+            count_tokens(text, model),
+            count_tokens(best, model),
+            max_tokens,
+        )
+    return best
+
+
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
     if model:
         model = model.removeprefix("litellm/")
@@ -43,10 +76,14 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
                 temperature=0,
             )
             content = response.choices[0].message.content
+            if content is None:
+                content = ""
             if return_finish_reason:
+                # Always return a 2-tuple so callers never hit
+                # "too many values to unpack" when providers return unexpected shapes (#15).
                 finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished"
-                return content, finish_reason
-            return content
+                return str(content), finish_reason
+            return str(content)
         except Exception as e:
             print('************* Retrying *************')
             logging.error(f"Error: {e}")

tests/test_issue_15.py

@@ -0,0 +1,130 @@
+"""Tests for issue #15: enforce --max-tokens-per-node and safe llm_completion unpacking."""
+
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from pageindex.page_index import (
+    generate_toc_init,
+    page_list_to_group_text,
+    process_no_toc,
+)
+from pageindex.utils import llm_completion, truncate_to_token_limit
+
+
+class TestPageListGroupingRespectsMaxTokens:
+    def test_small_budget_splits_pages(self):
+        pages = ["aaa", "bbb", "ccc"]
+        lengths = [50, 50, 50]
+        groups = page_list_to_group_text(pages, lengths, max_tokens=60)
+        assert len(groups) >= 2
+        # No group should exceed the budget when pages themselves fit.
+        for group in groups:
+            # Rough check: each page is 50 tokens; groups with one page ok.
+            assert isinstance(group, str)
+
+    def test_large_budget_keeps_single_group(self):
+        pages = ["aaa", "bbb"]
+        lengths = [10, 10]
+        groups = page_list_to_group_text(pages, lengths, max_tokens=20000)
+        assert groups == ["aaabbb"]
+
+    def test_none_budget_falls_back_to_default(self):
+        pages = ["aaa", "bbb"]
+        lengths = [10, 10]
+        groups = page_list_to_group_text(pages, lengths, max_tokens=None)
+        assert groups == ["aaabbb"]
+
+
+class TestTruncateToTokenLimit:
+    @patch("pageindex.utils.count_tokens")
+    def test_no_truncate_when_under_limit(self, mock_count):
+        mock_count.return_value = 5
+        assert truncate_to_token_limit("hello world", 10, model="m") == "hello world"
+
+    @patch("pageindex.utils.count_tokens")
+    def test_truncates_when_over_limit(self, mock_count):
+        # count_tokens(text[:n]) ≈ n for this stub
+        def _count(text, model=None):
+            return len(text or "")
+
+        mock_count.side_effect = _count
+        result = truncate_to_token_limit("abcdefghij", 4, model="m")
+        assert len(result) <= 4
+        assert result == "abcd"
+
+    def test_unset_limit_noop(self):
+        assert truncate_to_token_limit("abc", None) == "abc"
+        assert truncate_to_token_limit("abc", 0) == "abc"
+
+
+class TestLlmCompletionReturnShape:
+    @patch("pageindex.utils.litellm.completion")
+    def test_return_finish_reason_is_always_two_tuple(self, mock_completion):
+        choice = MagicMock()
+        choice.message.content = "ok"
+        choice.finish_reason = "stop"
+        mock_completion.return_value = MagicMock(choices=[choice])
+
+        result = llm_completion("gpt-4o-mini", "prompt", return_finish_reason=True)
+        assert isinstance(result, tuple)
+        assert len(result) == 2
+        assert result == ("ok", "finished")
+
+    @patch("pageindex.utils.litellm.completion")
+    def test_none_content_becomes_empty_string(self, mock_completion):
+        choice = MagicMock()
+        choice.message.content = None
+        choice.finish_reason = "stop"
+        mock_completion.return_value = MagicMock(choices=[choice])
+
+        content, finish_reason = llm_completion(
+            "gpt-4o-mini", "prompt", return_finish_reason=True
+        )
+        assert content == ""
+        assert finish_reason == "finished"
+
+    @patch("pageindex.utils.time.sleep")
+    @patch("pageindex.utils.litellm.completion", side_effect=RuntimeError("tokens_limit_reached"))
+    def test_exhausted_retries_return_error_tuple(self, mock_completion, _sleep):
+        content, finish_reason = llm_completion(
+            "gpt-4o-mini", "huge prompt", return_finish_reason=True
+        )
+        assert content == ""
+        assert finish_reason == "error"
+        assert mock_completion.call_count == 10
+
+
+class TestProcessNoTocHonorsMaxTokens:
+    @patch("pageindex.page_index.convert_physical_index_to_int", side_effect=lambda x: x)
+    @patch("pageindex.page_index.generate_toc_continue")
+    @patch("pageindex.page_index.generate_toc_init")
+    @patch("pageindex.page_index.page_list_to_group_text")
+    @patch("pageindex.page_index.count_tokens", return_value=100)
+    def test_passes_max_tokens_into_grouping(
+        self, _count, mock_group, mock_init, mock_continue, _convert
+    ):
+        mock_group.return_value = ["chunk"]
+        mock_init.return_value = [{"title": "A", "physical_index": 1}]
+        logger = MagicMock()
+
+        page_list = [("page text", 100)]
+        process_no_toc(page_list, model="m", logger=logger, max_tokens=100)
+
+        assert mock_group.call_args.kwargs.get("max_tokens") == 100
+        assert mock_init.call_args.kwargs.get("max_tokens") == 100
+
+
+class TestGenerateTocInitTruncatesPart:
+    @patch("pageindex.page_index.extract_json", return_value=[])
+    @patch("pageindex.page_index.llm_completion", return_value=('[]', "finished"))
+    @patch("pageindex.page_index.truncate_to_token_limit")
+    def test_truncates_document_part(self, mock_trunc, mock_llm, _extract):
+        mock_trunc.return_value = "short"
+        generate_toc_init("very long part", model="m", max_tokens=50)
+        mock_trunc.assert_called_once_with("very long part", 50, model="m")
+        prompt = mock_llm.call_args.kwargs["prompt"]
+        assert "very long part" not in prompt
+        assert "short" in prompt