merge upstream/main into dev; port #188 into the index pipeline

main advanced (litellm 1.84.0 #342, #188 TOC fixes, #281, README) while
dev turned pageindex/page_index.py and utils.py into deprecation shims
over pageindex/index/*. Both sides touched those two files, hence the
conflict.

Resolution:
- Keep dev's shims for the two top-level modules (the real implementation
  lives in pageindex/index/*). requirements.txt auto-merged to
  litellm==1.84.0.
- #188 ("prevent KeyError crash and context exhaustion in TOC
  processing") landed on main's top-level page_index.py, which is now a
  shim on dev — so its fixes were NOT in dev's index/page_index.py.
  Ported them into pageindex/index/page_index.py (preserving dev's
  IndexConfig/bool integration): .get() on the TOC check functions +
  detect_page_index, incremental-chat_history retry loops in
  extract_toc_content and toc_transformer, truncation moved before the
  loop, .get('table_of_contents', []) and the single_toc_item_index_fixer
  None guard.
- Repoint #188's merged test (tests/test_issue_163.py) at
  pageindex.index.page_index so its patches hit the module where the code
  now lives (they were silently hitting the shim → real LLM calls).

Full suite: 158 passed, 2 skipped.

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
mountain 2026-07-07 16:34:02 +08:00
commit a139bd53d9
10 changed files with 287 additions and 82 deletions

View file

@ -118,8 +118,8 @@ def toc_detector_single_page(content, model=None):
response = llm_completion(model=model, prompt=prompt)
# print('response', response)
json_content = extract_json(response)
return json_content['toc_detected']
json_content = extract_json(response)
return json_content.get('toc_detected', 'no')
def check_if_toc_extraction_is_complete(content, toc, model=None):
@ -137,7 +137,7 @@ def check_if_toc_extraction_is_complete(content, toc, model=None):
prompt = prompt + '\n Document:\n' + content + '\n Table of contents:\n' + toc
response = llm_completion(model=model, prompt=prompt)
json_content = extract_json(response)
return json_content['completed']
return json_content.get('completed', 'no')
def check_if_toc_transformation_is_complete(content, toc, model=None):
@ -155,7 +155,7 @@ def check_if_toc_transformation_is_complete(content, toc, model=None):
prompt = prompt + '\n Raw Table of contents:\n' + content + '\n Cleaned Table of contents:\n' + toc
response = llm_completion(model=model, prompt=prompt)
json_content = extract_json(response)
return json_content['completed']
return json_content.get('completed', 'no')
def extract_toc_content(content, model=None):
prompt = f"""
@ -175,27 +175,19 @@ def extract_toc_content(content, model=None):
{"role": "user", "content": prompt},
{"role": "assistant", "content": response},
]
prompt = f"""please continue the generation of table of contents , directly output the remaining part of the structure"""
new_response, finish_reason = llm_completion(model=model, prompt=prompt, chat_history=chat_history, return_finish_reason=True)
response = response + new_response
if_complete = check_if_toc_transformation_is_complete(content, response, model)
continue_prompt = "please continue the generation of table of contents, directly output the remaining part of the structure"
attempt = 0
max_attempts = 5
while not (if_complete == "yes" and finish_reason == "finished"):
attempt += 1
if attempt > max_attempts:
raise Exception('Failed to complete table of contents after maximum retries')
chat_history = [
{"role": "user", "content": prompt},
{"role": "assistant", "content": response},
]
prompt = f"""please continue the generation of table of contents , directly output the remaining part of the structure"""
new_response, finish_reason = llm_completion(model=model, prompt=prompt, chat_history=chat_history, return_finish_reason=True)
for attempt in range(max_attempts):
new_response, finish_reason = llm_completion(model=model, prompt=continue_prompt, chat_history=chat_history, return_finish_reason=True)
response = response + new_response
chat_history.append({"role": "user", "content": continue_prompt})
chat_history.append({"role": "assistant", "content": new_response})
if_complete = check_if_toc_transformation_is_complete(content, response, model)
if if_complete == "yes" and finish_reason == "finished":
break
else:
raise Exception('Failed to complete table of contents extraction after maximum retries')
return response
@ -217,7 +209,7 @@ def detect_page_index(toc_content, model=None):
response = llm_completion(model=model, prompt=prompt)
json_content = extract_json(response)
return json_content['page_index_given_in_toc']
return json_content.get('page_index_given_in_toc', 'no')
def toc_extractor(page_list, toc_page_list, model):
def transform_dots_to_colon(text):
@ -296,43 +288,41 @@ def toc_transformer(toc_content, model=None):
if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)
if if_complete == "yes" and finish_reason == "finished":
last_complete = extract_json(last_complete)
cleaned_response=convert_page_to_int(last_complete['table_of_contents'])
cleaned_response = convert_page_to_int(last_complete.get('table_of_contents', []))
return cleaned_response
last_complete = get_json_content(last_complete)
attempt = 0
chat_history = [
{"role": "user", "content": prompt},
{"role": "assistant", "content": last_complete},
]
continue_prompt = "Please continue the table of contents JSON structure from where you left off. Directly output only the remaining part."
position = last_complete.rfind('}')
if position != -1:
last_complete = last_complete[:position+2]
max_attempts = 5
while not (if_complete == "yes" and finish_reason == "finished"):
attempt += 1
if attempt > max_attempts:
raise Exception('Failed to complete toc transformation after maximum retries')
position = last_complete.rfind('}')
if position != -1:
last_complete = last_complete[:position+2]
prompt = f"""
Your task is to continue the table of contents json structure, directly output the remaining part of the json structure.
The response should be in the following JSON format:
for attempt in range(max_attempts):
The raw table of contents json structure is:
{toc_content}
The incomplete transformed table of contents json structure is:
{last_complete}
Please continue the json structure, directly output the remaining part of the json structure."""
new_complete, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True)
new_complete, finish_reason = llm_completion(model=model, prompt=continue_prompt, chat_history=chat_history, return_finish_reason=True)
if new_complete.startswith('```json'):
new_complete = get_json_content(new_complete)
last_complete = last_complete+new_complete
new_complete = get_json_content(new_complete)
last_complete = last_complete + new_complete
chat_history.append({"role": "user", "content": continue_prompt})
chat_history.append({"role": "assistant", "content": new_complete})
if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model)
if if_complete == "yes" and finish_reason == "finished":
break
else:
raise Exception('Failed to complete TOC transformation after maximum retries')
last_complete = extract_json(last_complete)
cleaned_response=convert_page_to_int(last_complete['table_of_contents'])
cleaned_response = convert_page_to_int(last_complete.get('table_of_contents', []))
return cleaned_response
@ -752,8 +742,11 @@ async def single_toc_item_index_fixer(section_title, content, model=None):
prompt = toc_extractor_prompt + '\nSection Title:\n' + str(section_title) + '\nDocument pages:\n' + content
response = await llm_acompletion(model=model, prompt=prompt)
json_content = extract_json(response)
return convert_physical_index_to_int(json_content['physical_index'])
json_content = extract_json(response)
physical_index = json_content.get('physical_index')
if physical_index is None:
return None
return convert_physical_index_to_int(physical_index)