diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 9004309fb..4d7c0b4da 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -7,68 +7,100 @@ from .utils import * import os from concurrent.futures import ThreadPoolExecutor, as_completed +from pydantic import BaseModel +from typing import Literal, Optional, List + +class TitleAppearance(BaseModel): + thinking: str + answer: Literal["yes", "no"] + +class TitleAppearanceInStart(BaseModel): + thinking: str + start_begin: Literal["yes", "no"] + +class TOCDetection(BaseModel): + thinking: str + toc_detected: Literal["yes", "no"] + +class TOCCompletionCheck(BaseModel): + thinking: str + completed: Literal["yes", "no"] + +class PageIndexDetection(BaseModel): + thinking: str + page_index_given_in_toc: Literal["yes", "no"] + +class TOCItem(BaseModel): + structure: Optional[str] + title: str + page: Optional[int] + +class TOCTransformation(BaseModel): + table_of_contents: List[TOCItem] + +class TOCIndexItem(BaseModel): + structure: Optional[str] + title: str + physical_index: Optional[str] + +class PageNumberItem(BaseModel): + structure: Optional[str] + title: str + start: Literal["yes", "no"] + physical_index: Optional[str] + +class PhysicalIndexResult(BaseModel): + thinking: str + physical_index: str + +class TOCStructureItem(BaseModel): + structure: str + title: str + physical_index: str + +class TOCStructureList(BaseModel): + items: List[TOCStructureItem] + +class TOCIndexList(BaseModel): + items: List[TOCIndexItem] + +class PageNumberList(BaseModel): + items: List[PageNumberItem] ################### check title in page ######################################################### -async def check_title_appearance(item, page_list, start_index=1, model=None): - title=item['title'] +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] + page_text = page_list[page_number - start_index][0] - prompt = f""" 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 {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 {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): +async def check_title_appearance_in_start(title, page_text, model=None, logger=None): prompt = 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. If there are other contents before the current section title, then the current section does not start in the beginning of the given page_text. If the current section title is the first content in the given page_text, then the current section starts in the beginning of 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 {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 {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): @@ -104,58 +136,38 @@ async def check_title_appearance_in_start_concurrent(structure, page_list, model def toc_detector_single_page(content, model=None): prompt = f""" Your job is to detect if there is a table of content provided in the given text. - Given text: {content} + Please note: abstract, summary, notation list, figure list, table list, etc. are not table of contents.""" - 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) - # print('response', response) - json_content = extract_json(response) - return json_content['toc_detected'] + 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.""" + 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. + Document: + {content} + Table of contents: + {toc}""" - 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'] + 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. + 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: + {content} + Cleaned Table of contents: + {toc}""" - Reply format: - {{ - "thinking": - "completed": "yes" or "no" - }} - Directly return the final JSON structure. Do not output anything else.""" - - 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'] + result = llm_structured(model=model, prompt=prompt, response_model=TOCCompletionCheck) + return result.completed def extract_toc_content(content, model=None): prompt = f""" @@ -203,21 +215,11 @@ def detect_page_index(toc_content, model=None): print('start detect_page_index') prompt = f""" You will be given a table of contents. - 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['page_index_given_in_toc'] + 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): @@ -242,98 +244,35 @@ def transform_dots_to_colon(text): 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. - + prompt = f""" + 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.""" + Table of contents: + {toc} + Document pages: + {content}""" - prompt = toc_extractor_prompt + '\nTable of contents:\n' + str(toc) + '\nDocument pages:\n' + content - response = llm_completion(model=model, prompt=prompt) - json_content = extract_json(response) - return json_content + result = llm_structured(model=model, prompt=prompt, response_model=TOCIndexList) + return [item.model_dump() for item in result.items] def toc_transformer(toc_content, model=None): 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>, - }, - ... - ], - } + prompt = f""" + You are given a table of contents. Your job is to transform the whole table of contents into a JSON format. + structure is the numeric system which represents the index of the hierarchy section. 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:' + 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['table_of_contents']) - return cleaned_response - - last_complete = get_json_content(last_complete) - 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 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: - - The raw table of contents json structure is: - {toc_content} - - The incomplete transformed table of contents json structure is: - {last_complete} + Given table of contents: + {toc_content}""" - Please continue the json structure, directly output the remaining part of the json structure.""" + result = llm_structured(model=model, prompt=prompt, response_model=TOCTransformation) + return convert_page_to_int([item.model_dump() for item in result.table_of_contents]) - new_complete, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - - if new_complete.startswith('```json'): - new_complete = get_json_content(new_complete) - last_complete = last_complete+new_complete - - if_complete = check_if_toc_transformation_is_complete(toc_content, last_complete, model) - - - last_complete = extract_json(last_complete) - - cleaned_response=convert_page_to_int(last_complete['table_of_contents']) - return cleaned_response @@ -459,35 +398,22 @@ def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, over return subsets 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.""" - - 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) - + prompt = f""" + 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. + Current Partial Document: + {part} + Given Structure: + {json.dumps(structure, indent=2)}""" + + 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 @@ -506,86 +432,59 @@ def remove_first_physical_index_section(text): ### add verify completeness def generate_toc_continue(toc_content, part, model=None): print('start generate_toc_continue') - prompt = """ + prompt = f""" 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.""" - - 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': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') + 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. + Given text: + {part} + Previous tree structure: + {json.dumps(toc_content, indent=2)}""" + + 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. - - 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. + prompt = f""" + You are an expert in extracting hierarchical tree structure. + Your task is to generate the tree structure of the document. - 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)" - }}, - - ], + 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. + For the title, extract the original title from the text, only fix space inconsistency. - Directly return the final JSON structure. Do not output anything else.""" + 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. - prompt = prompt + '\nGiven text\n:' + part - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) + Return your answer as a JSON object with a single key 'items' containing the list. - if finish_reason == 'finished': - return extract_json(response) - else: - raise Exception(f'finish reason: {finish_reason}') + Given text: + {part}""" + 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=[] - 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) 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) toc_with_page_number.extend(toc_with_page_number_additional) logger.info(f'generate_toc: {toc_with_page_number}') @@ -738,22 +637,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. + prompt = f""" + 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. + Section Title: + {section_title} + Document pages: + {content}""" - 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.""" + result = await llm_astructured(model=model, prompt=prompt, response_model=PhysicalIndexResult) + return convert_physical_index_to_int(result.physical_index) - 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']) diff --git a/pageindex/utils.py b/pageindex/utils.py index f00ccf3a7..85d93649d 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -16,6 +16,10 @@ import yaml from pathlib import Path from types import SimpleNamespace as config +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"): @@ -28,6 +32,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,53 +109,7 @@ async def llm_acompletion(model, prompt): 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):