Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 34 additions & 8 deletions pageindex/page_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -598,6 +598,34 @@ def remove_first_physical_index_section(text):
return text.replace(match.group(0), '', 1)
return text

# Shared title guidance for no-TOC structure extraction (issues #340/#341).
_TOC_TITLE_INSTRUCTION = """
For the title:
- Prefer a short natural heading from the text when one exists (only fix spacing).
- If there is no short heading (e.g. a numbered list item or a full sentence), synthesize a concise title of at most 12 words that captures the topic.
- Do NOT copy an entire long sentence or paragraph into the title field.
"""

MAX_TOC_TITLE_CHARS = 120


def clamp_toc_titles(toc, max_chars=MAX_TOC_TITLE_CHARS):
"""Truncate oversized TOC titles at a word boundary (safety net for #341)."""
if not isinstance(toc, list):
return toc
for item in toc:
if not isinstance(item, dict):
continue
title = item.get('title')
if not isinstance(title, str) or len(title) <= max_chars:
continue
truncated = title[:max_chars].rsplit(' ', 1)[0].rstrip()
if not truncated:
truncated = title[:max_chars]
item['title'] = truncated + '...'
return toc


### add verify completeness
def generate_toc_continue(toc_content, part, model=None):
print('start generate_toc_continue')
Expand All @@ -607,9 +635,7 @@ def generate_toc_continue(toc_content, part, model=None):
Your task is to continue the tree structure from the previous part to include the current part.

The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.

For the title, you need to extract the original title from the text, only fix the space inconsistency.

""" + _TOC_TITLE_INSTRUCTION + """
The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X. \

