apunkt/PageIndex
1
0
Fork 0
mirror of https://github.com/VectifyAI/PageIndex.git synced 2026-07-15 21:11:05 +02:00
Code Issues Projects Releases Packages Wiki Activity

refactor(index): dedupe the copied indexing pipeline behind deprecation shims

Browse source 
The new SDK copied the legacy indexing pipeline into pageindex/index/
instead of moving it, leaving two divergent copies of page_index.py /
page_index_md.py / utils.py. They had already drifted (the legacy copy
still compared IndexConfig booleans against 'yes' — a separate fix),
and every pipeline change had to be applied twice.

Make pageindex/index/ the single source of truth (same pattern as the
LegacyCloudAPI shim for the 0.2.x cloud SDK):

- pageindex/index/utils.py absorbs the 27 legacy-only helpers/classes
  (get_page_tokens, convert_page_to_int, ConfigLoader, PDF text helpers,
  ...) so it's the sole utils module. Reconciled the diverged funcs:
  kept the modern versions, backported the #331 get_leaf_nodes .get()
  fix, and restored remove_fields' max_len parameter (superset).
- index/page_index*.py now import `from .utils import *`;
  index/legacy_utils.py (a re-export of the old top-level utils) deleted.
- Top-level page_index.py / page_index_md.py / utils.py become thin
  re-export shims that emit PendingDeprecationWarning. The md_to_tree
  shim coerces legacy 'yes'/'no' string flags to bool (the canonical
  version is boolean-typed).
- ConfigLoader no longer reads the deleted config.yaml; it builds
  defaults from IndexConfig (was an unconditional FileNotFoundError).
- __init__.py and retrieve.py import from pageindex.index.* directly so
  `import pageindex` does not trip the shims.

Adds tests/test_legacy_shims.py pinning the contract: clean top-level
import doesn't warn, legacy submodule imports warn, symbols still
resolve, shim and canonical share one implementation, the #331 fix and
ConfigLoader-without-yaml both hold, and the md_to_tree coercion works.

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
mountain 2026-07-07 10:41:29 +08:00
parent 956147d864
commit 6c948a332b
11 changed files with 593 additions and 2288 deletions
Download patch file Download diff file Expand all files Collapse all files

5
.gitignore vendored
View file

@ -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

8
pageindex/__init__.py
View file

@ -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

2
pageindex/index/legacy_utils.py
View file

@ -1,2 +0,0 @@
# Re-export from the original utils.py for backward compatibility
from ..utils import *

2
pageindex/index/page_index.py
View file

@ -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

5
pageindex/index/page_index_md.py
View file

@ -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')

438
pageindex/index/utils.py
View file

@ -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"<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 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 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('<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_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))

1173
pageindex/page_index.py
View file

File diff suppressed because it is too large Load diff

370
pageindex/page_index_md.py
View file

@ -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}")
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)

4
pageindex/retrieve.py
View file

@ -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 ──────────────────────────────────────────────────────────────────

791
pageindex/utils.py
View file

@ -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

83
tests/test_legacy_shims.py Normal file
View file

@ -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