diff --git a/assets/generate.py b/assets/generate.py index e32193f..34e70c9 100644 --- a/assets/generate.py +++ b/assets/generate.py @@ -161,11 +161,102 @@ def create_search_targets(): doc.close() +def create_large_toc(): + """PDF with a very large multi-level TOC to test auto-trimming.""" + doc = fitz.open() + toc = [] + page_num = 0 + + for part in range(1, 6): + page_num += 1 + page = doc.new_page(width=595, height=842) + page.insert_text((72, 80), f"Part {part}", fontsize=18) + toc.append([1, f"Part {part}", page_num]) + + for chapter in range(1, 9): + page_num += 1 + page = doc.new_page(width=595, height=842) + title = f"Chapter {part}.{chapter}" + page.insert_text((72, 80), title, fontsize=14) + toc.append([2, title, page_num]) + + for section in range(1, 5): + page_num += 1 + page = doc.new_page(width=595, height=842) + title = f"Section {part}.{chapter}.{section}" + page.insert_text((72, 80), title, fontsize=12) + toc.append([3, title, page_num]) + + doc.set_toc(toc) + doc.save(ASSETS / "large_toc.pdf") + doc.close() + + +def create_mega_toc(): + """PDF with 400+ pages and 4-level deep TOC for comprehensive auto-trim testing. + + Structure: + 2 volumes (L1) × 8 chapters (L2) × 8 sections (L3) × 2 subsections (L4) + + Entry counts: + Level 1: 2 + Level 2: 16 (cumulative ≤L2: 18) + Level 3: 128 (cumulative ≤L3: 146) + Level 4: 256 (cumulative ≤L4: 402) + + Full-TOC auto-trim (threshold 100): + ≤L2 = 18 fits → best_level=2, total=18, auto_trimmed_to_level=2 + + Volume 1 subtree (parent="Volume 1"): + 8 chapters + 64 sections + 128 subsections = 200 children + ≤L2 = 8, ≤L3 = 72 fits → best_level=3, total=72, auto_trimmed_to_level=2 (relative) + + Chapter 1.1 subtree (parent="Chapter 1.1"): + 8 sections + 16 subsections = 24 children (< 100, no auto-trim) + """ + doc = fitz.open() + toc = [] + page_num = 0 + + for vol in range(1, 3): + page_num += 1 + page = doc.new_page(width=595, height=842) + page.insert_text((72, 80), f"Volume {vol}", fontsize=18) + toc.append([1, f"Volume {vol}", page_num]) + + for ch in range(1, 9): + page_num += 1 + page = doc.new_page(width=595, height=842) + title = f"Chapter {vol}.{ch}" + page.insert_text((72, 80), title, fontsize=14) + toc.append([2, title, page_num]) + + for sec in range(1, 9): + page_num += 1 + page = doc.new_page(width=595, height=842) + title = f"Section {vol}.{ch}.{sec}" + page.insert_text((72, 80), title, fontsize=12) + toc.append([3, title, page_num]) + + for sub in range(1, 3): + page_num += 1 + page = doc.new_page(width=595, height=842) + title = f"Subsection {vol}.{ch}.{sec}.{sub}" + page.insert_text((72, 80), title, fontsize=11) + toc.append([4, title, page_num]) + + doc.set_toc(toc) + doc.save(ASSETS / "mega_toc.pdf") + doc.close() + + if __name__ == "__main__": create_basic() create_with_diagrams() create_nested_toc() create_search_targets() + create_large_toc() + create_mega_toc() print(f"Generated test PDFs in {ASSETS}") for f in sorted(ASSETS.glob("*.pdf")): print(f" {f.name}") diff --git a/src/pdf_mcp/server.py b/src/pdf_mcp/server.py index 545ba41..b5ff010 100644 --- a/src/pdf_mcp/server.py +++ b/src/pdf_mcp/server.py @@ -13,14 +13,15 @@ mcp = FastMCP( ) HEADER_FOOTER_MARGIN_PTS = 50 +TOC_AUTO_TRIM_THRESHOLD = 100 @mcp.tool() def get_pdf_info(filename: str) -> dict: """Get metadata and basic info about a PDF file. - Returns page count, title, author, subject, creator, producer, - creation/modification dates, and encryption info. + Args: + filename: Path to a PDF file. """ doc = fitz.open(filename) info = { @@ -34,16 +35,94 @@ def get_pdf_info(filename: str) -> dict: @mcp.tool() -def get_table_of_contents(filename: str) -> list[dict]: +def get_table_of_contents( + filename: str, + parent: str | None = None, + max_level: int | None = None, +) -> dict: """Get the table of contents (bookmarks/outline) from a PDF. - Returns a list of entries with level, title, and page number. - These correspond to the sections shown in a PDF reader's sidebar. + When max_level is not specified and the TOC is very large, the depth is automatically + reduced to keep the response manageable. The response will include an + ``auto_trimmed_to_level`` field and a hint when this happens. + + Args: + filename: Path to a PDF file. + parent: Return only children of the entry whose title contains this string + (case-insensitive). For example, pass a chapter title to get its sections. + max_level: Only include entries up to this depth (1=chapters only, 2=sections, etc.). + When used with parent, levels are relative: 1 means direct children only, + 2 means children and grandchildren, etc. """ doc = fitz.open(filename) toc = doc.get_toc() doc.close() - return [{"level": level, "title": title, "page": page} for level, title, page in toc] + + # Track whether max_level was explicitly requested + explicit_max_level = max_level is not None + base_level = 0 + + if parent is not None: + parent_lower = parent.lower() + # Find the first entry whose title matches + parent_idx = None + parent_level = None + for i, (level, title, _page) in enumerate(toc): + if parent_lower in title.lower(): + parent_idx = i + parent_level = level + break + if parent_idx is None: + return { + "error": f"No TOC entry matching '{parent}'", + "total_entries": 0, + "entries": [], + } + # Collect all subsequent entries until we hit an entry at the same or higher level + children = [] + for level, title, page in toc[parent_idx + 1 :]: + if level <= parent_level: + break + children.append((level, title, page)) + toc = children + base_level = parent_level + # Make max_level relative to parent + if max_level is not None: + abs_max_level = parent_level + max_level + toc = [entry for entry in toc if entry[0] <= abs_max_level] + elif max_level is not None: + toc = [entry for entry in toc if entry[0] <= max_level] + + # Auto-trim depth when max_level was not explicitly requested and the TOC is large + applied_max_level = None + untrimmed_total = len(toc) + if not explicit_max_level and len(toc) > TOC_AUTO_TRIM_THRESHOLD: + levels_present = sorted(set(entry[0] for entry in toc)) + # Find the deepest level that keeps entries at or below the threshold. + # If even the shallowest level exceeds it, use that (can't go shallower). + best_level = levels_present[0] + for try_level in levels_present: + count = sum(1 for entry in toc if entry[0] <= try_level) + if count <= TOC_AUTO_TRIM_THRESHOLD: + best_level = try_level + else: + break + toc = [entry for entry in toc if entry[0] <= best_level] + applied_max_level = best_level - base_level + + entries = [ + {"level": level, "title": title, "page": page} for level, title, page in toc + ] + + result = {"total_entries": len(entries), "entries": entries} + if applied_max_level is not None: + result["auto_trimmed_to_level"] = applied_max_level + result["hint"] = ( + f"TOC has {untrimmed_total} entries; trimmed to level {applied_max_level} " + f"({len(entries)} entries). Use max_level to request a specific depth, " + "or use parent to narrow to a specific section." + ) + return result @mcp.tool() @@ -57,7 +136,7 @@ def get_page_text( """Extract text content from one or more PDF pages. Args: - filename: Path to the PDF file. + filename: Path to a PDF file. start_page: First page number (1-indexed, inclusive). end_page: Last page number (1-indexed, inclusive). Defaults to start_page. format: Output format. "json" returns structured page data with block/line/span @@ -132,7 +211,7 @@ def get_page_image( """Render a single PDF page as a PNG image. Args: - filename: Path to the PDF file. + filename: Path to a PDF file. page: Page number (1-indexed). dpi: Image resolution. Default 150 (good balance of readability and size). output: "base64" returns the image inline as MCP image content. @@ -166,7 +245,7 @@ def search_text( Returns a list of hits with page number and surrounding context. Args: - filename: Path to the PDF file. + filename: Path to a PDF file. query: Text to search for. context_chars: Characters of context to include around each hit. Default 100. """ @@ -185,10 +264,12 @@ def search_text( break ctx_start = max(0, idx - context_chars) ctx_end = min(len(text), idx + len(query) + context_chars) - hits.append({ - "page": page_idx + 1, - "context": text[ctx_start:ctx_end], - }) + hits.append( + { + "page": page_idx + 1, + "context": text[ctx_start:ctx_end], + } + ) pos = idx + 1 doc.close() diff --git a/tests/conftest.py b/tests/conftest.py index 72ae6dd..271f3ad 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,23 +1,36 @@ """Shared fixtures for pdf-mcp tests.""" +import asyncio import subprocess import sys from pathlib import Path import pytest +from pdf_mcp.server import mcp + ASSETS = Path(__file__).parent.parent / "assets" def pytest_configure(config): """Generate test PDFs before the test session if they don't exist.""" - expected = ["basic.pdf", "diagrams.pdf", "nested_toc.pdf", "search_targets.pdf"] + expected = ["basic.pdf", "diagrams.pdf", "nested_toc.pdf", "search_targets.pdf", "large_toc.pdf", "mega_toc.pdf"] if not all((ASSETS / f).exists() for f in expected): subprocess.check_call( [sys.executable, str(ASSETS / "generate.py")], ) +@pytest.fixture +def call_tool(): + """Call an MCP tool by name, going through the real FastMCP call path.""" + def _call(tool_name: str, **kwargs): + tools = mcp._tool_manager.list_tools() + tool = next(t for t in tools if t.name == tool_name) + return asyncio.run(tool.run(kwargs)) + return _call + + @pytest.fixture def basic_pdf() -> str: return str(ASSETS / "basic.pdf") @@ -36,3 +49,13 @@ def nested_toc_pdf() -> str: @pytest.fixture def search_pdf() -> str: return str(ASSETS / "search_targets.pdf") + + +@pytest.fixture +def large_toc_pdf() -> str: + return str(ASSETS / "large_toc.pdf") + + +@pytest.fixture +def mega_toc_pdf() -> str: + return str(ASSETS / "mega_toc.pdf") diff --git a/tests/test_get_page_image.py b/tests/test_get_page_image.py index b7759ae..54125e3 100644 --- a/tests/test_get_page_image.py +++ b/tests/test_get_page_image.py @@ -3,75 +3,70 @@ from pathlib import Path from mcp.server.fastmcp import Image -from pdf_mcp.server import get_page_image - class TestGetPageImageBase64: - def test_returns_image_object(self, basic_pdf): - result = get_page_image(basic_pdf, page=1) + def test_returns_image_object(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1) assert isinstance(result, Image) - def test_image_has_data(self, basic_pdf): - result = get_page_image(basic_pdf, page=1) + def test_image_has_data(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1) assert len(result.data) > 0 - def test_png_signature(self, basic_pdf): - result = get_page_image(basic_pdf, page=1) - # PNG files start with the 8-byte signature + def test_png_signature(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1) assert result.data[:4] == b"\x89PNG" - def test_different_pages(self, basic_pdf): - img1 = get_page_image(basic_pdf, page=1) - img2 = get_page_image(basic_pdf, page=2) + def test_different_pages(self, basic_pdf, call_tool): + img1 = call_tool("get_page_image", filename=basic_pdf, page=1) + img2 = call_tool("get_page_image", filename=basic_pdf, page=2) assert img1.data != img2.data class TestGetPageImageFile: - def test_returns_string_with_path(self, basic_pdf): - result = get_page_image(basic_pdf, page=1, output="file") + def test_returns_string_with_path(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1, output="file") assert isinstance(result, str) assert "Image saved to" in result - def test_file_exists(self, basic_pdf): - result = get_page_image(basic_pdf, page=1, output="file") - # Extract path from "Image saved to /path/to/file.png (WxHpx)" + def test_file_exists(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1, output="file") path = result.split("Image saved to ")[1].split(" (")[0] assert Path(path).exists() os.unlink(path) - def test_file_is_valid_png(self, basic_pdf): - result = get_page_image(basic_pdf, page=1, output="file") + def test_file_is_valid_png(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1, output="file") path = result.split("Image saved to ")[1].split(" (")[0] with open(path, "rb") as f: assert f.read(4) == b"\x89PNG" os.unlink(path) - def test_file_includes_dimensions(self, basic_pdf): - result = get_page_image(basic_pdf, page=1, output="file") - # Should contain something like "(1240x1754px)" + def test_file_includes_dimensions(self, basic_pdf, call_tool): + result = call_tool("get_page_image", filename=basic_pdf, page=1, output="file") assert "px)" in result class TestGetPageImageDpi: - def test_higher_dpi_larger_file(self, basic_pdf): - low = get_page_image(basic_pdf, page=1, dpi=72) - high = get_page_image(basic_pdf, page=1, dpi=300) + def test_higher_dpi_larger_file(self, basic_pdf, call_tool): + low = call_tool("get_page_image", filename=basic_pdf, page=1, dpi=72) + high = call_tool("get_page_image", filename=basic_pdf, page=1, dpi=300) assert len(high.data) > len(low.data) - def test_default_dpi(self, basic_pdf): - default = get_page_image(basic_pdf, page=1) - explicit_150 = get_page_image(basic_pdf, page=1, dpi=150) + def test_default_dpi(self, basic_pdf, call_tool): + default = call_tool("get_page_image", filename=basic_pdf, page=1) + explicit_150 = call_tool("get_page_image", filename=basic_pdf, page=1, dpi=150) assert len(default.data) == len(explicit_150.data) class TestGetPageImageDiagrams: - def test_diagrams_renders(self, diagrams_pdf): - result = get_page_image(diagrams_pdf, page=1) + def test_diagrams_renders(self, diagrams_pdf, call_tool): + result = call_tool("get_page_image", filename=diagrams_pdf, page=1) assert isinstance(result, Image) assert len(result.data) > 0 - def test_diagrams_file_output(self, diagrams_pdf): - result = get_page_image(diagrams_pdf, page=1, output="file") + def test_diagrams_file_output(self, diagrams_pdf, call_tool): + result = call_tool("get_page_image", filename=diagrams_pdf, page=1, output="file") path = result.split("Image saved to ")[1].split(" (")[0] assert Path(path).exists() assert Path(path).stat().st_size > 0 diff --git a/tests/test_get_page_text.py b/tests/test_get_page_text.py index e7c0792..ded31b6 100644 --- a/tests/test_get_page_text.py +++ b/tests/test_get_page_text.py @@ -1,120 +1,116 @@ -from pdf_mcp.server import get_page_text - - class TestGetPageTextJson: - def test_default_format_is_json(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1) + def test_default_format_is_json(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1) assert isinstance(result, list) assert len(result) == 1 assert isinstance(result[0], dict) - def test_json_has_page_number(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=2) + def test_json_has_page_number(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=2) assert result[0]["page_number"] == 2 - def test_json_has_blocks(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1) + def test_json_has_blocks(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1) assert "blocks" in result[0] assert len(result[0]["blocks"]) > 0 - def test_json_page_range(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, end_page=3, format="json") + def test_json_page_range(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, end_page=3, format="json") assert len(result) == 3 assert result[0]["page_number"] == 1 assert result[2]["page_number"] == 3 class TestGetPageTextPlain: - def test_text_format_returns_string(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="text") + def test_text_format_returns_string(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="text") assert isinstance(result, str) - def test_text_contains_content(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="text") + def test_text_contains_content(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="text") assert "Chapter 1" in result assert "introductory material" in result - def test_text_page_range_has_breaks(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, end_page=2, format="text") + def test_text_page_range_has_breaks(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, end_page=2, format="text") assert "--- Page Break ---" in result assert "Chapter 1" in result assert "Chapter 2" in result - def test_text_single_page_no_break(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="text") + def test_text_single_page_no_break(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="text") assert "--- Page Break ---" not in result class TestGetPageTextMarkdown: - def test_markdown_returns_string(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="markdown") + def test_markdown_returns_string(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="markdown") assert isinstance(result, str) - def test_markdown_contains_content(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="markdown") + def test_markdown_contains_content(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="markdown") assert "Chapter 1" in result - def test_markdown_page_range(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, end_page=3, format="markdown") + def test_markdown_page_range(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, end_page=3, format="markdown") assert "Chapter 1" in result assert "Chapter 3" in result class TestGetPageTextHtml: - def test_html_returns_string(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="html") + def test_html_returns_string(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="html") assert isinstance(result, str) - def test_html_has_tags(self, basic_pdf): - result = get_page_text(basic_pdf, start_page=1, format="html") + def test_html_has_tags(self, basic_pdf, call_tool): + result = call_tool("get_page_text", filename=basic_pdf, start_page=1, format="html") assert "