diff --git a/.gitignore b/.gitignore index 9311f58..4825407 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,8 @@ dist/ *.db venv/ uv.lock + +# local SDK test-run artifacts (generated by demos; keep tracked example json) +examples/workspace/files/ +examples/workspace/*.db +examples/documents/attention.pdf diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 802befa..166c33a 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,7 +1,9 @@ # pageindex/__init__.py -# Upstream exports (backward compatibility) -from .page_index import * -from .page_index_md import md_to_tree +# Upstream exports (backward compatibility). Import from the canonical +# pageindex.index.* modules directly so `import pageindex` does NOT trip the +# top-level deprecation shims (pageindex.page_index / .page_index_md / .utils). +from .index.page_index import * +from .index.page_index_md import md_to_tree from .retrieve import get_document, get_document_structure, get_page_content # SDK exports diff --git a/pageindex/index/legacy_utils.py b/pageindex/index/legacy_utils.py deleted file mode 100644 index 1d6aab5..0000000 --- a/pageindex/index/legacy_utils.py +++ /dev/null @@ -1,2 +0,0 @@ -# Re-export from the original utils.py for backward compatibility -from ..utils import * diff --git a/pageindex/index/page_index.py b/pageindex/index/page_index.py index 2913090..6eece2d 100644 --- a/pageindex/index/page_index.py +++ b/pageindex/index/page_index.py @@ -4,7 +4,7 @@ import copy import math import random import re -from .legacy_utils import * +from .utils import * import os from concurrent.futures import ThreadPoolExecutor, as_completed diff --git a/pageindex/index/page_index_md.py b/pageindex/index/page_index_md.py index e6078c2..f9e300a 100644 --- a/pageindex/index/page_index_md.py +++ b/pageindex/index/page_index_md.py @@ -2,10 +2,7 @@ import asyncio import json import re import os -try: - from .legacy_utils import * -except: - from legacy_utils import * +from .utils import * async def get_node_summary(node, summary_token_threshold=200, model=None): node_text = node.get('text') diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index 5e9700b..49e3878 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -1,11 +1,20 @@ import litellm import logging +import os +import textwrap import time import json import copy import re import asyncio import PyPDF2 +import pymupdf +import yaml +from datetime import datetime +from io import BytesIO +from pathlib import Path +from pprint import pprint +from types import SimpleNamespace as config from ..config import get_llm_params @@ -132,13 +141,15 @@ def write_node_id(data, node_id=0): return node_id -def remove_fields(data, fields=None): +def remove_fields(data, fields=None, max_len=None): fields = fields or ["text"] if isinstance(data, dict): - return {k: remove_fields(v, fields) + return {k: remove_fields(v, fields, max_len) for k, v in data.items() if k not in fields} elif isinstance(data, list): - return [remove_fields(item, fields) for item in data] + return [remove_fields(item, fields, max_len) for item in data] + elif isinstance(data, str): + return data[:max_len] + '...' if max_len is not None and len(data) > max_len else data return data @@ -174,7 +185,9 @@ def get_nodes(structure): def get_leaf_nodes(structure): if isinstance(structure, dict): - if not structure['nodes']: + # .get() — clean_node deletes the 'nodes' key on leaf nodes, so direct + # indexing raises KeyError on a standard tree (issue #330 / #331). + if not structure.get('nodes'): structure_node = copy.deepcopy(structure) structure_node.pop('nodes', None) return [structure_node] @@ -431,3 +444,420 @@ def get_md_page_content(structure: list, page_nums: list[int]) -> list[dict]: _traverse(structure) results.sort(key=lambda x: x['page']) return results + + + +# ───────────────────────────────────────────────────────────────────── +# Legacy 0.2.x / OSS utility API — kept here so this module is the single +# source of truth for the indexing pipeline. Previously duplicated in the +# top-level pageindex/utils.py (now a deprecation shim re-exporting this). +# ───────────────────────────────────────────────────────────────────── + +async def call_llm(prompt, api_key, model="gpt-4.1", temperature=0): + """Call an LLM to generate a response to a prompt. + + Kept for compatibility with the pageindex 0.2.x SDK utility API. + """ + import openai + + client = openai.AsyncOpenAI(api_key=api_key) + response = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + temperature=temperature, + ) + return response.choices[0].message.content.strip() + + +def is_leaf_node(data, node_id): + # Helper function to find the node by its node_id + def find_node(data, node_id): + if isinstance(data, dict): + if data.get('node_id') == node_id: + return data + for key in data.keys(): + if 'nodes' in key: + result = find_node(data[key], node_id) + if result: + return result + elif isinstance(data, list): + for item in data: + result = find_node(item, node_id) + if result: + return result + return None + + # Find the node with the given node_id + node = find_node(data, node_id) + + # Check if the node is a leaf node + if node and not node.get('nodes'): + return True + return False + + +def get_last_node(structure): + return structure[-1] + + +def extract_text_from_pdf(pdf_path): + pdf_reader = PyPDF2.PdfReader(pdf_path) + ###return text not list + text="" + for page_num in range(len(pdf_reader.pages)): + page = pdf_reader.pages[page_num] + text+=page.extract_text() + return text + + +def get_pdf_title(pdf_path): + pdf_reader = PyPDF2.PdfReader(pdf_path) + meta = pdf_reader.metadata + title = meta.title if meta and meta.title else 'Untitled' + return title + + +def get_text_of_pages(pdf_path, start_page, end_page, tag=True): + pdf_reader = PyPDF2.PdfReader(pdf_path) + text = "" + for page_num in range(start_page-1, end_page): + page = pdf_reader.pages[page_num] + page_text = page.extract_text() + if tag: + text += f"\n{page_text}\n\n" + else: + text += page_text + return text + + +def get_first_start_page_from_text(text): + start_page = -1 + start_page_match = re.search(r'', text) + if start_page_match: + start_page = int(start_page_match.group(1)) + return start_page + + +def get_last_start_page_from_text(text): + start_page = -1 + # Find all matches of start_index tags + start_page_matches = re.finditer(r'', text) + # Convert iterator to list and get the last match if any exist + matches_list = list(start_page_matches) + if matches_list: + start_page = int(matches_list[-1].group(1)) + return start_page + + +def sanitize_filename(filename, replacement='-'): + # In Linux, only '/' and '\0' (null) are invalid in filenames. + # Null can't be represented in strings, so we only handle '/'. + return filename.replace('/', replacement) + + +def get_pdf_name(pdf_path): + # Extract PDF name + if isinstance(pdf_path, str): + pdf_name = os.path.basename(pdf_path) + elif isinstance(pdf_path, BytesIO): + pdf_reader = PyPDF2.PdfReader(pdf_path) + meta = pdf_reader.metadata + pdf_name = meta.title if meta and meta.title else 'Untitled' + pdf_name = sanitize_filename(pdf_name) + return pdf_name + + +class JsonLogger: + def __init__(self, file_path): + # Extract PDF name for logger name + pdf_name = get_pdf_name(file_path) + + current_time = datetime.now().strftime("%Y%m%d_%H%M%S") + self.filename = f"{pdf_name}_{current_time}.json" + os.makedirs("./logs", exist_ok=True) + # Initialize empty list to store all messages + self.log_data = [] + + def log(self, level, message, **kwargs): + if isinstance(message, dict): + self.log_data.append(message) + else: + self.log_data.append({'message': message}) + # Add new message to the log data + + # Write entire log data to file + with open(self._filepath(), "w") as f: + json.dump(self.log_data, f, indent=2) + + def info(self, message, **kwargs): + self.log("INFO", message, **kwargs) + + def error(self, message, **kwargs): + self.log("ERROR", message, **kwargs) + + def debug(self, message, **kwargs): + self.log("DEBUG", message, **kwargs) + + def exception(self, message, **kwargs): + kwargs["exception"] = True + self.log("ERROR", message, **kwargs) + + def _filepath(self): + return os.path.join("logs", self.filename) + + +def add_preface_if_needed(data): + if not isinstance(data, list) or not data: + return data + + if data[0]['physical_index'] is not None and data[0]['physical_index'] > 1: + preface_node = { + "structure": "0", + "title": "Preface", + "physical_index": 1, + } + data.insert(0, preface_node) + return data + + +def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): + if pdf_parser == "PyPDF2": + pdf_reader = PyPDF2.PdfReader(pdf_path) + page_list = [] + for page_num in range(len(pdf_reader.pages)): + page = pdf_reader.pages[page_num] + page_text = page.extract_text() + token_length = litellm.token_counter(model=model, text=page_text) + page_list.append((page_text, token_length)) + return page_list + elif pdf_parser == "PyMuPDF": + if isinstance(pdf_path, BytesIO): + pdf_stream = pdf_path + doc = pymupdf.open(stream=pdf_stream, filetype="pdf") + elif isinstance(pdf_path, str) and os.path.isfile(pdf_path) and pdf_path.lower().endswith(".pdf"): + doc = pymupdf.open(pdf_path) + page_list = [] + for page in doc: + page_text = page.get_text() + token_length = litellm.token_counter(model=model, text=page_text) + page_list.append((page_text, token_length)) + return page_list + else: + raise ValueError(f"Unsupported PDF parser: {pdf_parser}") + + +def get_text_of_pdf_pages(pdf_pages, start_page, end_page): + text = "" + for page_num in range(start_page-1, end_page): + text += pdf_pages[page_num][0] + return text + + +def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page): + text = "" + for page_num in range(start_page-1, end_page): + text += f"\n{pdf_pages[page_num][0]}\n\n" + return text + + +def get_number_of_pages(pdf_path): + pdf_reader = PyPDF2.PdfReader(pdf_path) + num = len(pdf_reader.pages) + return num + + +def clean_structure_post(data): + if isinstance(data, dict): + data.pop('page_number', None) + data.pop('start_index', None) + data.pop('end_index', None) + if 'nodes' in data: + clean_structure_post(data['nodes']) + elif isinstance(data, list): + for section in data: + clean_structure_post(section) + return data + + +def print_toc(tree, indent=0): + for node in tree: + print(' ' * indent + node['title']) + if node.get('nodes'): + print_toc(node['nodes'], indent + 1) + + +def print_json(data, max_len=40, indent=2): + def simplify_data(obj): + if isinstance(obj, dict): + return {k: simplify_data(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [simplify_data(item) for item in obj] + elif isinstance(obj, str) and len(obj) > max_len: + return obj[:max_len] + '...' + else: + return obj + + simplified = simplify_data(data) + print(json.dumps(simplified, indent=indent, ensure_ascii=False)) + + +def check_token_limit(structure, limit=110000): + list = structure_to_list(structure) + for node in list: + num_tokens = count_tokens(node['text'], model=None) + if num_tokens > limit: + print(f"Node ID: {node['node_id']} has {num_tokens} tokens") + print("Start Index:", node['start_index']) + print("End Index:", node['end_index']) + print("Title:", node['title']) + print("\n") + + +def convert_physical_index_to_int(data): + if isinstance(data, list): + for i in range(len(data)): + # Check if item is a dictionary and has 'physical_index' key + if isinstance(data[i], dict) and 'physical_index' in data[i]: + if isinstance(data[i]['physical_index'], str): + if data[i]['physical_index'].startswith('').strip()) + elif data[i]['physical_index'].startswith('physical_index_'): + data[i]['physical_index'] = int(data[i]['physical_index'].split('_')[-1].strip()) + elif isinstance(data, str): + if data.startswith('').strip()) + elif data.startswith('physical_index_'): + data = int(data.split('_')[-1].strip()) + # Check data is int + if isinstance(data, int): + return data + else: + return None + return data + + +def convert_page_to_int(data): + for item in data: + if 'page' in item and isinstance(item['page'], str): + try: + item['page'] = int(item['page']) + except ValueError: + # Keep original value if conversion fails + pass + return data + + +def add_node_text_with_labels(node, pdf_pages): + if isinstance(node, dict): + start_page = node.get('start_index') + end_page = node.get('end_index') + node['text'] = get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page) + if 'nodes' in node: + add_node_text_with_labels(node['nodes'], pdf_pages) + elif isinstance(node, list): + for index in range(len(node)): + add_node_text_with_labels(node[index], pdf_pages) + return + + +class ConfigLoader: + """Legacy 0.2.x config helper. Defaults now come from IndexConfig — the + old ``config.yaml`` no longer ships. Prefer ``pageindex.IndexConfig``. + """ + + def __init__(self, default_path=None): + from ..config import IndexConfig + self._default_dict = IndexConfig().model_dump() + + def _validate_keys(self, user_dict): + unknown_keys = set(user_dict) - set(self._default_dict) + if unknown_keys: + raise ValueError(f"Unknown config keys: {unknown_keys}") + + def load(self, user_opt=None) -> config: + """Merge user options over IndexConfig defaults, returning a namespace.""" + if user_opt is None: + user_dict = {} + elif isinstance(user_opt, config): + user_dict = vars(user_opt) + elif isinstance(user_opt, dict): + user_dict = user_opt + else: + raise TypeError("user_opt must be dict, config(SimpleNamespace) or None") + + self._validate_keys(user_dict) + merged = {**self._default_dict, **user_dict} + return config(**merged) + + +def create_node_mapping(tree, include_page_ranges=False, max_page=None): + """Create a mapping of node_id to node for quick lookup. + + The optional page-range arguments are kept for compatibility with the + pageindex 0.2.x SDK utility API. + """ + def get_all_nodes(nodes): + if isinstance(nodes, dict): + return [nodes] + [ + child_node + for child in nodes.get('nodes', []) + for child_node in get_all_nodes(child) + ] + elif isinstance(nodes, list): + return [ + child_node + for item in nodes + for child_node in get_all_nodes(item) + ] + return [] + + all_nodes = get_all_nodes(tree) + + if not include_page_ranges: + return {node["node_id"]: node for node in all_nodes if node.get("node_id")} + + mapping = {} + for i, node in enumerate(all_nodes): + if not node.get("node_id"): + continue + start_page = node.get("page_index", node.get("start_index")) + if node.get("end_index") is not None: + end_page = node.get("end_index") + elif i + 1 < len(all_nodes): + next_node = all_nodes[i + 1] + end_page = next_node.get("page_index", next_node.get("start_index")) + else: + end_page = max_page + + mapping[node["node_id"]] = { + "node": node, + "start_index": start_page, + "end_index": end_page, + } + + return mapping + + +def print_tree(tree, exclude_fields=None, indent=None): + if exclude_fields is None: + exclude_fields = ['text', 'page_index'] + if isinstance(exclude_fields, int): + indent = exclude_fields + exclude_fields = None + if indent is None and exclude_fields is not None: + cleaned_tree = remove_fields(copy.deepcopy(tree), exclude_fields, max_len=40) + pprint(cleaned_tree, sort_dicts=False, width=100) + return + + indent = indent or 0 + for node in tree: + summary = node.get('summary') or node.get('prefix_summary', '') + summary_str = f" — {summary[:60]}..." if summary else "" + print(' ' * indent + f"[{node.get('node_id', '?')}] {node.get('title', '')}{summary_str}") + if node.get('nodes'): + print_tree(node['nodes'], exclude_fields=exclude_fields, indent=indent + 1) + + +def print_wrapped(text, width=100): + for line in text.splitlines(): + print(textwrap.fill(line, width=width)) diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 911cc27..974ce11 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -1,1158 +1,15 @@ -import os -import json -import copy -import math -import random -import re -from .utils import * -import os -from concurrent.futures import ThreadPoolExecutor, as_completed - - -################### 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} - - - page_number = item['physical_index'] - 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} - - -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) - if logger: - logger.info(f"Response: {response}") - return response.get("start_begin", "no") - - -async def check_title_appearance_in_start_concurrent(structure, page_list, model=None, logger=None): - if logger: - logger.info("Checking title appearance in start concurrently") - - # skip items without physical_index - for item in structure: - if item.get('physical_index') is None: - item['appear_start'] = 'no' - - # only for items with valid physical_index - tasks = [] - valid_items = [] - for item in structure: - if item.get('physical_index') is not None: - page_text = page_list[item['physical_index'] - 1][0] - tasks.append(check_title_appearance_in_start(item['title'], page_text, model=model, logger=logger)) - valid_items.append(item) - - results = await asyncio.gather(*tasks, return_exceptions=True) - for item, result in zip(valid_items, results): - if isinstance(result, Exception): - if logger: - logger.error(f"Error checking start for {item['title']}: {result}") - item['appear_start'] = 'no' - else: - item['appear_start'] = result - - return structure - - -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} - - 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'] - - -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 = 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'] - - -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.""" - - 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'] - -def extract_toc_content(content, model=None): - prompt = f""" - Your job is to extract the full table of contents from the given text, replace ... with : - - Given text: {content} - - Directly return the full table of contents content. Do not output anything else.""" - - response, finish_reason = llm_completion(model=model, prompt=prompt, return_finish_reason=True) - - if_complete = check_if_toc_transformation_is_complete(content, response, model) - if if_complete == "yes" and finish_reason == "finished": - return response - - 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) - response = response + new_response - if_complete = check_if_toc_transformation_is_complete(content, response, model) - - 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) - response = response + new_response - if_complete = check_if_toc_transformation_is_complete(content, response, model) - - return response - -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} - - 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'] - -def toc_extractor(page_list, toc_page_list, model): - def transform_dots_to_colon(text): - text = re.sub(r'\.{5,}', ': ', text) - # Handle dots separated by spaces - text = re.sub(r'(?:\. ){5,}\.?', ': ', text) - return text - - toc_content = "" - for page_index in toc_page_list: - toc_content += page_list[page_index][0] - toc_content = transform_dots_to_colon(toc_content) - has_page_index = detect_page_index(toc_content, model=model) - - return { - "toc_content": toc_content, - "page_index_given_in_toc": has_page_index - } - - - - -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. - - 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) - }, - ... - ] - - 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.""" - - 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 - - - -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>, - }, - ... - ], - } - 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} - - 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) - - 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 - - - - -def find_toc_pages(start_page_index, page_list, opt, logger=None): - print('start find_toc_pages') - last_page_is_yes = False - toc_page_list = [] - i = start_page_index - - while i < len(page_list): - # Only check beyond max_pages if we're still finding TOC pages - if i >= opt.toc_check_page_num and not last_page_is_yes: - break - detected_result = toc_detector_single_page(page_list[i][0],model=opt.model) - if detected_result == 'yes': - if logger: - logger.info(f'Page {i} has toc') - toc_page_list.append(i) - last_page_is_yes = True - elif detected_result == 'no' and last_page_is_yes: - if logger: - logger.info(f'Found the last page with toc: {i-1}') - break - i += 1 - - if not toc_page_list and logger: - logger.info('No toc found') - - return toc_page_list - -def remove_page_number(data): - if isinstance(data, dict): - data.pop('page_number', None) - for key in list(data.keys()): - if 'nodes' in key: - remove_page_number(data[key]) - elif isinstance(data, list): - for item in data: - remove_page_number(item) - return data - -def extract_matching_page_pairs(toc_page, toc_physical_index, start_page_index): - pairs = [] - for phy_item in toc_physical_index: - for page_item in toc_page: - if phy_item.get('title') == page_item.get('title'): - physical_index = phy_item.get('physical_index') - if physical_index is not None and int(physical_index) >= start_page_index: - pairs.append({ - 'title': phy_item.get('title'), - 'page': page_item.get('page'), - 'physical_index': physical_index - }) - return pairs - - -def calculate_page_offset(pairs): - differences = [] - for pair in pairs: - try: - physical_index = pair['physical_index'] - page_number = pair['page'] - difference = physical_index - page_number - differences.append(difference) - except (KeyError, TypeError): - continue - - if not differences: - return None - - difference_counts = {} - for diff in differences: - difference_counts[diff] = difference_counts.get(diff, 0) + 1 - - most_common = max(difference_counts.items(), key=lambda x: x[1])[0] - - return most_common - -def add_page_offset_to_toc_json(data, offset): - 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 - del data[i]['page'] - - return data - - - -def page_list_to_group_text(page_contents, token_lengths, max_tokens=20000, overlap_page=1): - num_tokens = sum(token_lengths) - - if num_tokens <= max_tokens: - # merge all pages into one text - page_text = "".join(page_contents) - return [page_text] - - subsets = [] - current_subset = [] - current_token_count = 0 - - expected_parts_num = math.ceil(num_tokens / max_tokens) - average_tokens_per_part = math.ceil(((num_tokens / expected_parts_num) + max_tokens) / 2) - - for i, (page_content, page_tokens) in enumerate(zip(page_contents, token_lengths)): - if current_token_count + page_tokens > average_tokens_per_part: - - subsets.append(''.join(current_subset)) - # Start new subset from overlap if specified - overlap_start = max(i - overlap_page, 0) - current_subset = page_contents[overlap_start:i] - current_token_count = sum(token_lengths[overlap_start:i]) - - # Add current page to the subset - current_subset.append(page_content) - current_token_count += page_tokens - - # Add the last subset if it contains any pages - if current_subset: - subsets.append(''.join(current_subset)) - - print('divide page_list to groups', len(subsets)) - 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) - - for item in json_result: - if 'start' in item: - del item['start'] - return json_result - - -def remove_first_physical_index_section(text): - """ - Removes the first section between <physical_index_X> and <physical_index_X> tags, - and returns the remaining text. - """ - pattern = r'<physical_index_\d+>.*?<physical_index_\d+>' - match = re.search(pattern, text, re.DOTALL) - if match: - # Remove the first matched section - return text.replace(match.group(0), '', 1) - return text - -### add verify completeness -def generate_toc_continue(toc_content, part, model=None): - print('start generate_toc_continue') - prompt = """ - 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}') - -### 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. - - 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 final JSON structure. Do not output anything else.""" - - prompt = prompt + '\nGiven text\n:' + 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}') - -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_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) - 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.extend(toc_with_page_number_additional) - logger.info(f'generate_toc: {toc_with_page_number}') - - toc_with_page_number = convert_physical_index_to_int(toc_with_page_number) - logger.info(f'convert_physical_index_to_int: {toc_with_page_number}') - - return toc_with_page_number - -def process_toc_no_page_numbers(toc_content, toc_page_list, page_list, start_index=1, model=None, logger=None): - page_contents=[] - token_lengths=[] - toc_content = toc_transformer(toc_content, model) - logger.info(f'toc_transformer: {toc_content}') - 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=copy.deepcopy(toc_content) - for group_text in group_texts: - toc_with_page_number = add_page_number_to_toc(group_text, toc_with_page_number, model) - logger.info(f'add_page_number_to_toc: {toc_with_page_number}') - - toc_with_page_number = convert_physical_index_to_int(toc_with_page_number) - logger.info(f'convert_physical_index_to_int: {toc_with_page_number}') - - return toc_with_page_number - - - -def process_toc_with_page_numbers(toc_content, toc_page_list, page_list, toc_check_page_num=None, model=None, logger=None): - toc_with_page_number = toc_transformer(toc_content, model) - logger.info(f'toc_with_page_number: {toc_with_page_number}') - - toc_no_page_number = remove_page_number(copy.deepcopy(toc_with_page_number)) - - start_page_index = toc_page_list[-1] + 1 - main_content = "" - for page_index in range(start_page_index, min(start_page_index + toc_check_page_num, len(page_list))): - main_content += f"<physical_index_{page_index+1}>\n{page_list[page_index][0]}\n<physical_index_{page_index+1}>\n\n" - - toc_with_physical_index = toc_index_extractor(toc_no_page_number, main_content, model) - logger.info(f'toc_with_physical_index: {toc_with_physical_index}') - - toc_with_physical_index = convert_physical_index_to_int(toc_with_physical_index) - logger.info(f'toc_with_physical_index: {toc_with_physical_index}') - - matching_pairs = extract_matching_page_pairs(toc_with_page_number, toc_with_physical_index, start_page_index) - logger.info(f'matching_pairs: {matching_pairs}') - - offset = calculate_page_offset(matching_pairs) - logger.info(f'offset: {offset}') - - toc_with_page_number = add_page_offset_to_toc_json(toc_with_page_number, offset) - logger.info(f'toc_with_page_number: {toc_with_page_number}') - - toc_with_page_number = process_none_page_numbers(toc_with_page_number, page_list, model=model) - logger.info(f'toc_with_page_number: {toc_with_page_number}') - - return toc_with_page_number - - - -##check if needed to process none page numbers -def process_none_page_numbers(toc_items, page_list, start_index=1, model=None): - for i, item in enumerate(toc_items): - if "physical_index" not in item: - # logger.info(f"fix item: {item}") - # Find previous physical_index - prev_physical_index = 0 # Default if no previous item exists - for j in range(i - 1, -1, -1): - if toc_items[j].get('physical_index') is not None: - prev_physical_index = toc_items[j]['physical_index'] - break - - # Find next physical_index - next_physical_index = -1 # Default if no next item exists - for j in range(i + 1, len(toc_items)): - if toc_items[j].get('physical_index') is not None: - next_physical_index = toc_items[j]['physical_index'] - break - - page_contents = [] - for page_index in range(prev_physical_index, next_physical_index+1): - # Add bounds checking to prevent IndexError - list_index = page_index - start_index - if list_index >= 0 and list_index < len(page_list): - page_text = f"<physical_index_{page_index}>\n{page_list[list_index][0]}\n<physical_index_{page_index}>\n\n" - page_contents.append(page_text) - else: - continue - - item_copy = copy.deepcopy(item) - del item_copy['page'] - result = add_page_number_to_toc(page_contents, item_copy, model) - if isinstance(result[0]['physical_index'], str) and result[0]['physical_index'].startswith('<physical_index'): - item['physical_index'] = int(result[0]['physical_index'].split('_')[-1].rstrip('>').strip()) - del item['page'] - - return toc_items - - - - -def check_toc(page_list, opt=None): - toc_page_list = find_toc_pages(start_page_index=0, page_list=page_list, opt=opt) - if len(toc_page_list) == 0: - print('no toc found') - return {'toc_content': None, 'toc_page_list': [], 'page_index_given_in_toc': 'no'} - else: - print('toc found') - toc_json = toc_extractor(page_list, toc_page_list, opt.model) - - if toc_json['page_index_given_in_toc'] == 'yes': - print('index found') - return {'toc_content': toc_json['toc_content'], 'toc_page_list': toc_page_list, 'page_index_given_in_toc': 'yes'} - else: - current_start_index = toc_page_list[-1] + 1 - - while (toc_json['page_index_given_in_toc'] == 'no' and - current_start_index < len(page_list) and - current_start_index < opt.toc_check_page_num): - - additional_toc_pages = find_toc_pages( - start_page_index=current_start_index, - page_list=page_list, - opt=opt - ) - - if len(additional_toc_pages) == 0: - break - - additional_toc_json = toc_extractor(page_list, additional_toc_pages, opt.model) - if additional_toc_json['page_index_given_in_toc'] == 'yes': - print('index found') - return {'toc_content': additional_toc_json['toc_content'], 'toc_page_list': additional_toc_pages, 'page_index_given_in_toc': 'yes'} - - else: - current_start_index = additional_toc_pages[-1] + 1 - print('index not found') - return {'toc_content': toc_json['toc_content'], 'toc_page_list': toc_page_list, 'page_index_given_in_toc': 'no'} - - - - - - -################### 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.""" - - 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']) - - - -async def fix_incorrect_toc(toc_with_page_number, page_list, incorrect_results, start_index=1, model=None, logger=None): - print(f'start fix_incorrect_toc with {len(incorrect_results)} incorrect results') - incorrect_indices = {result['list_index'] for result in incorrect_results} - - end_index = len(page_list) + start_index - 1 - - incorrect_results_and_range_logs = [] - # Helper function to process and check a single incorrect item - async def process_and_check_item(incorrect_item): - list_index = incorrect_item['list_index'] - - # Check if list_index is valid - if list_index < 0 or list_index >= len(toc_with_page_number): - # Return an invalid result for out-of-bounds indices - return { - 'list_index': list_index, - 'title': incorrect_item['title'], - 'physical_index': incorrect_item.get('physical_index'), - 'is_valid': False - } - - # Find the previous correct item - prev_correct = None - for i in range(list_index-1, -1, -1): - if i not in incorrect_indices and i >= 0 and i < len(toc_with_page_number): - physical_index = toc_with_page_number[i].get('physical_index') - if physical_index is not None: - prev_correct = physical_index - break - # If no previous correct item found, use start_index - if prev_correct is None: - prev_correct = start_index - 1 - - # Find the next correct item - next_correct = None - for i in range(list_index+1, len(toc_with_page_number)): - if i not in incorrect_indices and i >= 0 and i < len(toc_with_page_number): - physical_index = toc_with_page_number[i].get('physical_index') - if physical_index is not None: - next_correct = physical_index - break - # If no next correct item found, use end_index - if next_correct is None: - next_correct = end_index - - incorrect_results_and_range_logs.append({ - 'list_index': list_index, - 'title': incorrect_item['title'], - 'prev_correct': prev_correct, - 'next_correct': next_correct - }) - - page_contents=[] - for page_index in range(prev_correct, next_correct+1): - # Add bounds checking to prevent IndexError - page_list_idx = page_index - start_index - if page_list_idx >= 0 and page_list_idx < len(page_list): - page_text = f"<physical_index_{page_index}>\n{page_list[page_list_idx][0]}\n<physical_index_{page_index}>\n\n" - page_contents.append(page_text) - else: - continue - content_range = ''.join(page_contents) - - physical_index_int = await single_toc_item_index_fixer(incorrect_item['title'], content_range, model) - - # Check if the result is correct - check_item = incorrect_item.copy() - check_item['physical_index'] = physical_index_int - check_result = await check_title_appearance(check_item, page_list, start_index, model) - - return { - 'list_index': list_index, - 'title': incorrect_item['title'], - 'physical_index': physical_index_int, - 'is_valid': check_result['answer'] == 'yes' - } - - # Process incorrect items concurrently - tasks = [ - process_and_check_item(item) - for item in incorrect_results - ] - results = await asyncio.gather(*tasks, return_exceptions=True) - for item, result in zip(incorrect_results, results): - if isinstance(result, Exception): - print(f"Processing item {item} generated an exception: {result}") - continue - results = [result for result in results if not isinstance(result, Exception)] - - # Update the toc_with_page_number with the fixed indices and check for any invalid results - invalid_results = [] - for result in results: - if result['is_valid']: - # Add bounds checking to prevent IndexError - list_idx = result['list_index'] - if 0 <= list_idx < len(toc_with_page_number): - toc_with_page_number[list_idx]['physical_index'] = result['physical_index'] - else: - # Index is out of bounds, treat as invalid - invalid_results.append({ - 'list_index': result['list_index'], - 'title': result['title'], - 'physical_index': result['physical_index'], - }) - else: - invalid_results.append({ - 'list_index': result['list_index'], - 'title': result['title'], - 'physical_index': result['physical_index'], - }) - - logger.info(f'incorrect_results_and_range_logs: {incorrect_results_and_range_logs}') - logger.info(f'invalid_results: {invalid_results}') - - return toc_with_page_number, invalid_results - - - -async def fix_incorrect_toc_with_retries(toc_with_page_number, page_list, incorrect_results, start_index=1, max_attempts=3, model=None, logger=None): - print('start fix_incorrect_toc') - fix_attempt = 0 - current_toc = toc_with_page_number - current_incorrect = incorrect_results - - while current_incorrect: - print(f"Fixing {len(current_incorrect)} incorrect results") - - current_toc, current_incorrect = await fix_incorrect_toc(current_toc, page_list, current_incorrect, start_index, model, logger) - - fix_attempt += 1 - if fix_attempt >= max_attempts: - logger.info("Maximum fix attempts reached") - break - - return current_toc, current_incorrect - - - - -################### verify toc ######################################################### -async def verify_toc(page_list, list_result, start_index=1, N=None, model=None): - print('start verify_toc') - # Find the last non-None physical_index - last_physical_index = None - for item in reversed(list_result): - if item.get('physical_index') is not None: - last_physical_index = item['physical_index'] - break - - # Early return if we don't have valid physical indices - if last_physical_index is None or last_physical_index < len(page_list)/2: - return 0, [] - - # Determine which items to check - if N is None: - print('check all items') - sample_indices = range(0, len(list_result)) - else: - N = min(N, len(list_result)) - print(f'check {N} items') - sample_indices = random.sample(range(0, len(list_result)), N) - - # Prepare items with their list indices - indexed_sample_list = [] - for idx in sample_indices: - item = list_result[idx] - # Skip items with None physical_index (these were invalidated by validate_and_truncate_physical_indices) - if item.get('physical_index') is not None: - item_with_index = item.copy() - item_with_index['list_index'] = idx # Add the original index in list_result - indexed_sample_list.append(item_with_index) - - # Run checks concurrently - tasks = [ - check_title_appearance(item, page_list, start_index, model) - for item in indexed_sample_list - ] - results = await asyncio.gather(*tasks) - - # Process results - correct_count = 0 - incorrect_results = [] - for result in results: - if result['answer'] == 'yes': - correct_count += 1 - else: - incorrect_results.append(result) - - # Calculate accuracy - checked_count = len(results) - accuracy = correct_count / checked_count if checked_count > 0 else 0 - print(f"accuracy: {accuracy*100:.2f}%") - return accuracy, incorrect_results - - - - - -################### main process ######################################################### -async def meta_processor(page_list, mode=None, toc_content=None, toc_page_list=None, start_index=1, opt=None, logger=None): - print(mode) - print(f'start_index: {start_index}') - - if mode == 'process_toc_with_page_numbers': - toc_with_page_number = process_toc_with_page_numbers(toc_content, toc_page_list, page_list, toc_check_page_num=opt.toc_check_page_num, model=opt.model, logger=logger) - elif mode == 'process_toc_no_page_numbers': - toc_with_page_number = process_toc_no_page_numbers(toc_content, toc_page_list, page_list, model=opt.model, logger=logger) - else: - toc_with_page_number = process_no_toc(page_list, start_index=start_index, model=opt.model, logger=logger) - - toc_with_page_number = [item for item in toc_with_page_number if item.get('physical_index') is not None] - - toc_with_page_number = validate_and_truncate_physical_indices( - toc_with_page_number, - len(page_list), - start_index=start_index, - logger=logger - ) - - accuracy, incorrect_results = await verify_toc(page_list, toc_with_page_number, start_index=start_index, model=opt.model) - - logger.info({ - 'mode': 'process_toc_with_page_numbers', - 'accuracy': accuracy, - 'incorrect_results': incorrect_results - }) - if accuracy == 1.0 and len(incorrect_results) == 0: - return toc_with_page_number - if accuracy > 0.6 and len(incorrect_results) > 0: - toc_with_page_number, incorrect_results = await fix_incorrect_toc_with_retries(toc_with_page_number, page_list, incorrect_results,start_index=start_index, max_attempts=3, model=opt.model, logger=logger) - return toc_with_page_number - else: - if mode == 'process_toc_with_page_numbers': - return await meta_processor(page_list, mode='process_toc_no_page_numbers', toc_content=toc_content, toc_page_list=toc_page_list, start_index=start_index, opt=opt, logger=logger) - elif mode == 'process_toc_no_page_numbers': - return await meta_processor(page_list, mode='process_no_toc', start_index=start_index, opt=opt, logger=logger) - else: - raise Exception('Processing failed') - - -async def process_large_node_recursively(node, page_list, opt=None, logger=None): - node_page_list = page_list[node['start_index']-1:node['end_index']] - token_num = sum([page[1] for page in node_page_list]) - - if node['end_index'] - node['start_index'] > opt.max_page_num_each_node and token_num >= opt.max_token_num_each_node: - print('large node:', node['title'], 'start_index:', node['start_index'], 'end_index:', node['end_index'], 'token_num:', token_num) - - node_toc_tree = await meta_processor(node_page_list, mode='process_no_toc', start_index=node['start_index'], opt=opt, logger=logger) - node_toc_tree = await check_title_appearance_in_start_concurrent(node_toc_tree, page_list, model=opt.model, logger=logger) - - # Filter out items with None physical_index before post_processing - valid_node_toc_items = [item for item in node_toc_tree if item.get('physical_index') is not None] - - if valid_node_toc_items and node['title'].strip() == valid_node_toc_items[0]['title'].strip(): - node['nodes'] = post_processing(valid_node_toc_items[1:], node['end_index']) - node['end_index'] = valid_node_toc_items[1]['start_index'] if len(valid_node_toc_items) > 1 else node['end_index'] - else: - node['nodes'] = post_processing(valid_node_toc_items, node['end_index']) - node['end_index'] = valid_node_toc_items[0]['start_index'] if valid_node_toc_items else node['end_index'] - - if 'nodes' in node and node['nodes']: - tasks = [ - process_large_node_recursively(child_node, page_list, opt, logger=logger) - for child_node in node['nodes'] - ] - await asyncio.gather(*tasks) - - return node - -async def tree_parser(page_list, opt, doc=None, logger=None): - check_toc_result = check_toc(page_list, opt) - logger.info(check_toc_result) - - if check_toc_result.get("toc_content") and check_toc_result["toc_content"].strip() and check_toc_result["page_index_given_in_toc"] == "yes": - toc_with_page_number = await meta_processor( - page_list, - mode='process_toc_with_page_numbers', - start_index=1, - toc_content=check_toc_result['toc_content'], - toc_page_list=check_toc_result['toc_page_list'], - opt=opt, - logger=logger) - else: - toc_with_page_number = await meta_processor( - page_list, - mode='process_no_toc', - start_index=1, - opt=opt, - logger=logger) - - toc_with_page_number = add_preface_if_needed(toc_with_page_number) - toc_with_page_number = await check_title_appearance_in_start_concurrent(toc_with_page_number, page_list, model=opt.model, logger=logger) - - # Filter out items with None physical_index before post_processings - valid_toc_items = [item for item in toc_with_page_number if item.get('physical_index') is not None] - - toc_tree = post_processing(valid_toc_items, len(page_list)) - tasks = [ - process_large_node_recursively(node, page_list, opt, logger=logger) - for node in toc_tree - ] - await asyncio.gather(*tasks) - - return toc_tree - - -def page_index_main(doc, opt=None): - logger = JsonLogger(doc) - - is_valid_pdf = ( - (isinstance(doc, str) and os.path.isfile(doc) and doc.lower().endswith(".pdf")) or - isinstance(doc, BytesIO) - ) - if not is_valid_pdf: - raise ValueError("Unsupported input type. Expected a PDF file path or BytesIO object.") - - print('Parsing PDF...') - page_list = get_page_tokens(doc, model=opt.model) - - logger.info({'total_page_number': len(page_list)}) - logger.info({'total_token': sum([page[1] for page in page_list])}) - - async def page_index_builder(): - structure = await tree_parser(page_list, opt, doc=doc, logger=logger) - # IndexConfig fields are booleans (pydantic coerces legacy 'yes'/'no' - # strings at the boundary) — comparing against 'yes' here would be - # always-False and silently skip every enhancement. - if opt.if_add_node_id: - write_node_id(structure) - if opt.if_add_node_text: - add_node_text(structure, page_list) - if opt.if_add_node_summary: - if not opt.if_add_node_text: - add_node_text(structure, page_list) - await generate_summaries_for_structure(structure, model=opt.model) - if not opt.if_add_node_text: - remove_structure_text(structure) - if opt.if_add_doc_description: - # Create a clean structure without unnecessary fields for description generation - clean_structure = create_clean_structure_for_description(structure) - doc_description = generate_doc_description(clean_structure, model=opt.model) - structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'summary', 'text', 'nodes']) - return { - 'doc_name': get_pdf_name(doc), - 'doc_description': doc_description, - 'structure': structure, - } - structure = format_structure(structure, order=['title', 'node_id', 'start_index', 'end_index', 'summary', 'text', 'nodes']) - return { - 'doc_name': get_pdf_name(doc), - 'structure': structure, - } - - return asyncio.run(page_index_builder()) - - -def page_index(doc, model=None, toc_check_page_num=None, max_page_num_each_node=None, max_token_num_each_node=None, - if_add_node_id=None, if_add_node_summary=None, if_add_doc_description=None, if_add_node_text=None): - - from .config import IndexConfig - user_opt = { - arg: value for arg, value in locals().items() - if arg != "doc" and value is not None - } - opt = IndexConfig(**user_opt) - return page_index_main(doc, opt) - - -def validate_and_truncate_physical_indices(toc_with_page_number, page_list_length, start_index=1, logger=None): - """ - Validates and truncates physical indices that exceed the actual document length. - This prevents errors when TOC references pages that don't exist in the document (e.g. the file is broken or incomplete). - """ - if not toc_with_page_number: - return toc_with_page_number - - max_allowed_page = page_list_length + start_index - 1 - truncated_items = [] - - for i, item in enumerate(toc_with_page_number): - if item.get('physical_index') is not None: - original_index = item['physical_index'] - if original_index > max_allowed_page: - item['physical_index'] = None - truncated_items.append({ - 'title': item.get('title', 'Unknown'), - 'original_index': original_index - }) - if logger: - logger.info(f"Removed physical_index for '{item.get('title', 'Unknown')}' (was {original_index}, too far beyond document)") - - if truncated_items and logger: - logger.info(f"Total removed items: {len(truncated_items)}") - - print(f"Document validation: {page_list_length} pages, max allowed index: {max_allowed_page}") - if truncated_items: - print(f"Truncated {len(truncated_items)} TOC items that exceeded document length") - - return toc_with_page_number \ No newline at end of file +# pageindex/page_index.py +# Deprecation shim. The PDF indexing pipeline now lives in +# pageindex/index/page_index.py (the single source of truth). This module +# re-exports it so legacy imports (`from pageindex.page_index import ...`, +# `from pageindex import page_index`) keep working. +import warnings + +warnings.warn( + "pageindex.page_index has moved to pageindex.index.page_index; importing it " + "from the top level is deprecated and will be removed in a future release.", + PendingDeprecationWarning, + stacklevel=2, +) + +from .index.page_index import * # noqa: F401,F403,E402 diff --git a/pageindex/page_index_md.py b/pageindex/page_index_md.py index 5a59716..55ca073 100644 --- a/pageindex/page_index_md.py +++ b/pageindex/page_index_md.py @@ -1,342 +1,38 @@ -import asyncio -import json -import re -import os -try: - from .utils import * -except: - from utils import * +# pageindex/page_index_md.py +# Deprecation shim. The Markdown indexing pipeline now lives in +# pageindex/index/page_index_md.py (the single source of truth). This module +# re-exports it so legacy imports keep working. +# +# The canonical md_to_tree takes booleans; legacy callers passed 'yes'/'no' +# strings, so the wrapper below coerces them (a bare 'no' is otherwise truthy). +import warnings -async def get_node_summary(node, summary_token_threshold=200, model=None): - node_text = node.get('text') - num_tokens = count_tokens(node_text, model=model) - if num_tokens < summary_token_threshold: - return node_text - else: - return await generate_node_summary(node, model=model) +warnings.warn( + "pageindex.page_index_md has moved to pageindex.index.page_index_md; " + "importing it from the top level is deprecated and will be removed in a " + "future release.", + PendingDeprecationWarning, + stacklevel=2, +) + +from .index.page_index_md import * # noqa: F401,F403,E402 +from .index.page_index_md import md_to_tree as _md_to_tree # noqa: E402 + +_BOOL_PARAMS = ( + "if_thinning", "if_add_node_summary", "if_add_doc_description", + "if_add_node_text", "if_add_node_id", +) -async def generate_summaries_for_structure_md(structure, summary_token_threshold, model=None): - nodes = structure_to_list(structure) - tasks = [get_node_summary(node, summary_token_threshold=summary_token_threshold, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) - - for node, summary in zip(nodes, summaries): - if not node.get('nodes'): - node['summary'] = summary - else: - node['prefix_summary'] = summary - return structure +def _coerce_bool(value): + if isinstance(value, str): + return value.strip().lower() in ("yes", "true", "1", "y", "on") + return bool(value) -def extract_nodes_from_markdown(markdown_content): - header_pattern = r'^(#{1,6})\s+(.+)$' - code_block_pattern = r'^```' - node_list = [] - - lines = markdown_content.split('\n') - in_code_block = False - - for line_num, line in enumerate(lines, 1): - stripped_line = line.strip() - - # Check for code block delimiters (triple backticks) - if re.match(code_block_pattern, stripped_line): - in_code_block = not in_code_block - continue - - # Skip empty lines - if not stripped_line: - continue - - # Only look for headers when not inside a code block - if not in_code_block: - match = re.match(header_pattern, stripped_line) - if match: - title = match.group(2).strip() - node_list.append({'node_title': title, 'line_num': line_num}) - - return node_list, lines - - -def extract_node_text_content(node_list, markdown_lines): - all_nodes = [] - for node in node_list: - line_content = markdown_lines[node['line_num'] - 1] - header_match = re.match(r'^(#{1,6})', line_content) - - if header_match is None: - print(f"Warning: Line {node['line_num']} does not contain a valid header: '{line_content}'") - continue - - processed_node = { - 'title': node['node_title'], - 'line_num': node['line_num'], - 'level': len(header_match.group(1)) - } - all_nodes.append(processed_node) - - for i, node in enumerate(all_nodes): - start_line = node['line_num'] - 1 - if i + 1 < len(all_nodes): - end_line = all_nodes[i + 1]['line_num'] - 1 - else: - end_line = len(markdown_lines) - - node['text'] = '\n'.join(markdown_lines[start_line:end_line]).strip() - return all_nodes - -def update_node_list_with_text_token_count(node_list, model=None): - - def find_all_children(parent_index, parent_level, node_list): - """Find all direct and indirect children of a parent node""" - children_indices = [] - - # Look for children after the parent - for i in range(parent_index + 1, len(node_list)): - current_level = node_list[i]['level'] - - # If we hit a node at same or higher level than parent, stop - if current_level <= parent_level: - break - - # This is a descendant - children_indices.append(i) - - return children_indices - - # Make a copy to avoid modifying the original - result_list = node_list.copy() - - # Process nodes from end to beginning to ensure children are processed before parents - for i in range(len(result_list) - 1, -1, -1): - current_node = result_list[i] - current_level = current_node['level'] - - # Get all children of this node - children_indices = find_all_children(i, current_level, result_list) - - # Start with the node's own text - node_text = current_node.get('text', '') - total_text = node_text - - # Add all children's text - for child_index in children_indices: - child_text = result_list[child_index].get('text', '') - if child_text: - total_text += '\n' + child_text - - # Calculate token count for combined text - result_list[i]['text_token_count'] = count_tokens(total_text, model=model) - - return result_list - - -def tree_thinning_for_index(node_list, min_node_token=None, model=None): - def find_all_children(parent_index, parent_level, node_list): - children_indices = [] - - for i in range(parent_index + 1, len(node_list)): - current_level = node_list[i]['level'] - - if current_level <= parent_level: - break - - children_indices.append(i) - - return children_indices - - result_list = node_list.copy() - nodes_to_remove = set() - - for i in range(len(result_list) - 1, -1, -1): - if i in nodes_to_remove: - continue - - current_node = result_list[i] - current_level = current_node['level'] - - total_tokens = current_node.get('text_token_count', 0) - - if total_tokens < min_node_token: - children_indices = find_all_children(i, current_level, result_list) - - children_texts = [] - for child_index in sorted(children_indices): - if child_index not in nodes_to_remove: - child_text = result_list[child_index].get('text', '') - if child_text.strip(): - children_texts.append(child_text) - nodes_to_remove.add(child_index) - - if children_texts: - parent_text = current_node.get('text', '') - merged_text = parent_text - for child_text in children_texts: - if merged_text and not merged_text.endswith('\n'): - merged_text += '\n\n' - merged_text += child_text - - result_list[i]['text'] = merged_text - - result_list[i]['text_token_count'] = count_tokens(merged_text, model=model) - - for index in sorted(nodes_to_remove, reverse=True): - result_list.pop(index) - - return result_list - - -def build_tree_from_nodes(node_list): - if not node_list: - return [] - - stack = [] - root_nodes = [] - node_counter = 1 - - for node in node_list: - current_level = node['level'] - - tree_node = { - 'title': node['title'], - 'node_id': str(node_counter).zfill(4), - 'text': node['text'], - 'line_num': node['line_num'], - 'nodes': [] - } - node_counter += 1 - - while stack and stack[-1][1] >= current_level: - stack.pop() - - if not stack: - root_nodes.append(tree_node) - else: - parent_node, parent_level = stack[-1] - parent_node['nodes'].append(tree_node) - - stack.append((tree_node, current_level)) - - return root_nodes - - -def clean_tree_for_output(tree_nodes): - cleaned_nodes = [] - - for node in tree_nodes: - cleaned_node = { - 'title': node['title'], - 'node_id': node['node_id'], - 'text': node['text'], - 'line_num': node['line_num'] - } - - if node['nodes']: - cleaned_node['nodes'] = clean_tree_for_output(node['nodes']) - - cleaned_nodes.append(cleaned_node) - - return cleaned_nodes - - -async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_add_node_summary='no', summary_token_threshold=None, model=None, if_add_doc_description='no', if_add_node_text='no', if_add_node_id='yes'): - with open(md_path, 'r', encoding='utf-8') as f: - markdown_content = f.read() - line_count = markdown_content.count('\n') + 1 - - print(f"Extracting nodes from markdown...") - node_list, markdown_lines = extract_nodes_from_markdown(markdown_content) - - print(f"Extracting text content from nodes...") - nodes_with_content = extract_node_text_content(node_list, markdown_lines) - - if if_thinning: - nodes_with_content = update_node_list_with_text_token_count(nodes_with_content, model=model) - print(f"Thinning nodes...") - nodes_with_content = tree_thinning_for_index(nodes_with_content, min_token_threshold, model=model) - - print(f"Building tree from nodes...") - tree_structure = build_tree_from_nodes(nodes_with_content) - - if if_add_node_id == 'yes': - write_node_id(tree_structure) - - print(f"Formatting tree structure...") - - if if_add_node_summary == 'yes': - # Always include text for summary generation - tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'line_num', 'summary', 'prefix_summary', 'text', 'nodes']) - - print(f"Generating summaries for each node...") - tree_structure = await generate_summaries_for_structure_md(tree_structure, summary_token_threshold=summary_token_threshold, model=model) - - if if_add_node_text == 'no': - # Remove text after summary generation if not requested - tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'line_num', 'summary', 'prefix_summary', 'nodes']) - - if if_add_doc_description == 'yes': - print(f"Generating document description...") - # Create a clean structure without unnecessary fields for description generation - clean_structure = create_clean_structure_for_description(tree_structure) - doc_description = generate_doc_description(clean_structure, model=model) - return { - 'doc_name': os.path.splitext(os.path.basename(md_path))[0], - 'doc_description': doc_description, - 'line_count': line_count, - 'structure': tree_structure, - } - else: - # No summaries needed, format based on text preference - if if_add_node_text == 'yes': - tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'line_num', 'summary', 'prefix_summary', 'text', 'nodes']) - else: - tree_structure = format_structure(tree_structure, order = ['title', 'node_id', 'line_num', 'summary', 'prefix_summary', 'nodes']) - - return { - 'doc_name': os.path.splitext(os.path.basename(md_path))[0], - 'line_count': line_count, - 'structure': tree_structure, - } - - -if __name__ == "__main__": - import os - import json - - # MD_NAME = 'Detect-Order-Construct' - MD_NAME = 'cognitive-load' - MD_PATH = os.path.join(os.path.dirname(__file__), '..', 'examples/documents/', f'{MD_NAME}.md') - - - MODEL="gpt-4.1" - IF_THINNING=False - THINNING_THRESHOLD=5000 - SUMMARY_TOKEN_THRESHOLD=200 - IF_SUMMARY=True - - tree_structure = asyncio.run(md_to_tree( - md_path=MD_PATH, - if_thinning=IF_THINNING, - min_token_threshold=THINNING_THRESHOLD, - if_add_node_summary='yes' if IF_SUMMARY else 'no', - summary_token_threshold=SUMMARY_TOKEN_THRESHOLD, - model=MODEL)) - - print('\n' + '='*60) - print('TREE STRUCTURE') - print('='*60) - print_json(tree_structure) - - print('\n' + '='*60) - print('TABLE OF CONTENTS') - print('='*60) - print_toc(tree_structure['structure']) - - output_path = os.path.join(os.path.dirname(__file__), '..', 'results', f'{MD_NAME}_structure.json') - os.makedirs(os.path.dirname(output_path), exist_ok=True) - - with open(output_path, 'w', encoding='utf-8') as f: - json.dump(tree_structure, f, indent=2, ensure_ascii=False) - - print(f"\nTree structure saved to: {output_path}") \ No newline at end of file +async def md_to_tree(*args, **kwargs): + """Legacy wrapper: coerce 'yes'/'no' string flags to bool, then delegate.""" + for key in _BOOL_PARAMS: + if key in kwargs: + kwargs[key] = _coerce_bool(kwargs[key]) + return await _md_to_tree(*args, **kwargs) diff --git a/pageindex/retrieve.py b/pageindex/retrieve.py index 55c3850..e0e1537 100644 --- a/pageindex/retrieve.py +++ b/pageindex/retrieve.py @@ -2,9 +2,9 @@ import json import PyPDF2 try: - from .utils import get_number_of_pages, remove_fields + from .index.utils import get_number_of_pages, remove_fields except ImportError: - from utils import get_number_of_pages, remove_fields + from index.utils import get_number_of_pages, remove_fields # ── Helpers ────────────────────────────────────────────────────────────────── diff --git a/pageindex/utils.py b/pageindex/utils.py index c40e6ba..fe403d2 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -1,777 +1,14 @@ -import litellm -import logging -import os -import textwrap -from datetime import datetime -import time -import json -import PyPDF2 -import copy -import asyncio -import pymupdf -from io import BytesIO -from dotenv import load_dotenv -load_dotenv() -import logging -import yaml -from pathlib import Path -from pprint import pprint -from types import SimpleNamespace as config - -from .config import get_llm_params - -# 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"): - os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") - - -async def call_llm(prompt, api_key, model="gpt-4.1", temperature=0): - """Call an LLM to generate a response to a prompt. - - Kept for compatibility with the pageindex 0.2.x SDK utility API. - """ - import openai - - client = openai.AsyncOpenAI(api_key=api_key) - response = await client.chat.completions.create( - model=model, - messages=[{"role": "user", "content": prompt}], - temperature=temperature, - ) - return response.choices[0].message.content.strip() - - -def count_tokens(text, model=None): - if not text: - return 0 - return litellm.token_counter(model=model, text=text) - - -def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): - if model: - model = model.removeprefix("litellm/") - max_retries = 10 - messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] - for i in range(max_retries): - try: - response = litellm.completion( - model=model, - messages=messages, - # Per-call litellm kwargs (default temperature=0, drop_params=True); - # configure via config.set_llm_params(...) — never the litellm global. - **get_llm_params(), - ) - content = response.choices[0].message.content - if return_finish_reason: - finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished" - return content, finish_reason - return content - except Exception as e: - print('************* Retrying *************') - logging.error(f"Error: {e}") - if i < max_retries - 1: - time.sleep(1) - else: - logging.error('Max retries reached for prompt: ' + prompt) - if return_finish_reason: - return "", "error" - return "" - - - -async def llm_acompletion(model, prompt): - if model: - model = model.removeprefix("litellm/") - max_retries = 10 - messages = [{"role": "user", "content": prompt}] - for i in range(max_retries): - try: - response = await litellm.acompletion( - model=model, - messages=messages, - **get_llm_params(), # per-call kwargs; never the litellm global - ) - return response.choices[0].message.content - except Exception as e: - print('************* Retrying *************') - logging.error(f"Error: {e}") - if i < max_retries - 1: - await asyncio.sleep(1) - 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): - data['node_id'] = str(node_id).zfill(4) - node_id += 1 - for key in list(data.keys()): - if 'nodes' in key: - node_id = write_node_id(data[key], node_id) - elif isinstance(data, list): - for index in range(len(data)): - node_id = write_node_id(data[index], node_id) - return node_id - -def get_nodes(structure): - if isinstance(structure, dict): - structure_node = copy.deepcopy(structure) - structure_node.pop('nodes', None) - nodes = [structure_node] - for key in list(structure.keys()): - if 'nodes' in key: - nodes.extend(get_nodes(structure[key])) - return nodes - elif isinstance(structure, list): - nodes = [] - for item in structure: - nodes.extend(get_nodes(item)) - return nodes - -def structure_to_list(structure): - if isinstance(structure, dict): - nodes = [] - nodes.append(structure) - if 'nodes' in structure: - nodes.extend(structure_to_list(structure['nodes'])) - return nodes - elif isinstance(structure, list): - nodes = [] - for item in structure: - nodes.extend(structure_to_list(item)) - return nodes - - -def get_leaf_nodes(structure): - if isinstance(structure, dict): - if not structure.get('nodes'): - structure_node = copy.deepcopy(structure) - structure_node.pop('nodes', None) - return [structure_node] - else: - leaf_nodes = [] - for key in list(structure.keys()): - if 'nodes' in key: - leaf_nodes.extend(get_leaf_nodes(structure[key])) - return leaf_nodes - elif isinstance(structure, list): - leaf_nodes = [] - for item in structure: - leaf_nodes.extend(get_leaf_nodes(item)) - return leaf_nodes - -def is_leaf_node(data, node_id): - # Helper function to find the node by its node_id - def find_node(data, node_id): - if isinstance(data, dict): - if data.get('node_id') == node_id: - return data - for key in data.keys(): - if 'nodes' in key: - result = find_node(data[key], node_id) - if result: - return result - elif isinstance(data, list): - for item in data: - result = find_node(item, node_id) - if result: - return result - return None - - # Find the node with the given node_id - node = find_node(data, node_id) - - # Check if the node is a leaf node - if node and not node.get('nodes'): - return True - return False - -def get_last_node(structure): - return structure[-1] - - -def extract_text_from_pdf(pdf_path): - pdf_reader = PyPDF2.PdfReader(pdf_path) - ###return text not list - text="" - for page_num in range(len(pdf_reader.pages)): - page = pdf_reader.pages[page_num] - text+=page.extract_text() - return text - -def get_pdf_title(pdf_path): - pdf_reader = PyPDF2.PdfReader(pdf_path) - meta = pdf_reader.metadata - title = meta.title if meta and meta.title else 'Untitled' - return title - -def get_text_of_pages(pdf_path, start_page, end_page, tag=True): - pdf_reader = PyPDF2.PdfReader(pdf_path) - text = "" - for page_num in range(start_page-1, end_page): - page = pdf_reader.pages[page_num] - page_text = page.extract_text() - if tag: - text += f"<start_index_{page_num+1}>\n{page_text}\n<end_index_{page_num+1}>\n" - else: - text += page_text - return text - -def get_first_start_page_from_text(text): - start_page = -1 - start_page_match = re.search(r'<start_index_(\d+)>', text) - if start_page_match: - start_page = int(start_page_match.group(1)) - return start_page - -def get_last_start_page_from_text(text): - start_page = -1 - # Find all matches of start_index tags - start_page_matches = re.finditer(r'<start_index_(\d+)>', text) - # Convert iterator to list and get the last match if any exist - matches_list = list(start_page_matches) - if matches_list: - start_page = int(matches_list[-1].group(1)) - return start_page - - -def sanitize_filename(filename, replacement='-'): - # In Linux, only '/' and '\0' (null) are invalid in filenames. - # Null can't be represented in strings, so we only handle '/'. - return filename.replace('/', replacement) - -def get_pdf_name(pdf_path): - # Extract PDF name - if isinstance(pdf_path, str): - pdf_name = os.path.basename(pdf_path) - elif isinstance(pdf_path, BytesIO): - pdf_reader = PyPDF2.PdfReader(pdf_path) - meta = pdf_reader.metadata - pdf_name = meta.title if meta and meta.title else 'Untitled' - pdf_name = sanitize_filename(pdf_name) - return pdf_name - - -class JsonLogger: - def __init__(self, file_path): - # Extract PDF name for logger name - pdf_name = get_pdf_name(file_path) - - current_time = datetime.now().strftime("%Y%m%d_%H%M%S") - self.filename = f"{pdf_name}_{current_time}.json" - os.makedirs("./logs", exist_ok=True) - # Initialize empty list to store all messages - self.log_data = [] - - def log(self, level, message, **kwargs): - if isinstance(message, dict): - self.log_data.append(message) - else: - self.log_data.append({'message': message}) - # Add new message to the log data - - # Write entire log data to file - with open(self._filepath(), "w") as f: - json.dump(self.log_data, f, indent=2) - - def info(self, message, **kwargs): - self.log("INFO", message, **kwargs) - - def error(self, message, **kwargs): - self.log("ERROR", message, **kwargs) - - def debug(self, message, **kwargs): - self.log("DEBUG", message, **kwargs) - - def exception(self, message, **kwargs): - kwargs["exception"] = True - self.log("ERROR", message, **kwargs) - - def _filepath(self): - return os.path.join("logs", self.filename) - - - - -def list_to_tree(data): - def get_parent_structure(structure): - """Helper function to get the parent structure code""" - if not structure: - return None - parts = str(structure).split('.') - return '.'.join(parts[:-1]) if len(parts) > 1 else None - - # First pass: Create nodes and track parent-child relationships - nodes = {} - root_nodes = [] - - for item in data: - structure = item.get('structure') - node = { - 'title': item.get('title'), - 'start_index': item.get('start_index'), - 'end_index': item.get('end_index'), - 'nodes': [] - } - - nodes[structure] = node - - # Find parent - parent_structure = get_parent_structure(structure) - - if parent_structure: - # Add as child to parent if parent exists - if parent_structure in nodes: - nodes[parent_structure]['nodes'].append(node) - else: - root_nodes.append(node) - else: - # No parent, this is a root node - root_nodes.append(node) - - # Helper function to clean empty children arrays - def clean_node(node): - if not node['nodes']: - del node['nodes'] - else: - for child in node['nodes']: - clean_node(child) - return node - - # Clean and return the tree - return [clean_node(node) for node in root_nodes] - -def add_preface_if_needed(data): - if not isinstance(data, list) or not data: - return data - - if data[0]['physical_index'] is not None and data[0]['physical_index'] > 1: - preface_node = { - "structure": "0", - "title": "Preface", - "physical_index": 1, - } - data.insert(0, preface_node) - return data - - - -def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): - if pdf_parser == "PyPDF2": - pdf_reader = PyPDF2.PdfReader(pdf_path) - page_list = [] - for page_num in range(len(pdf_reader.pages)): - page = pdf_reader.pages[page_num] - page_text = page.extract_text() - token_length = litellm.token_counter(model=model, text=page_text) - page_list.append((page_text, token_length)) - return page_list - elif pdf_parser == "PyMuPDF": - if isinstance(pdf_path, BytesIO): - pdf_stream = pdf_path - doc = pymupdf.open(stream=pdf_stream, filetype="pdf") - elif isinstance(pdf_path, str) and os.path.isfile(pdf_path) and pdf_path.lower().endswith(".pdf"): - doc = pymupdf.open(pdf_path) - page_list = [] - for page in doc: - page_text = page.get_text() - token_length = litellm.token_counter(model=model, text=page_text) - page_list.append((page_text, token_length)) - return page_list - else: - raise ValueError(f"Unsupported PDF parser: {pdf_parser}") - - - -def get_text_of_pdf_pages(pdf_pages, start_page, end_page): - text = "" - for page_num in range(start_page-1, end_page): - text += pdf_pages[page_num][0] - return text - -def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page): - text = "" - for page_num in range(start_page-1, end_page): - text += f"<physical_index_{page_num+1}>\n{pdf_pages[page_num][0]}\n<physical_index_{page_num+1}>\n" - return text - -def get_number_of_pages(pdf_path): - pdf_reader = PyPDF2.PdfReader(pdf_path) - num = len(pdf_reader.pages) - return num - - - -def post_processing(structure, end_physical_index): - # First convert page_number to start_index in flat list - for i, item in enumerate(structure): - item['start_index'] = item.get('physical_index') - if i < len(structure) - 1: - if structure[i + 1].get('appear_start') == 'yes': - item['end_index'] = structure[i + 1]['physical_index']-1 - else: - item['end_index'] = structure[i + 1]['physical_index'] - else: - item['end_index'] = end_physical_index - tree = list_to_tree(structure) - if len(tree)!=0: - return tree - else: - ### remove appear_start - for node in structure: - node.pop('appear_start', None) - node.pop('physical_index', None) - return structure - -def clean_structure_post(data): - if isinstance(data, dict): - data.pop('page_number', None) - data.pop('start_index', None) - data.pop('end_index', None) - if 'nodes' in data: - clean_structure_post(data['nodes']) - elif isinstance(data, list): - for section in data: - clean_structure_post(section) - return data - -def remove_fields(data, fields=['text'], max_len=None): - if isinstance(data, dict): - return {k: remove_fields(v, fields, max_len) - for k, v in data.items() if k not in fields} - elif isinstance(data, list): - return [remove_fields(item, fields, max_len) for item in data] - elif isinstance(data, str): - return data[:max_len] + '...' if max_len is not None and len(data) > max_len else data - return data - -def print_toc(tree, indent=0): - for node in tree: - print(' ' * indent + node['title']) - if node.get('nodes'): - print_toc(node['nodes'], indent + 1) - -def print_json(data, max_len=40, indent=2): - def simplify_data(obj): - if isinstance(obj, dict): - return {k: simplify_data(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [simplify_data(item) for item in obj] - elif isinstance(obj, str) and len(obj) > max_len: - return obj[:max_len] + '...' - else: - return obj - - simplified = simplify_data(data) - print(json.dumps(simplified, indent=indent, ensure_ascii=False)) - - -def remove_structure_text(data): - if isinstance(data, dict): - data.pop('text', None) - if 'nodes' in data: - remove_structure_text(data['nodes']) - elif isinstance(data, list): - for item in data: - remove_structure_text(item) - return data - - -def check_token_limit(structure, limit=110000): - list = structure_to_list(structure) - for node in list: - num_tokens = count_tokens(node['text'], model=None) - if num_tokens > limit: - print(f"Node ID: {node['node_id']} has {num_tokens} tokens") - print("Start Index:", node['start_index']) - print("End Index:", node['end_index']) - print("Title:", node['title']) - print("\n") - - -def convert_physical_index_to_int(data): - if isinstance(data, list): - for i in range(len(data)): - # Check if item is a dictionary and has 'physical_index' key - if isinstance(data[i], dict) and 'physical_index' in data[i]: - if isinstance(data[i]['physical_index'], str): - if data[i]['physical_index'].startswith('<physical_index_'): - data[i]['physical_index'] = int(data[i]['physical_index'].split('_')[-1].rstrip('>').strip()) - elif data[i]['physical_index'].startswith('physical_index_'): - data[i]['physical_index'] = int(data[i]['physical_index'].split('_')[-1].strip()) - elif isinstance(data, str): - if data.startswith('<physical_index_'): - data = int(data.split('_')[-1].rstrip('>').strip()) - elif data.startswith('physical_index_'): - data = int(data.split('_')[-1].strip()) - # Check data is int - if isinstance(data, int): - return data - else: - return None - return data - - -def convert_page_to_int(data): - for item in data: - if 'page' in item and isinstance(item['page'], str): - try: - item['page'] = int(item['page']) - except ValueError: - # Keep original value if conversion fails - pass - return data - - -def add_node_text(node, pdf_pages): - if isinstance(node, dict): - start_page = node.get('start_index') - end_page = node.get('end_index') - node['text'] = get_text_of_pdf_pages(pdf_pages, start_page, end_page) - if 'nodes' in node: - add_node_text(node['nodes'], pdf_pages) - elif isinstance(node, list): - for index in range(len(node)): - add_node_text(node[index], pdf_pages) - return - - -def add_node_text_with_labels(node, pdf_pages): - if isinstance(node, dict): - start_page = node.get('start_index') - end_page = node.get('end_index') - node['text'] = get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page) - if 'nodes' in node: - add_node_text_with_labels(node['nodes'], pdf_pages) - elif isinstance(node, list): - for index in range(len(node)): - add_node_text_with_labels(node[index], pdf_pages) - return - - -async def generate_node_summary(node, model=None): - prompt = f"""You are given a part of a document, your task is to generate a description of the partial document about what are main points covered in the partial document. - - Partial Document Text: {node['text']} - - Directly return the description, do not include any other text. - """ - response = await llm_acompletion(model, prompt) - return response - - -async def generate_summaries_for_structure(structure, model=None): - nodes = structure_to_list(structure) - tasks = [generate_node_summary(node, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) - - for node, summary in zip(nodes, summaries): - node['summary'] = summary - return structure - - -def create_clean_structure_for_description(structure): - """ - Create a clean structure for document description generation, - excluding unnecessary fields like 'text'. - """ - if isinstance(structure, dict): - clean_node = {} - # Only include essential fields for description - for key in ['title', 'node_id', 'summary', 'prefix_summary']: - if key in structure: - clean_node[key] = structure[key] - - # Recursively process child nodes - if 'nodes' in structure and structure['nodes']: - clean_node['nodes'] = create_clean_structure_for_description(structure['nodes']) - - return clean_node - elif isinstance(structure, list): - return [create_clean_structure_for_description(item) for item in structure] - else: - return structure - - -def generate_doc_description(structure, model=None): - prompt = f"""Your are an expert in generating descriptions for a document. - You are given a structure of a document. Your task is to generate a one-sentence description for the document, which makes it easy to distinguish the document from other documents. - - Document Structure: {structure} - - Directly return the description, do not include any other text. - """ - response = llm_completion(model, prompt) - return response - - -def reorder_dict(data, key_order): - if not key_order: - return data - return {key: data[key] for key in key_order if key in data} - - -def format_structure(structure, order=None): - if not order: - return structure - if isinstance(structure, dict): - if 'nodes' in structure: - structure['nodes'] = format_structure(structure['nodes'], order) - if not structure.get('nodes'): - structure.pop('nodes', None) - structure = reorder_dict(structure, order) - elif isinstance(structure, list): - structure = [format_structure(item, order) for item in structure] - return structure - - -class ConfigLoader: - def __init__(self, default_path: str = None): - if default_path is None: - default_path = Path(__file__).parent / "config.yaml" - self._default_dict = self._load_yaml(default_path) - - @staticmethod - def _load_yaml(path): - with open(path, "r", encoding="utf-8") as f: - return yaml.safe_load(f) or {} - - def _validate_keys(self, user_dict): - unknown_keys = set(user_dict) - set(self._default_dict) - if unknown_keys: - raise ValueError(f"Unknown config keys: {unknown_keys}") - - def load(self, user_opt=None) -> config: - """ - Load the configuration, merging user options with default values. - """ - if user_opt is None: - user_dict = {} - elif isinstance(user_opt, config): - user_dict = vars(user_opt) - elif isinstance(user_opt, dict): - user_dict = user_opt - else: - raise TypeError("user_opt must be dict, config(SimpleNamespace) or None") - - self._validate_keys(user_dict) - merged = {**self._default_dict, **user_dict} - return config(**merged) - -def create_node_mapping(tree, include_page_ranges=False, max_page=None): - """Create a mapping of node_id to node for quick lookup. - - The optional page-range arguments are kept for compatibility with the - pageindex 0.2.x SDK utility API. - """ - def get_all_nodes(nodes): - if isinstance(nodes, dict): - return [nodes] + [ - child_node - for child in nodes.get('nodes', []) - for child_node in get_all_nodes(child) - ] - elif isinstance(nodes, list): - return [ - child_node - for item in nodes - for child_node in get_all_nodes(item) - ] - return [] - - all_nodes = get_all_nodes(tree) - - if not include_page_ranges: - return {node["node_id"]: node for node in all_nodes if node.get("node_id")} - - mapping = {} - for i, node in enumerate(all_nodes): - if not node.get("node_id"): - continue - start_page = node.get("page_index", node.get("start_index")) - if node.get("end_index") is not None: - end_page = node.get("end_index") - elif i + 1 < len(all_nodes): - next_node = all_nodes[i + 1] - end_page = next_node.get("page_index", next_node.get("start_index")) - else: - end_page = max_page - - mapping[node["node_id"]] = { - "node": node, - "start_index": start_page, - "end_index": end_page, - } - - return mapping - -def print_tree(tree, exclude_fields=None, indent=None): - if exclude_fields is None: - exclude_fields = ['text', 'page_index'] - if isinstance(exclude_fields, int): - indent = exclude_fields - exclude_fields = None - if indent is None and exclude_fields is not None: - cleaned_tree = remove_fields(copy.deepcopy(tree), exclude_fields, max_len=40) - pprint(cleaned_tree, sort_dicts=False, width=100) - return - - indent = indent or 0 - for node in tree: - summary = node.get('summary') or node.get('prefix_summary', '') - summary_str = f" — {summary[:60]}..." if summary else "" - print(' ' * indent + f"[{node.get('node_id', '?')}] {node.get('title', '')}{summary_str}") - if node.get('nodes'): - print_tree(node['nodes'], exclude_fields=exclude_fields, indent=indent + 1) - -def print_wrapped(text, width=100): - for line in text.splitlines(): - print(textwrap.fill(line, width=width)) +# pageindex/utils.py +# Deprecation shim. The indexing utilities now live in pageindex/index/utils.py, +# which is the single source of truth. This module re-exports them so legacy +# imports (`from pageindex.utils import ...`) keep working. +import warnings + +warnings.warn( + "pageindex.utils has moved to pageindex.index.utils; importing it from the " + "top level is deprecated and will be removed in a future release.", + PendingDeprecationWarning, + stacklevel=2, +) + +from .index.utils import * # noqa: F401,F403,E402 diff --git a/tests/test_legacy_shims.py b/tests/test_legacy_shims.py new file mode 100644 index 0000000..4ebb3fd --- /dev/null +++ b/tests/test_legacy_shims.py @@ -0,0 +1,83 @@ +"""The top-level pageindex.page_index / .page_index_md / .utils modules are +now deprecation shims over the canonical pageindex.index.* modules. These +tests pin the compatibility contract.""" +import asyncio +import importlib +import warnings + +import pytest + + +def test_plain_import_pageindex_does_not_warn(): + # `import pageindex` must not route through the deprecation shims. + with warnings.catch_warnings(): + warnings.simplefilter("error", PendingDeprecationWarning) + importlib.import_module("pageindex") + + +@pytest.mark.parametrize("mod", [ + "pageindex.utils", + "pageindex.page_index", + "pageindex.page_index_md", +]) +def test_legacy_submodule_import_warns(mod): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + importlib.reload(importlib.import_module(mod)) + assert any(issubclass(w.category, PendingDeprecationWarning) for w in caught) + + +def test_legacy_symbols_resolve_through_shims(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + from pageindex.utils import ( # noqa: F401 + get_page_tokens, ConfigLoader, convert_page_to_int, + get_leaf_nodes, remove_fields, + ) + from pageindex.page_index import page_index, page_index_main # noqa: F401 + from pageindex.page_index_md import md_to_tree # noqa: F401 + + +def test_canonical_and_shim_share_one_implementation(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + import pageindex.utils as shim + import pageindex.index.utils as canonical + # Same function object -> a single source of truth (no divergence possible). + assert shim.get_leaf_nodes is canonical.get_leaf_nodes + assert shim.get_page_tokens is canonical.get_page_tokens + + +def test_get_leaf_nodes_has_331_fix(): + """Canonical get_leaf_nodes must use .get('nodes'); clean_node deletes the + key on leaf nodes so [...]['nodes'] would KeyError (issue #330).""" + from pageindex.index.utils import get_leaf_nodes + # A leaf node with the 'nodes' key deleted (as clean_node leaves it). + leaves = get_leaf_nodes({"title": "Leaf", "start_index": 1, "end_index": 2}) + assert leaves == [{"title": "Leaf", "start_index": 1, "end_index": 2}] + + +def test_configloader_no_longer_needs_config_yaml(): + """config.yaml was removed; ConfigLoader must build defaults from IndexConfig.""" + from pageindex.index.utils import ConfigLoader + cfg = ConfigLoader().load({"model": "gpt-5.4"}) + assert cfg.model == "gpt-5.4" + assert cfg.if_add_node_summary is True # IndexConfig default + with pytest.raises(ValueError, match="Unknown config keys"): + ConfigLoader().load({"nope": 1}) + + +def test_md_to_tree_shim_coerces_yes_no_strings(monkeypatch): + """Canonical md_to_tree takes booleans; the shim must coerce legacy + 'yes'/'no' strings so a bare 'no' doesn't read as truthy True.""" + import pageindex.page_index_md as shim + captured = {} + + async def fake(*args, **kwargs): + captured.update(kwargs) + return {"ok": True} + + monkeypatch.setattr(shim, "_md_to_tree", fake) + asyncio.run(shim.md_to_tree(md_path="x.md", if_add_node_summary="no", if_add_node_id="yes")) + assert captured["if_add_node_summary"] is False + assert captured["if_add_node_id"] is True