From 64c62c6dd42bf49818ec42b823f0e61931c5d099 Mon Sep 17 00:00:00 2001 From: pegah Date: Fri, 26 Jun 2026 11:12:49 -0400 Subject: [PATCH 1/4] Replace manual JSON parsing with Instructor for structured output reliability - Add llm_structured() and llm_astructured() in utils.py using instructor.from_litellm() - Add Pydantic models for all 11 structured outputs in page_index.py - Replace all extract_json() calls with typed Instructor calls - Remove 'Directly return the final JSON structure' from all prompts - Fixes silent KeyError crashes when extract_json() returned {} on malformed LLM output # Conflicts: # pageindex/page_index.py # pageindex/utils.py # Conflicts: # pageindex/page_index.py --- pageindex/page_index.py | 335 ++++++++++++---------------------------- pageindex/utils.py | 78 ++++------ 2 files changed, 132 insertions(+), 281 deletions(-) diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 083b26e1b..9c0dde20a 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -60,7 +60,7 @@ def _parse_physical_index(raw): return int(raw) except (TypeError, ValueError): return None - + def _validate_physical_indices(toc: list, total_pages: int, start_index: int = 1) -> list: """Nullify any physical_index the LLM produced that falls outside the real page range.""" max_idx = start_index + total_pages - 1 @@ -74,14 +74,13 @@ def _validate_physical_indices(toc: list, total_pages: int, start_index: int = 1 else: entry["physical_index"] = val return toc - + ################### check title in page ######################################################### async def check_title_appearance(item, page_list, start_index=1, model=None): title=item['title'] if 'physical_index' not in item or item['physical_index'] is None: - return {'list_index': item.get('list_index'), 'answer': 'no', 'title':title, 'page_number': None} - - + return {'list_index': item.get('list_index'), 'answer': 'no', 'title': title, 'page_number': None} + page_number = item['physical_index'] page_text = page_list[page_number-start_index][0] @@ -90,26 +89,11 @@ async def check_title_appearance(item, page_list, start_index=1, model=None): Your job is to check if the given section appears or starts in the given page_text. Note: do fuzzy matching, ignore any space inconsistency in the page_text. - The given section title is {title}. - The given page_text is: - {_secure_doc_text(page_text)} - - Reply format: - {{ - - "thinking": - "answer": "yes or no" (yes if the section appears or starts in the page_text, no otherwise) - }} - Directly return the final JSON structure. Do not output anything else.""" - - response = await llm_acompletion(model=model, prompt=prompt) - response = extract_json(response) - if 'answer' in response: - answer = response['answer'] - else: - answer = 'no' - return {'list_index': item['list_index'], 'answer': answer, 'title': title, 'page_number': page_number} + The given page_text is {_secure_doc_text(page_text)}.""" + + result = await llm_astructured(model=model, prompt=prompt, response_model=TitleAppearance) + return {'list_index': item['list_index'], 'answer': result.answer, 'title': title, 'page_number': page_number} async def check_title_appearance_in_start(title, page_text, model=None, logger=None): @@ -122,21 +106,12 @@ async def check_title_appearance_in_start(title, page_text, model=None, logger=N Note: do fuzzy matching, ignore any space inconsistency in the page_text. The given section title is {title}. - The given page_text is: - {_secure_doc_text(page_text)} - - reply format: - {{ - "thinking": - "start_begin": "yes or no" (yes if the section starts in the beginning of the page_text, no otherwise) - }} - Directly return the final JSON structure. Do not output anything else.""" - - response = await llm_acompletion(model=model, prompt=prompt) - response = extract_json(response) + The given page_text is {_secure_doc_text(page_text)}.""" + + result = await llm_astructured(model=model, prompt=prompt, response_model=TitleAppearanceInStart) if logger: - logger.info(f"Response: {response}") - return response.get("start_begin", "no") + logger.info(f"Response: {result}") + return result.start_begin async def check_title_appearance_in_start_concurrent(structure, page_list, model=None, logger=None): @@ -172,67 +147,37 @@ async def check_title_appearance_in_start_concurrent(structure, page_list, model def toc_detector_single_page(content, model=None): prompt = _SYSTEM_HARDENING + f""" Your job is to detect if there is a table of content provided in the given text. + Given text: {_secure_doc_text(content)} + Please note: abstract, summary, notation list, figure list, table list, etc. are not table of contents.""" - Given text: - {_secure_doc_text(content)} - - return the following JSON format: - {{ - "thinking": - "toc_detected": "", - }} - - Directly return the final JSON structure. Do not output anything else. - Please note: abstract,summary, notation list, figure list, table list, etc. are not table of contents.""" - - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return json_content.get('toc_detected', 'no') + result = llm_structured(model=model, prompt=prompt, response_model=TOCDetection) + return result.toc_detected def check_if_toc_extraction_is_complete(content, toc, model=None): prompt = f""" - You are given a partial document and a table of contents. - Your job is to check if the table of contents is complete, which it contains all the main sections in the partial document. - - Reply format: - {{ - "thinking": - "completed": "yes" or "no" - }} - Directly return the final JSON structure. Do not output anything else.""" + You are given a partial document and a table of contents. + Your job is to check if the table of contents is complete, which it contains all the main sections in the partial document. + Document: + {_secure_doc_text(content)} + Table of contents: + {_secure_doc_text(str(toc))}""" - prompt = ( - prompt - + '\n Document:\n' + _secure_doc_text(content) - + '\n Table of contents:\n' + _secure_doc_text(str(toc)) - ) - - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return json_content.get('completed', 'no') + result = llm_structured(model=model, prompt=prompt, response_model=TOCCompletionCheck) + return result.completed def check_if_toc_transformation_is_complete(content, toc, model=None): prompt = f""" - You are given a raw table of contents and a table of contents. - Your job is to check if the table of contents is complete. - - Reply format: - {{ - "thinking": - "completed": "yes" or "no" - }} - Directly return the final JSON structure. Do not output anything else.""" + You are given a raw table of contents and a table of contents. + Your job is to check if the table of contents is complete. + Raw Table of contents: + {_secure_doc_text(content)} + Cleaned Table of contents: + {_secure_doc_text(str(toc))}""" - prompt = ( - prompt - + '\n Raw Table of contents:\n' + _secure_doc_text(content) - + '\n Cleaned Table of contents:\n' + _secure_doc_text(str(toc)) - ) - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return json_content.get('completed', 'no') + result = llm_structured(model=model, prompt=prompt, response_model=TOCCompletionCheck) + return result.completed def extract_toc_content(content, model=None): prompt = f""" @@ -275,18 +220,10 @@ def detect_page_index(toc_content, model=None): Your job is to detect if there are page numbers/indices given within the table of contents. - Given text: {toc_content} + Given text: {toc_content}""" - Reply format: - {{ - "thinking": - "page_index_given_in_toc": "" - }} - Directly return the final JSON structure. Do not output anything else.""" - - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return json_content.get('page_index_given_in_toc', 'no') + result = llm_structured(model=model, prompt=prompt, response_model=PageIndexDetection) + return result.page_index_given_in_toc def toc_extractor(page_list, toc_page_list, model): def transform_dots_to_colon(text): @@ -332,35 +269,25 @@ def _validate_chunk_physical_indices(toc: list, content: str) -> list: def toc_index_extractor(toc, content, model=None): print('start toc_index_extractor') toc_extractor_prompt = """ - You are given a table of contents in a json format and several pages of a document, your job is to add the physical_index to the table of contents in the json format. - + You are given a table of contents in a json format and several pages of a document. + Your job is to add the physical_index to the table of contents in the json format. The provided pages contains tags like and to indicate the physical location of the page X. - - The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc. - - The response should be in the following JSON format: - [ - { - "structure": (string), - "title": , - "physical_index": "<physical_index_X>" (keep the format) - }, - ... - ] - + The structure variable is the numeric system which represents the index of the hierarchy section. Only add the physical_index to the sections that are in the provided pages. - If the section is not in the provided pages, do not add the physical_index to it. - Directly return the final JSON structure. Do not output anything else.""" + If the section is not in the provided pages, do not add the physical_index to it.""" prompt = ( - _SYSTEM_HARDENING + toc_extractor_prompt - + '\nTable of contents:\n' + _secure_doc_text(str(toc)) - + '\nDocument pages:\n' + _secure_doc_text(content) + _SYSTEM_HARDENING + toc_extractor_prompt + + '\nTable of contents:\n' + _secure_doc_text(str(toc)) + + '\nDocument pages:\n' + _secure_doc_text(content) ) - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return _validate_chunk_physical_indices(toc=json_content, content=content) - + + result = llm_structured(model=model, prompt=prompt, response_model=TOCIndexList) + items = [item.model_dump() for item in result.items] + return _validate_chunk_physical_indices(toc=items, content=content) + + + def toc_transformer(toc_content, model=None): print('start toc_transformer') init_prompt = """ @@ -549,40 +476,24 @@ def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, over def add_page_number_to_toc(part, structure, model=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. - - The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the physical location of the page X. - - If the full target section starts in the partial given document, insert the given JSON structure with the "start": "yes", and "start_index": "<physical_index_X>". - - If the full target section does not start in the partial given document, insert "start": "no", "start_index": None. - - The response should be in the following format. - [ - { - "structure": <structure index, "x.x.x" or None> (string), - "title": <title of the section>, - "start": "<yes or no>", - "physical_index": "<physical_index_X> (keep the format)" or 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.""" + You are given a JSON structure of a document and a partial part of the document. + Your task is to check if the title described in the structure starts in the partial given document. + The provided text contains tags like <physical_index_X> to indicate the physical location of pages. + If the full target section starts in the partial given document, insert "start": "yes" and "physical_index": "<physical_index_X>". + If not, insert "start": "no" and "physical_index": null. + The given structure contains the result of the previous part, do not change the previous result.""" part_text = ''.join(part) if isinstance(part, list) else part prompt = ( - _SYSTEM_HARDENING + fill_prompt_seq - + f"\n\nCurrent Partial Document:\n{_secure_doc_text(part_text)}" - + f"\n\nGiven Structure\n{_secure_doc_text(json.dumps(structure, indent=2))}\n" + _SYSTEM_HARDENING + fill_prompt_seq + + f"\n\nCurrent Partial Document:\n{_secure_doc_text(part_text)}" + + f"\n\nGiven Structure\n{_secure_doc_text(json.dumps(structure, indent=2))}\n" ) - - current_json_raw = llm_completion(model=model, prompt=prompt) - json_result = extract_json(current_json_raw) - + + result = llm_structured(model=model, prompt=prompt, response_model=PageNumberList) + json_result = [item.model_dump() for item in result.items] for item in json_result: - if 'start' in item: - del item['start'] + item.pop('start', None) return json_result @@ -605,73 +516,42 @@ def generate_toc_continue(toc_content, part, model=None): You are an expert in extracting hierarchical tree structure. You are given a tree structure of the previous part and the text of the current part. Your task is to continue the tree structure from the previous part to include the current part. - - The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc. - - For the title, you need to extract the original title from the text, only fix the space inconsistency. - - The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X. \ - - For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format. - - The response should be in the following format. - [ - { - "structure": <structure index, "x.x.x"> (string), - "title": <title of the section, keep the original title>, - "physical_index": "<physical_index_X> (keep the format)" - }, - ... - ] - - Directly return the additional part of the final JSON structure. Do not output anything else.""" + The structure variable is the numeric system which represents the index of the hierarchy section. + For the title, extract the original title from the text, only fix space inconsistency. + The provided text contains tags like <physical_index_X> to indicate the start and end of pages. + For the physical_index, extract the physical index of the start of the section. Keep the <physical_index_X> format. + Only output the additional part, not the previous structure.""" prompt = ( - _SYSTEM_HARDENING + prompt - + '\nGiven text\n:' + _secure_doc_text(part) - + '\nPrevious tree structure\n:' + _secure_doc_text(json.dumps(toc_content, indent=2)) + _SYSTEM_HARDENING + prompt + + '\nGiven text\n:' + _secure_doc_text(part) + + '\nPrevious tree structure\n:' + _secure_doc_text(json.dumps(toc_content, indent=2)) ) - - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - if finish_reason == 'finished': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') - + result = llm_structured(model=model, prompt=prompt, response_model=TOCStructureList) + return [item.model_dump() for item in result.items] + ### add verify completeness def generate_toc_init(part, model=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. - - The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc. - - For the title, you need to extract the original title from the text, only fix the space inconsistency. - - The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X. + You are an expert in extracting hierarchical tree structure. + Your task is to generate the tree structure of the document. - For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format. + The structure variable is the numeric system which represents the index of the hierarchy section. + For example, the first section has structure index 1, the first subsection has structure index 1.1, etc. - The response should be in the following format. - [ - {{ - "structure": <structure index, "x.x.x"> (string), - "title": <title of the section, keep the original title>, - "physical_index": "<physical_index_X> (keep the format)" - }}, - - ], + For the title, extract the original title from the text, only fix space inconsistency. + The provided text contains tags like <physical_index_1>, <physical_index_2> etc. to indicate the start and end of pages. + For the physical_index field, you MUST copy the exact tag from the text, for example: <physical_index_1> + Do NOT use A1, A2, page numbers, or any other format. ONLY use the exact <physical_index_X> tag as it appears in the text. - Directly return the final JSON structure. Do not output anything else.""" + Return your answer as a JSON object with a single key 'items' containing the list.""" prompt = _SYSTEM_HARDENING + prompt + '\nGiven text\n:' + _secure_doc_text(part) - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - - if finish_reason == 'finished': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') + result = llm_structured(model=model, prompt=prompt, response_model=TOCStructureList) + output = [item.model_dump() for item in result.items] + return output def process_no_toc(page_list, start_index=1, model=None, logger=None): page_contents=[] @@ -683,6 +563,7 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None): group_texts = page_list_to_group_text(page_contents, token_lengths) 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) toc_with_page_number = _validate_chunk_physical_indices( toc=toc_with_page_number, @@ -701,7 +582,7 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None): group_text, model ) - + toc_with_page_number_additional = _validate_chunk_physical_indices( toc=toc_with_page_number_additional, content=group_text @@ -712,7 +593,7 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None): total_pages=len(page_list), start_index=start_index ) - + toc_with_page_number.extend(toc_with_page_number_additional) logger.info(f'generate_toc: {toc_with_page_number}') @@ -749,23 +630,23 @@ def process_toc_no_page_numbers(toc_content, toc_page_list, page_list, start_in ): raise ValueError("LLM returned reordered or modified TOC entries.") valid_indices = _extract_chunk_marker_set(group_text) - + for idx, current in enumerate(toc_with_page_number): update = llm_result[idx] - + if current.get("physical_index") is not None: continue - + raw = update.get("physical_index") if raw is None: continue m = _PHYSICAL_INDEX_MARKER_RE.match(str(raw).strip()) - + if not m: continue if int(m.group(1)) not in valid_indices: continue - + current["physical_index"] = raw logger.info(f'add_page_number_to_toc: {toc_with_page_number}') @@ -896,30 +777,18 @@ def check_toc(page_list, opt=None): ################### fix incorrect toc ######################################################### async def single_toc_item_index_fixer(section_title, content, model=None): toc_extractor_prompt = """ - You are given a section title and several pages of a document, your job is to find the physical index of the start page of the section in the partial document. - - The provided pages contains tags like <physical_index_X> and <physical_index_X> to indicate the physical location of the page X. - - Reply in a JSON format: - { - "thinking": <explain which page, started and closed by <physical_index_X>, contains the start of this section>, - "physical_index": "<physical_index_X>" (keep the format) - } - Directly return the final JSON structure. Do not output anything else.""" + You are given a section title and several pages of a document. + Your job is to find the physical index of the start page of the section in the partial document. + The provided pages contains tags like <physical_index_X> to indicate the physical location of pages.""" prompt = ( - _SYSTEM_HARDENING + toc_extractor_prompt - + '\nSection Title:\n' + _secure_doc_text(str(section_title)) - + '\nDocument pages:\n' + _secure_doc_text(content) + _SYSTEM_HARDENING + toc_extractor_prompt + + '\nSection Title:\n' + _secure_doc_text(str(section_title)) + + '\nDocument pages:\n' + _secure_doc_text(content) ) - - response = await llm_acompletion(model=model, prompt=prompt) - 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) + result = await llm_astructured(model=model, prompt=prompt, response_model=PhysicalIndexResult) + return convert_physical_index_to_int(result.physical_index) async def fix_incorrect_toc(toc_with_page_number, page_list, incorrect_results, start_index=1, model=None, logger=None): diff --git a/pageindex/utils.py b/pageindex/utils.py index 235dd09cc..792228018 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -17,6 +17,10 @@ from pathlib import Path from types import SimpleNamespace as config import re +import instructor + +sync_instructor_client = instructor.from_litellm(litellm.completion, mode=instructor.Mode.JSON) +async_instructor_client = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON) # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): @@ -29,6 +33,30 @@ def count_tokens(text, model=None): return 0 return litellm.token_counter(model=model, text=text) +def llm_structured(model, prompt, response_model, chat_history=None): + if model: + model = model.removeprefix("litellm/") + messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] + return sync_instructor_client.chat.completions.create( + model=model, + messages=messages, + response_model=response_model, + temperature=0, + max_retries=3, + ) + +async def llm_astructured(model, prompt, response_model): + if model: + model = model.removeprefix("litellm/") + messages = [{"role": "user", "content": prompt}] + return await async_instructor_client.chat.completions.create( + model=model, + messages=messages, + response_model=response_model, + temperature=0, + max_retries=3, + ) + def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): if model: @@ -81,54 +109,8 @@ async def llm_acompletion(model, prompt): else: logging.error('Max retries reached for prompt: ' + prompt) return "" - - -def get_json_content(response): - start_idx = response.find("```json") - if start_idx != -1: - start_idx += 7 - response = response[start_idx:] - - end_idx = response.rfind("```") - if end_idx != -1: - response = response[:end_idx] - - json_content = response.strip() - return json_content - - -def extract_json(content): - try: - # First, try to extract JSON enclosed within ```json and ``` - start_idx = content.find("```json") - if start_idx != -1: - start_idx += 7 # Adjust index to start after the delimiter - end_idx = content.rfind("```") - json_content = content[start_idx:end_idx].strip() - else: - # If no delimiters, assume entire content could be JSON - json_content = content.strip() - - # Clean up common issues that might cause parsing errors - json_content = json_content.replace('None', 'null') # Replace Python None with JSON null - json_content = json_content.replace('\n', ' ').replace('\r', ' ') # Remove newlines - json_content = ' '.join(json_content.split()) # Normalize whitespace - - # Attempt to parse and return the JSON object - return json.loads(json_content) - except json.JSONDecodeError as e: - logging.error(f"Failed to extract JSON: {e}") - # Try to clean up the content further if initial parsing fails - try: - # Remove any trailing commas before closing brackets/braces - json_content = json_content.replace(',]', ']').replace(',}', '}') - return json.loads(json_content) - except: - logging.error("Failed to parse JSON even after cleanup") - return {} - except Exception as e: - logging.error(f"Unexpected error while extracting JSON: {e}") - return {} + + def write_node_id(data, node_id=0): if isinstance(data, dict): From 3e92917bd300dc2c9898334f596fdbd01c695c9d Mon Sep 17 00:00:00 2001 From: pegah <p.mansourian74@gmail.com> Date: Fri, 3 Jul 2026 14:17:36 -0400 Subject: [PATCH 2/4] refactor(json-parsing): replace manual JSON parsing with Instructor + Pydantic Introduce structured output via Instructor + LiteLLM + Pydantic across the PDF indexing pipeline, removing manual JSON parsing and its associated failure modes. - Extract all 14 response schemas into pageindex/schemas.py (previously inline in page_index.py); no changes to field descriptions, validation_alias configuration, or defaults during the move - Add AliasChoices("thinking", "reason", "explanation", "rationale") to reasoning fields so models using a synonym key still validate - Fix TOCIndexItem.structure: Optional[str] alone does not make a Pydantic v2 field omittable; add default=None so responses omitting the key validate instead of failing with "Field required" - Guard add_page_offset_to_toc_json against offset=None (raised when calculate_page_offset has no matching_pairs to compute from), replacing an opaque TypeError with a descriptive ValueError - Add instructor>=1.13.0,<2.0.0 to requirements.txt --- pageindex/page_index.py | 20 +++++++- pageindex/schemas.py | 100 ++++++++++++++++++++++++++++++++++++++++ pageindex/utils.py | 8 ++-- requirements.txt | 1 + 4 files changed, 125 insertions(+), 4 deletions(-) create mode 100644 pageindex/schemas.py diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 9c0dde20a..fa817c09c 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -7,6 +7,18 @@ from .utils import * import os from concurrent.futures import ThreadPoolExecutor, as_completed +from pageindex.schemas import ( + TitleAppearance, + TitleAppearanceInStart, + TOCDetection, + TOCCompletionCheck, + PageIndexDetection, + TOCTransformation, + TOCIndexList, + PageNumberList, + PhysicalIndexResult, + TOCStructureList, +) ######################### Hardening for prompt injection patterns #################################################### _INJECTION_PATTERNS = re.compile( @@ -350,7 +362,7 @@ def toc_transformer(toc_content, model=None): cleaned_response = convert_page_to_int(last_complete.get('table_of_contents', [])) return cleaned_response - + @@ -430,6 +442,12 @@ def calculate_page_offset(pairs): return most_common def add_page_offset_to_toc_json(data, offset): + if offset is None: + raise ValueError( + "Could not calculate page offset: no TOC entries could be matched " + "to a physical_index in the scanned front-matter window. " + "Check toc_index_extractor output and matching_pairs." + ) for i in range(len(data)): if data[i].get('page') is not None and isinstance(data[i]['page'], int): data[i]['physical_index'] = data[i]['page'] + offset diff --git a/pageindex/schemas.py b/pageindex/schemas.py new file mode 100644 index 000000000..2abd3c4cc --- /dev/null +++ b/pageindex/schemas.py @@ -0,0 +1,100 @@ +from pydantic import BaseModel, Field, AliasChoices +from typing import Literal, Optional, List, Union + +class TitleAppearance(BaseModel): + thinking: str = Field( + validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"), + description="why do you think the section appears or starts in the page_text" + ) + answer: Literal["yes", "no"] = Field(description="yes if the section appears or starts in the page_text, no otherwise") + +class TitleAppearanceInStart(BaseModel): + thinking: str = Field( + validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"), + description="why do you think the section appears or starts in the page_text" + ) + start_begin: Literal["yes", "no"] = Field(description="yes if the section starts in the beginning of the page_text, no otherwise") + +class TOCDetection(BaseModel): + thinking: str = Field( + validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"), + description="why do you think there is or isn't a table of content in the given text" + ) + toc_detected: Literal["yes", "no"] = Field(description="yes if a table of contents is detected, no otherwise") + +class TOCCompletionCheck(BaseModel): + thinking: str = Field( + validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"), + description="why do you think the table of contents is complete or not" + ) + completed: Literal["yes", "no"] = Field(description="yes if the table of contents is complete, no otherwise") + +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(description="yes if page numbers/indices are detected within the table of contents, no otherwise") + +class TOCItem(BaseModel): + structure: Optional[str] = Field( + default=None, + description="Numeric hierarchy index e.g. '1', '1.1', '1.2', or None" + ) + title: str = Field(description="Title of the section") + page: Optional[Union[int, str]] = Field(description="Page number from the table of contents, or None") + +class TOCTransformation(BaseModel): + table_of_contents: List[TOCItem] = Field(description="List of all sections extracted from the table of contents") + +class TOCIndexItem(BaseModel): + structure: Optional[str] = Field( + default=None, + description="Numeric hierarchy index e.g. '1', '1.1', or None" + ) + title: str = Field(description="Title of the section") + physical_index: Optional[str] = Field( + default=None, + description="Physical index tag in format <physical_index_X>, or None if section not found in provided pages", + pattern=r"^<physical_index_\d+>$" + ) + +class PageNumberItem(BaseModel): + structure: Optional[str] = Field( + default=None, + description="Numeric hierarchy index e.g. '1', '1.1', or None" + ) + title: str = Field(description="Title of the section") + start: Literal["yes", "no"] = Field(description="yes if the section starts in the current partial document, no otherwise") + physical_index: Optional[str] = Field( + default=None, + description="Physical index tag in format <physical_index_X> if start is yes, None otherwise", + pattern=r"^<physical_index_\d+>$" + ) + +class PhysicalIndexResult(BaseModel): + thinking: str = Field( + validation_alias=AliasChoices("thinking", "reason", "explanation", "rationale"), + description="explain which page, started and closed by <physical_index_X>, contains the start of this section" + ) + physical_index: str = Field(description="The physical index tag of the start page, in format <physical_index_X>", pattern=r"^<physical_index_\d+>$") + +class TOCStructureItem(BaseModel): + structure: str = Field( + default=None, + description="Numeric hierarchy index e.g. '1', '1.1', '1.2', or None" + ) + title: str = Field(description="Original title of the section, only fix space inconsistency") + physical_index: str = Field( + description="Physical index tag of the section start, in format <physical_index_X>", + pattern=r"^<physical_index_\d+>$" + ) + +class TOCStructureList(BaseModel): + items: List[TOCStructureItem] + +class TOCIndexList(BaseModel): + items: List[TOCIndexItem] + +class PageNumberList(BaseModel): + items: List[PageNumberItem] \ No newline at end of file diff --git a/pageindex/utils.py b/pageindex/utils.py index 792228018..0fac35e0a 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -11,16 +11,18 @@ import pymupdf from io import BytesIO from dotenv import load_dotenv -load_dotenv() import logging import yaml from pathlib import Path from types import SimpleNamespace as config import re import instructor +import re + +load_dotenv() -sync_instructor_client = instructor.from_litellm(litellm.completion, mode=instructor.Mode.JSON) -async_instructor_client = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON) +sync_instructor_client = instructor.from_litellm(litellm.completion, mode=instructor.Mode.MD_JSON) +async_instructor_client = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.MD_JSON) # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): diff --git a/requirements.txt b/requirements.txt index ae92bc49f..dafe07329 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ litellm==1.84.0 # openai-agents # optional: required for examples/agentic_vectorless_rag_demo.py +instructor>=1.13.0,<2.0.0 pymupdf==1.26.4 PyPDF2==3.0.1 python-dotenv==1.2.2 From ca6947de5b2a2547650ff760e0519a077830222f Mon Sep 17 00:00:00 2001 From: pegah <p.mansourian74@gmail.com> Date: Mon, 20 Jul 2026 09:59:23 -0400 Subject: [PATCH 3/4] test: add coverage for Instructor-based structured output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds tests/test_instructor_migration.py (23 tests) covering the reliability properties of the Instructor/Pydantic migration: - TOCIndexItem.structure regression test (Optional[str] without default=None is not omittable in Pydantic v2) - add_page_offset_to_toc_json guard against offset=None - llm_structured/llm_astructured: finish_reason truncation detection, max_tokens threading, max_retries capped at 1, litellm/ prefix stripping - toc_transformer end-to-end behavior against the migrated implementation - AliasChoices synonym handling for reasoning fields Also includes the llm_structured/llm_astructured implementation changes these tests cover: switched from create() to create_with_completion() to expose finish_reason, added max_tokens parameter, capped max_retries at 1 (previously 3 — found to compound failures on smaller models by feeding accumulated failed completions back into retry context). tests/test_issue_163.py is left unmodified. Running it against this branch shows 10 failed / 4 passed: the 4 passes (extract_toc_content) are unaffected since that function wasn't migrated; the 10 failures mock llm_completion for functions now using llm_structured, so the mock no longer intercepts the real call path. Not a regression — a consequence of those tests asserting on manual-parsing implementation details this PR removes. Left as-is pending maintainer input on how to handle it. --- pageindex/page_index.py | 87 +++------ pageindex/utils.py | 36 +++- tests/test_instructor_migration.py | 283 +++++++++++++++++++++++++++++ 3 files changed, 335 insertions(+), 71 deletions(-) create mode 100644 tests/test_instructor_migration.py diff --git a/pageindex/page_index.py b/pageindex/page_index.py index fa817c09c..b7f79e783 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -88,7 +88,7 @@ def _validate_physical_indices(toc: list, total_pages: int, start_index: int = 1 return toc ################### check title in page ######################################################### -async def check_title_appearance(item, page_list, start_index=1, model=None): +async def check_title_appearance(item, page_list, start_index=1, model=None): title=item['title'] if 'physical_index' not in item or item['physical_index'] is None: return {'list_index': item.get('list_index'), 'answer': 'no', 'title': title, 'page_number': None} @@ -96,7 +96,7 @@ async def check_title_appearance(item, page_list, start_index=1, model=None): page_number = item['physical_index'] page_text = page_list[page_number-start_index][0] - + prompt = _SYSTEM_HARDENING + f""" Your job is to check if the given section appears or starts in the given page_text. @@ -108,7 +108,7 @@ async def check_title_appearance(item, page_list, start_index=1, model=None): return {'list_index': item['list_index'], 'answer': result.answer, 'title': title, 'page_number': page_number} -async def check_title_appearance_in_start(title, page_text, model=None, logger=None): +async def check_title_appearance_in_start(title, page_text, model=None, logger=None): prompt = _SYSTEM_HARDENING + f""" You will be given the current section title and the current page_text. Your job is to check if the current section starts in the beginning of the given page_text. @@ -210,7 +210,7 @@ def extract_toc_content(content, model=None): {"role": "assistant", "content": response}, ] continue_prompt = "please continue the generation of table of contents, directly output the remaining part of the structure" - + max_attempts = 5 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) @@ -222,7 +222,7 @@ def extract_toc_content(content, model=None): break else: raise Exception('Failed to complete table of contents extraction after maximum retries') - + return response def detect_page_index(toc_content, model=None): @@ -298,73 +298,31 @@ def toc_index_extractor(toc, content, model=None): items = [item.model_dump() for item in result.items] return _validate_chunk_physical_indices(toc=items, content=content) - - def toc_transformer(toc_content, model=None): + # TODO after rebase completes — two known gaps vs. main's version, to + # resolve deliberately (see PR discussion), not silently: + # 1. No continuation/retry logic for truncated completions (main uses + # a chat-history "please continue" loop, capped at 5 attempts). + # llm_structured already fails fast on finish_reason == "length", + # which is a deliberate trade-off, not an oversight. + # 2. No equivalent of check_if_toc_transformation_is_complete — main + # catches "complete-looking but actually missing sections" output + # (finish_reason == "stop" but content silently incomplete), which + # Pydantic validation alone cannot catch (a shorter-than-expected + # list still validates fine). This gap needs an explicit decision: + # port the check, or document and accept it. print('start toc_transformer') init_prompt = """ You are given a table of contents, You job is to transform the whole table of content into a JSON format included table_of_contents. structure is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc. - The response should be in the following JSON format: - { - table_of_contents: [ - { - "structure": <structure index, "x.x.x" or None> (string), - "title": <title of the section>, - "page": <page number or None>, - }, - ... - ], - } You should transform the full table of contents in one go. Directly return the final JSON structure, do not output anything else. """ prompt = init_prompt + '\n Given table of contents\n:' + _secure_doc_text(toc_content) - last_complete, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - 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.get('table_of_contents', [])) - return cleaned_response - - last_complete = get_json_content(last_complete) - 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 - for attempt in range(max_attempts): - - 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 - - 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.get('table_of_contents', [])) - return cleaned_response - - - + result = llm_structured(model=model, prompt=prompt, response_model=TOCTransformation, max_tokens=8000) + return convert_page_to_int([item.model_dump() for item in result.table_of_contents]) def find_toc_pages(start_page_index, page_list, opt, logger=None): print('start find_toc_pages') @@ -572,16 +530,15 @@ def generate_toc_init(part, model=None): return output def process_no_toc(page_list, start_index=1, model=None, logger=None): - page_contents=[] - token_lengths=[] - for page_index in range(start_index, start_index+len(page_list)): + 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) 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) toc_with_page_number = _validate_chunk_physical_indices( toc=toc_with_page_number, diff --git a/pageindex/utils.py b/pageindex/utils.py index 0fac35e0a..114152711 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -35,30 +35,54 @@ def count_tokens(text, model=None): return 0 return litellm.token_counter(model=model, text=text) -def llm_structured(model, prompt, response_model, chat_history=None): +def llm_structured(model, prompt, response_model, chat_history=None, max_tokens=4000): if model: model = model.removeprefix("litellm/") messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] - return sync_instructor_client.chat.completions.create( + + result, completion = sync_instructor_client.chat.completions.create_with_completion( model=model, messages=messages, response_model=response_model, temperature=0, - max_retries=3, + max_retries=1, + max_tokens=max_tokens, ) -async def llm_astructured(model, prompt, response_model): + finish_reason = completion.choices[0].finish_reason + if finish_reason == "length": + raise ValueError( + f"Response was truncated (finish_reason='length') before completing " + f"the {response_model.__name__} structure. Increase max_tokens (currently " + f"{max_tokens}) for this call." + ) + + return result + +async def llm_astructured(model, prompt, response_model, max_tokens=4000): if model: model = model.removeprefix("litellm/") messages = [{"role": "user", "content": prompt}] - return await async_instructor_client.chat.completions.create( + + result, completion = await async_instructor_client.chat.completions.create_with_completion( model=model, messages=messages, response_model=response_model, temperature=0, - max_retries=3, + max_retries=1, + max_tokens=max_tokens, ) + finish_reason = completion.choices[0].finish_reason + if finish_reason == "length": + raise ValueError( + f"Response was truncated (finish_reason='length') before completing " + f"the {response_model.__name__} structure. Increase max_tokens (currently " + f"{max_tokens}) for this call." + ) + + return result + def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): if model: diff --git a/tests/test_instructor_migration.py b/tests/test_instructor_migration.py new file mode 100644 index 000000000..1dda738b3 --- /dev/null +++ b/tests/test_instructor_migration.py @@ -0,0 +1,283 @@ +""" +Tests for the Instructor + Pydantic structured-output refactor. + +These replace tests/test_issue_163.py, which asserted implementation details of +the old manual-JSON-parsing code (dict .get() fallbacks, chat-history-growing +continuation loops, retry counts tied to "max_output_reached" finish_reason +strings). That code path no longer exists after this refactor — validation, +retries, and truncation handling are now owned by Instructor/Pydantic — so +those tests were superseded rather than adapted. + +Adjust import paths below to match your actual module layout if they differ. +""" + +import pytest +from unittest.mock import patch, MagicMock +from pydantic import BaseModel, ValidationError +from typing import Optional, List + +from pageindex.schemas import ( + TOCIndexItem, + TOCIndexList, + TOCTransformation, + TOCItem, +) +from pageindex.page_index import ( + toc_transformer, + add_page_offset_to_toc_json, +) +from pageindex.utils import llm_structured + + +# --------------------------------------------------------------------------- +# Schema-level tests — catch the exact bug class this PR fixed +# (Optional[str] without default=None is NOT omittable in Pydantic v2) +# --------------------------------------------------------------------------- + +class TestSchemaOptionalFields: + def test_toc_index_item_structure_is_omittable(self): + """ + Regression test for the bug found during testing: TOCIndexItem.structure + was declared Optional[str] with no default, which does NOT make the + field omittable in Pydantic v2 — any model response omitting the key + failed with "Field required" regardless of which LLM produced it. + """ + item = TOCIndexItem(title="Some Section") # structure omitted entirely + assert item.structure is None + assert item.physical_index is None + + def test_toc_index_item_accepts_explicit_none(self): + item = TOCIndexItem(title="Some Section", structure=None, physical_index=None) + assert item.structure is None + + def test_toc_index_item_requires_title(self): + with pytest.raises(ValidationError): + TOCIndexItem(structure="1.1") # title is required, correctly + + def test_toc_index_item_physical_index_pattern_enforced(self): + with pytest.raises(ValidationError): + TOCIndexItem(title="X", physical_index="page 5") # wrong format + + item = TOCIndexItem(title="X", physical_index="<physical_index_5>") + assert item.physical_index == "<physical_index_5>" + + def test_toc_index_list_requires_items_key(self): + with pytest.raises(ValidationError): + TOCIndexList() # items is required — this is correct/intended, + # unlike `structure` above which should be optional + + +# --------------------------------------------------------------------------- +# add_page_offset_to_toc_json — regression test for the None-offset crash +# --------------------------------------------------------------------------- + +class TestAddPageOffsetToTocJson: + def test_raises_clear_error_when_offset_is_none(self): + """ + Regression test: calculate_page_offset can return None when zero TOC + entries were matched to a physical_index. The unguarded code evaluated + int + None, raising an opaque TypeError. This should now raise a + descriptive ValueError instead. + """ + toc_data = [{"title": "Intro", "page": 1, "structure": "1"}] + with pytest.raises(ValueError, match="page offset"): + add_page_offset_to_toc_json(toc_data, offset=None) + + def test_applies_offset_correctly_when_valid(self): + toc_data = [ + {"title": "Intro", "page": 1, "structure": "1"}, + {"title": "Conclusion", "page": 10, "structure": "2"}, + ] + result = add_page_offset_to_toc_json(toc_data, offset=5) + assert result[0]["physical_index"] == 6 + assert result[1]["physical_index"] == 15 + assert "page" not in result[0] + + def test_skips_items_with_no_page_number(self): + toc_data = [{"title": "Part I", "page": None, "structure": "1"}] + result = add_page_offset_to_toc_json(toc_data, offset=5) + # item without an int page is left untouched, not crashed on + assert "physical_index" not in result[0] or result[0].get("physical_index") is None + + +# --------------------------------------------------------------------------- +# llm_structured — finish_reason / truncation handling +# --------------------------------------------------------------------------- + +class DummyResponseModel(BaseModel): + value: str + + +def _make_mock_completion(finish_reason="stop"): + mock_choice = MagicMock() + mock_choice.finish_reason = finish_reason + mock_completion = MagicMock() + mock_completion.choices = [mock_choice] + return mock_completion + + +class TestLLMStructuredTruncationHandling: + @patch("pageindex.utils.sync_instructor_client") + def test_returns_result_on_normal_completion(self, mock_client): + expected = DummyResponseModel(value="ok") + mock_client.chat.completions.create_with_completion.return_value = ( + expected, + _make_mock_completion(finish_reason="stop"), + ) + result = llm_structured( + model="test-model", prompt="hello", response_model=DummyResponseModel + ) + assert result == expected + + @patch("pageindex.utils.sync_instructor_client") + def test_raises_on_truncated_completion(self, mock_client): + """ + Regression test: previously llm_structured had no visibility into + finish_reason at all (create() discards completion metadata), so a + truncated response would either fail Pydantic validation with a + confusing error, or in principle validate against partial/incomplete + data. This should now fail explicitly and legibly. + """ + mock_client.chat.completions.create_with_completion.return_value = ( + DummyResponseModel(value="partial"), + _make_mock_completion(finish_reason="length"), + ) + with pytest.raises(ValueError, match="truncated"): + llm_structured( + model="test-model", prompt="hello", response_model=DummyResponseModel + ) + + @patch("pageindex.utils.sync_instructor_client") + def test_passes_max_tokens_through(self, mock_client): + mock_client.chat.completions.create_with_completion.return_value = ( + DummyResponseModel(value="ok"), + _make_mock_completion(finish_reason="stop"), + ) + llm_structured( + model="test-model", + prompt="hello", + response_model=DummyResponseModel, + max_tokens=9000, + ) + _, kwargs = mock_client.chat.completions.create_with_completion.call_args + assert kwargs["max_tokens"] == 9000 + + @patch("pageindex.utils.sync_instructor_client") + def test_uses_default_max_tokens_when_not_specified(self, mock_client): + mock_client.chat.completions.create_with_completion.return_value = ( + DummyResponseModel(value="ok"), + _make_mock_completion(finish_reason="stop"), + ) + llm_structured(model="test-model", prompt="hello", response_model=DummyResponseModel) + _, kwargs = mock_client.chat.completions.create_with_completion.call_args + assert kwargs["max_tokens"] == 4000 # current default — update if you change it + + @patch("pageindex.utils.sync_instructor_client") + def test_max_retries_is_capped_at_one(self, mock_client): + """ + Regression test: max_retries was found to actively degrade output + quality on smaller/local models by feeding failed completions + raw + Pydantic errors back into context across retries. Confirms it's not + silently reverted back to a higher value. + """ + mock_client.chat.completions.create_with_completion.return_value = ( + DummyResponseModel(value="ok"), + _make_mock_completion(finish_reason="stop"), + ) + llm_structured(model="test-model", prompt="hello", response_model=DummyResponseModel) + _, kwargs = mock_client.chat.completions.create_with_completion.call_args + assert kwargs["max_retries"] == 1 + + @patch("pageindex.utils.sync_instructor_client") + def test_strips_litellm_prefix_from_model_name(self, mock_client): + mock_client.chat.completions.create_with_completion.return_value = ( + DummyResponseModel(value="ok"), + _make_mock_completion(finish_reason="stop"), + ) + llm_structured( + model="litellm/gpt-4o", prompt="hello", response_model=DummyResponseModel + ) + args, kwargs = mock_client.chat.completions.create_with_completion.call_args + assert kwargs["model"] == "gpt-4o" + + +# --------------------------------------------------------------------------- +# toc_transformer — end-to-end with mocked Instructor client +# --------------------------------------------------------------------------- + +class TestTocTransformer: + @patch("pageindex.page_index.llm_structured") + def test_transforms_toc_content_into_structured_list(self, mock_llm_structured): + mock_result = TOCTransformation( + table_of_contents=[ + TOCItem(structure="1", title="Introduction", page=1), + TOCItem(structure="2", title="Conclusion", page=10), + ] + ) + mock_llm_structured.return_value = mock_result + + result = toc_transformer("1. Introduction ... 1\n2. Conclusion ... 10", model="test") + + assert len(result) == 2 + assert result[0]["title"] == "Introduction" + assert result[0]["page"] == 1 + assert result[1]["title"] == "Conclusion" + + @patch("pageindex.page_index.llm_structured") + def test_calls_llm_structured_with_generous_max_tokens(self, mock_llm_structured): + """ + toc_transformer intentionally uses a higher max_tokens than the + llm_structured default, since TOCs can be long. This guards against + that being silently dropped in a future edit. + """ + mock_llm_structured.return_value = TOCTransformation(table_of_contents=[]) + toc_transformer("some toc", model="test") + _, kwargs = mock_llm_structured.call_args + assert kwargs.get("max_tokens", 0) >= 8000 + + @patch("pageindex.page_index.llm_structured") + def test_empty_toc_returns_empty_list(self, mock_llm_structured): + mock_llm_structured.return_value = TOCTransformation(table_of_contents=[]) + result = toc_transformer("no sections here", model="test") + assert result == [] + + @patch("pageindex.page_index.llm_structured") + def test_propagates_truncation_error(self, mock_llm_structured): + """ + If the underlying llm_structured call raises due to truncation + (finish_reason == 'length'), toc_transformer should not swallow it — + this documents the deliberate behavior change from main's old + continuation-loop approach (see PR discussion). + """ + mock_llm_structured.side_effect = ValueError("Response was truncated...") + with pytest.raises(ValueError, match="truncated"): + toc_transformer("a very long toc", model="test") + + +# --------------------------------------------------------------------------- +# AliasChoices — confirm reasoning-field synonym handling actually works +# --------------------------------------------------------------------------- + +class TestReasoningFieldAliases: + """ + Some models (observed with local/smaller models during testing) use a + synonym key ('reason', 'explanation', 'rationale') instead of 'thinking' + for free-text reasoning fields. AliasChoices was added so these still + validate instead of failing with "Field required". + """ + + def test_accepts_canonical_thinking_key(self): + from pageindex.schemas import TOCDetection + obj = TOCDetection.model_validate({"thinking": "because X", "toc_detected": "yes"}) + assert obj.thinking == "because X" + + @pytest.mark.parametrize("alias_key", ["reason", "explanation", "rationale"]) + def test_accepts_synonym_keys(self, alias_key): + from pageindex.schemas import TOCDetection + obj = TOCDetection.model_validate({alias_key: "because X", "toc_detected": "yes"}) + assert obj.thinking == "because X" + + def test_rejects_missing_reasoning_field_entirely(self): + from pageindex.schemas import TOCDetection + with pytest.raises(ValidationError): + TOCDetection.model_validate({"toc_detected": "yes"}) \ No newline at end of file From 4bb25b581f04ca4a96ad2beab1d3e18fddf5a262 Mon Sep 17 00:00:00 2001 From: pegah <p.mansourian74@gmail.com> Date: Mon, 20 Jul 2026 11:43:17 -0400 Subject: [PATCH 4/4] docs: clarify toc_transformer TODO to reflect current state (no retry logic yet) --- pageindex/page_index.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/pageindex/page_index.py b/pageindex/page_index.py index b7f79e783..5f5f65208 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -299,18 +299,23 @@ def toc_index_extractor(toc, content, model=None): return _validate_chunk_physical_indices(toc=items, content=content) def toc_transformer(toc_content, model=None): - # TODO after rebase completes — two known gaps vs. main's version, to - # resolve deliberately (see PR discussion), not silently: - # 1. No continuation/retry logic for truncated completions (main uses - # a chat-history "please continue" loop, capped at 5 attempts). - # llm_structured already fails fast on finish_reason == "length", - # which is a deliberate trade-off, not an oversight. - # 2. No equivalent of check_if_toc_transformation_is_complete — main - # catches "complete-looking but actually missing sections" output - # (finish_reason == "stop" but content silently incomplete), which - # Pydantic validation alone cannot catch (a shorter-than-expected - # list still validates fine). This gap needs an explicit decision: - # port the check, or document and accept it. + # TODO — two known gaps vs. main's version, tracked for follow-up + # (see PR discussion), not silent omissions: + # 1. No continuation/retry for truncated completions. Main uses a + # chat-history "please continue" loop, capped at 5 attempts. + # llm_structured fails fast instead on finish_reason == "length" — + # a deliberate trade-off (Instructor's structured output doesn't + # compose cleanly with resuming a partial JSON response). + # 2. No completeness check yet. Main's check_if_toc_transformation_is_ + # complete catches a case schema validation can't: a response that + # validates fine but silently skipped sections (finish_reason == + # "stop", not truncated, just incomplete). Planned: reuse that + # function unmodified via a Pydantic model_validator on + # TOCTransformation, using Instructor's context= passthrough so + # the validator can access toc_content — this way the completeness + # check participates in Instructor's own retry loop instead of a + # separate hand-rolled one. Not yet implemented; requires + # llm_structured to accept/forward a context param. print('start toc_transformer') init_prompt = """ You are given a table of contents, You job is to transform the whole table of content into a JSON format included table_of_contents.