For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format.
Expand All @@ -618,7 +644,7 @@ def generate_toc_continue(toc_content, part, model=None):
[
{
"structure": <structure index, "x.x.x"> (string),
"title": <title of the section, keep the original title>,
"title": <short section title>,
"physical_index": "<physical_index_X> (keep the format)"
},
...
Expand All @@ -645,9 +671,7 @@ def generate_toc_init(part, model=None):
You are an expert in extracting hierarchical tree structure, your task is to generate the tree structure of the document.

The structure variable is the numeric system which represents the index of the hierarchy section in the table of contents. For example, the first section has structure index 1, the first subsection has structure index 1.1, the second subsection has structure index 1.2, etc.

For the title, you need to extract the original title from the text, only fix the space inconsistency.

""" + _TOC_TITLE_INSTRUCTION + """
The provided text contains tags like <physical_index_X> and <physical_index_X> to indicate the start and end of page X.

For the physical_index, you need to extract the physical index of the start of the section from the text. Keep the <physical_index_X> format.
Expand All @@ -656,7 +680,7 @@ def generate_toc_init(part, model=None):
[
{{
"structure": <structure index, "x.x.x"> (string),
"title": <title of the section, keep the original title>,
"title": <short section title>,
"physical_index": "<physical_index_X> (keep the format)"
}},

Expand Down Expand Up @@ -718,6 +742,8 @@ def process_no_toc(page_list, start_index=1, model=None, logger=None):

toc_with_page_number = convert_physical_index_to_int(toc_with_page_number)
logger.info(f'convert_physical_index_to_int: {toc_with_page_number}')
toc_with_page_number = clamp_toc_titles(toc_with_page_number)
logger.info(f'clamp_toc_titles: {toc_with_page_number}')

return toc_with_page_number

Expand Down
134 changes: 129 additions & 5 deletions pageindex/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -557,12 +557,100 @@ def add_node_text(node, pdf_pages):
node['text'] = get_text_of_pdf_pages(pdf_pages, start_page, end_page)
if 'nodes' in node:
add_node_text(node['nodes'], pdf_pages)
refine_overlapping_node_text(node)
elif isinstance(node, list):
for index in range(len(node)):
add_node_text(node[index], pdf_pages)
refine_overlapping_sibling_texts(node)
return


def _find_title_anchor(page_text, title):
"""Return the start offset of title (or a prefix) in page_text, or -1."""
if not page_text or not title:
return -1
haystack = page_text.lower()
needle = title.strip()
for length in (len(needle), 80, 60, 40, 20):
fragment = needle[:length].strip()
if len(fragment) < 8:
break
idx = haystack.find(fragment.lower())
if idx >= 0:
return idx
return -1


def partition_text_by_titles(page_text, titles):
"""
Split page_text into exclusive slices for each title using first-occurrence anchors.
Returns None when anchors cannot be resolved reliably.
"""
if not page_text or len(titles) < 2:
return None

anchors = [_find_title_anchor(page_text, t) for t in titles]
if sum(1 for a in anchors if a >= 0) < 2:
return None

# Build ordered cuts for titles we found; unknown titles keep full text later.
indexed = [(i, a) for i, a in enumerate(anchors) if a >= 0]
indexed.sort(key=lambda x: x[1])

slices = [page_text] * len(titles)
for order, (idx, start) in enumerate(indexed):
end = indexed[order + 1][1] if order + 1 < len(indexed) else len(page_text)
if end < start:
continue
slices[idx] = page_text[start:end].strip()
return slices


def refine_overlapping_sibling_texts(siblings):
"""When siblings share identical whole-page text, partition by title anchors."""
if not isinstance(siblings, list) or len(siblings) < 2:
return
texts = [s.get('text') for s in siblings]
if not texts or any(t is None for t in texts):
return
if len(set(texts)) != 1 or not texts[0]:
return
titles = [s.get('title', '') for s in siblings]
slices = partition_text_by_titles(texts[0], titles)
if not slices:
return
for sibling, text in zip(siblings, slices):
sibling['text'] = text


def refine_overlapping_node_text(parent):
"""Partition same-page children under a parent and shrink parent preface text."""
children = parent.get('nodes') or []
if len(children) < 2:
return

parent_text = parent.get('text') or ''
child_texts = [c.get('text') for c in children]
shared_with_parent = (
len(set(child_texts)) == 1
and bool(child_texts[0])
and (not parent_text or parent_text == child_texts[0])
)

refine_overlapping_sibling_texts(children)

if not shared_with_parent or not parent_text:
return

first_title = (children[0].get('title') or '').strip()
idx = _find_title_anchor(parent_text, first_title)
if idx > 0:
parent['text'] = parent_text[:idx].strip()
elif idx == 0:
# No exclusive preface; leave empty so summarization uses child titles.
parent['text'] = ''


def add_node_text_with_labels(node, pdf_pages):
if isinstance(node, dict):
start_page = node.get('start_index')
Expand All @@ -587,13 +675,49 @@ async def generate_node_summary(node, model=None):
return response


def _container_summary_from_children(node):
titles = [c.get('title', '').strip() for c in (node.get('nodes') or []) if c.get('title')]
section = (node.get('title') or '').strip() or 'Section'
if titles:
return f'{section} covering: ' + '; '.join(titles)
return section


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

# Pre-assign container summaries when parent text is empty or duplicates children.
pending = []
for node in nodes:
children = node.get('nodes') or []
text = node.get('text') or ''
if children and (
not text.strip()
or any(text == (c.get('text') or '') for c in children)
):
node['summary'] = _container_summary_from_children(node)
else:
pending.append(node)

# Deduplicate LLM calls for identical text (same-page siblings before partition, etc.).
unique_texts = []
text_to_nodes = {}
for node in pending:
key = node.get('text') or ''
if key not in text_to_nodes:
text_to_nodes[key] = []
unique_texts.append(key)
text_to_nodes[key].append(node)

if unique_texts:
reps = [text_to_nodes[t][0] for t in unique_texts]
summaries = await asyncio.gather(
*[generate_node_summary(node, model=model) for node in reps]
)
for key, summary in zip(unique_texts, summaries):
for node in text_to_nodes[key]:
node['summary'] = summary

return structure


Expand Down
158 changes: 158 additions & 0 deletions tests/test_toc_title_summary.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
import asyncio
import importlib
import inspect
import unittest
from unittest.mock import AsyncMock, patch

from pageindex.page_index import clamp_toc_titles
from pageindex.utils import (
add_node_text,
generate_summaries_for_structure,
partition_text_by_titles,
)

page_index_mod = importlib.import_module("pageindex.page_index")


class ClampTocTitlesTest(unittest.TestCase):
def test_truncates_long_titles_at_word_boundary(self):
long_title = (
"1.1 The iOS version of the software requires that the device's "
"running memory is not less than 2GB and many more words follow here "
"to exceed the maximum title character limit substantially"
)
toc = [{"title": long_title, "structure": "1.1"}]

clamp_toc_titles(toc, max_chars=80)

self.assertLessEqual(len(toc[0]["title"]), 83) # 80 + "..."
self.assertTrue(toc[0]["title"].endswith("..."))
self.assertNotIn("substantially", toc[0]["title"])

def test_leaves_short_titles_unchanged(self):
toc = [{"title": "Software Installation", "structure": "1"}]
clamp_toc_titles(toc)
self.assertEqual(toc[0]["title"], "Software Installation")


class TocTitlePromptTest(unittest.TestCase):
def test_toc_prompts_ask_for_concise_titles(self):
instruction = page_index_mod._TOC_TITLE_INSTRUCTION
self.assertIn("synthesize a concise title", instruction)
self.assertIn("Do NOT copy an entire long sentence", instruction)

init_src = inspect.getsource(page_index_mod.generate_toc_init)
continue_src = inspect.getsource(page_index_mod.generate_toc_continue)
self.assertIn("_TOC_TITLE_INSTRUCTION", init_src)
self.assertIn("_TOC_TITLE_INSTRUCTION", continue_src)
self.assertNotIn("only fix the space inconsistency", init_src)
self.assertNotIn("only fix the space inconsistency", continue_src)


class PartitionTextByTitlesTest(unittest.TestCase):
def test_partitions_same_page_siblings_by_title_anchor(self):
page = (
"Software Installation\n"
"1.1 Memory requirements need 2GB RAM.\n"
"1.2 Storage requirements need 10GB free space.\n"
"1.3 Network requirements need Wi-Fi.\n"
)
titles = [
"1.1 Memory requirements",
"1.2 Storage requirements",
"1.3 Network requirements",
]

slices = partition_text_by_titles(page, titles)

self.assertIsNotNone(slices)
self.assertEqual(len(slices), 3)
self.assertIn("2GB RAM", slices[0])
self.assertNotIn("10GB", slices[0])
self.assertIn("10GB", slices[1])
self.assertNotIn("Wi-Fi", slices[1])
self.assertIn("Wi-Fi", slices[2])


class AddNodeTextOverlapTest(unittest.TestCase):
def test_same_page_children_get_distinct_text(self):
page_text = (
"Software Installation preface.\n"
"1.1 Memory requirements need 2GB RAM.\n"
"1.2 Storage requirements need 10GB free space.\n"
)
pdf_pages = [[page_text]]
tree = [{
"title": "Software Installation",
"start_index": 1,
"end_index": 1,
"nodes": [
{
"title": "1.1 Memory requirements",
"start_index": 1,
"end_index": 1,
},
{
"title": "1.2 Storage requirements",
"start_index": 1,
"end_index": 1,
},
],
}]

add_node_text(tree, pdf_pages)

parent, child1, child2 = tree[0], tree[0]["nodes"][0], tree[0]["nodes"][1]
self.assertIn("preface", parent["text"])
self.assertNotEqual(child1["text"], child2["text"])
self.assertIn("2GB RAM", child1["text"])
self.assertIn("10GB", child2["text"])


class GenerateSummariesDedupeTest(unittest.TestCase):
def test_reuses_summary_for_identical_text(self):
structure = [
{"title": "A", "text": "same text", "node_id": "1"},
{"title": "B", "text": "same text", "node_id": "2"},
{"title": "C", "text": "other text", "node_id": "3"},
]

with patch(
"pageindex.utils.generate_node_summary",
new_callable=AsyncMock,
side_effect=["summary-same", "summary-other"],
) as mock_summary:
asyncio.run(generate_summaries_for_structure(structure, model="dummy"))

self.assertEqual(mock_summary.await_count, 2)
self.assertEqual(structure[0]["summary"], "summary-same")
self.assertEqual(structure[1]["summary"], "summary-same")
self.assertEqual(structure[2]["summary"], "summary-other")

def test_container_with_duplicate_child_text_skips_llm(self):
structure = [{
"title": "Software Installation",
"text": "dup",
"nodes": [
{"title": "1.1 Memory", "text": "dup"},
{"title": "1.2 Storage", "text": "dup"},
],
}]

with patch(
"pageindex.utils.generate_node_summary",
new_callable=AsyncMock,
return_value="child-summary",
) as mock_summary:
asyncio.run(generate_summaries_for_structure(structure, model="dummy"))

# Parent gets container summary; one LLM call for shared child text.
self.assertEqual(mock_summary.await_count, 1)
self.assertIn("covering:", structure[0]["summary"])
self.assertIn("1.1 Memory", structure[0]["summary"])
self.assertEqual(structure[0]["nodes"][0]["summary"], "child-summary")
self.assertEqual(structure[0]["nodes"][1]["summary"], "child-summary")


if __name__ == "__main__":
unittest.main()