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

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

.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

pageindex/__init__.py

@@ -1,7 +1,9 @@
 # pageindex/__init__.py
-# Upstream exports (backward compatibility)
+# Upstream exports (backward compatibility). Import from the canonical
-from .page_index import *
+# pageindex.index.* modules directly so `import pageindex` does NOT trip the
-from .page_index_md import md_to_tree
+# 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

pageindex/index/legacy_utils.py

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

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
 

pageindex/index/page_index_md.py

@@ -2,10 +2,7 @@ import asyncio
 import json
 import re
 import os
-try:
+from .utils import *
-    from .legacy_utils import *
-except:
-    from legacy_utils import *
 
 async def get_node_summary(node, summary_token_threshold=200, model=None):
     node_text = node.get('text')

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

pageindex/page_index_md.py

@@ -1,342 +1,38 @@
-import asyncio
+# pageindex/page_index_md.py
-import json
+# Deprecation shim. The Markdown indexing pipeline now lives in
-import re
+# pageindex/index/page_index_md.py (the single source of truth). This module
-import os
+# re-exports it so legacy imports keep working.
-try:
+#
-    from .utils import *
+# The canonical md_to_tree takes booleans; legacy callers passed 'yes'/'no'
-except:
+# strings, so the wrapper below coerces them (a bare 'no' is otherwise truthy).
-    from utils import *
+import warnings
 
-async def get_node_summary(node, summary_token_threshold=200, model=None):
+warnings.warn(
-    node_text = node.get('text')
+    "pageindex.page_index_md has moved to pageindex.index.page_index_md; "
-    num_tokens = count_tokens(node_text, model=model)
+    "importing it from the top level is deprecated and will be removed in a "
-    if num_tokens < summary_token_threshold:
+    "future release.",
-        return node_text
+    PendingDeprecationWarning,
-    else:
+    stacklevel=2,
-        return await generate_node_summary(node, model=model)
+)
+
+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):
+def _coerce_bool(value):
-    nodes = structure_to_list(structure)
+    if isinstance(value, str):
-    tasks = [get_node_summary(node, summary_token_threshold=summary_token_threshold, model=model) for node in nodes]
+        return value.strip().lower() in ("yes", "true", "1", "y", "on")
-    summaries = await asyncio.gather(*tasks)
+    return bool(value)
-    
-    for node, summary in zip(nodes, summaries):
-        if not node.get('nodes'):
-            node['summary'] = summary
-        else:
-            node['prefix_summary'] = summary
-    return structure
 
 
-def extract_nodes_from_markdown(markdown_content):
+async def md_to_tree(*args, **kwargs):
-    header_pattern = r'^(#{1,6})\s+(.+)$'
+    """Legacy wrapper: coerce 'yes'/'no' string flags to bool, then delegate."""
-    code_block_pattern = r'^```'
+    for key in _BOOL_PARAMS:
-    node_list = []
+        if key in kwargs:
-    
+            kwargs[key] = _coerce_bool(kwargs[key])
-    lines = markdown_content.split('\n')
+    return await _md_to_tree(*args, **kwargs)
-    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}")

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

pageindex/utils.py

@@ -1,777 +1,14 @@
-import litellm
+# pageindex/utils.py
-import logging
+# Deprecation shim. The indexing utilities now live in pageindex/index/utils.py,
-import os
+# which is the single source of truth. This module re-exports them so legacy
-import textwrap
+# imports (`from pageindex.utils import ...`) keep working.
-from datetime import datetime
+import warnings
-import time
+
-import json
+warnings.warn(
-import PyPDF2
+    "pageindex.utils has moved to pageindex.index.utils; importing it from the "
-import copy
+    "top level is deprecated and will be removed in a future release.",
-import asyncio
+    PendingDeprecationWarning,
-import pymupdf
+    stacklevel=2,
-from io import BytesIO
+)
-from dotenv import load_dotenv
+
-load_dotenv()
+from .index.utils import *  # noqa: F401,F403,E402
-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))

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