mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
fix: CMYK image drop, empty-doc crash, page_index shadowing, sqlite hardening, flaky tests
Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272. - pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without- alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png then raises "unsupported colorspace", silently dropped by the surrounding except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves RGBA untouched). - pipeline.py: detect_strategy([]) (an empty/whitespace-only source file) returned "content_based", routing into the PDF-oriented TOC-detection pipeline -- wasting a real LLM call before raising IndexingError. Empty node lists now route to level_based, whose build_tree_from_levels([]) returns an empty structure instantly with zero LLM calls. - page_index.py (shim): pageindex/__init__.py binds the canonical `page_index` function as the package attribute, but this file is ALSO a real submodule of the same name -- importing it anywhere (import machinery, unconditional) overwrites that attribute with the module object, breaking `from pageindex import page_index; page_index(x)` for the rest of the process. Made the shim module itself callable (delegates to the real function via a ModuleType subclass), so whichever object ends up in that slot is callable regardless of import order. - storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape on a duplicate name (new CollectionAlreadyExistsError); the collections table's CHECK constraint only validated the name's first character (GLOB '*' is a wildcard, not a regex quantifier over the preceding class) -- fixed to validate the whole string, and SQLiteStorage now also validates in Python (it's a public StorageEngine usable directly, bypassing LocalBackend's own check). - tests/test_review_fixes_2.py: two tests used a ContentNode with no `level` set, so build_index took the content_based path and made real (retried, slow, and -- with a valid key -- billable) LLM calls instead of testing the text-stripping logic they claimed to. Mocked out _content_based_pipeline. - retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of the canonical parse_pages/get_pdf_page_content that had already drifted (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical now, so the legacy pageindex.get_page_content path can't silently regress again. - parser/markdown.py: a leading UTF-8 BOM broke first-header detection (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only backtick fences were recognized as code blocks, so a '#'-prefixed line inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading -- recognize both fence styles. - run_pageindex.py: --if-thinning wasn't migrated to the bare-flag + legacy-yes/no convention the other four --if-add-* flags got; bare usage raised an argparse error and it never went through the shared coercion. - types.py: DocumentDetail's `structure` field was inside the class's total=False body, so TypedDict rules made it optional even though every backend always populates it. Split into a required base class. Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped (one pre-existing, unrelated flaky cloud-streaming test). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
parent
b9d021916f
commit
4e6a13576d
18 changed files with 368 additions and 36 deletions
|
|
@ -8,6 +8,11 @@ class CollectionNotFoundError(PageIndexError):
|
|||
pass
|
||||
|
||||
|
||||
class CollectionAlreadyExistsError(PageIndexError):
|
||||
"""Collection already exists (create_collection, not get_or_create)."""
|
||||
pass
|
||||
|
||||
|
||||
class DocumentNotFoundError(PageIndexError):
|
||||
"""Document ID not found."""
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -5,6 +5,13 @@ from ..parser.protocol import ContentNode, ParsedDocument
|
|||
|
||||
def detect_strategy(nodes: list[ContentNode]) -> str:
|
||||
"""Determine which indexing strategy to use based on node data."""
|
||||
if not nodes:
|
||||
# No content at all (e.g. an empty/whitespace-only source file) ->
|
||||
# level_based's build_tree_from_levels([]) returns an empty structure
|
||||
# immediately with zero LLM calls. content_based's TOC-detection
|
||||
# pipeline needs real page content; on an empty page_list it wastes an
|
||||
# LLM call and then still raises, for no benefit.
|
||||
return "level_based"
|
||||
if any(n.level is not None for n in nodes):
|
||||
return "level_based"
|
||||
return "content_based"
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@
|
|||
# pageindex/index/page_index.py (the single source of truth). This module
|
||||
# re-exports it so legacy imports (`from pageindex.page_index import ...`,
|
||||
# `from pageindex import page_index`) keep working.
|
||||
import sys
|
||||
import types
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
|
|
@ -13,3 +15,26 @@ warnings.warn(
|
|||
)
|
||||
|
||||
from .index.page_index import * # noqa: F401,F403,E402
|
||||
|
||||
# pageindex/__init__.py binds the FUNCTION `page_index` as the package
|
||||
# attribute `pageindex.page_index` (`from .index.page_index import *`). But
|
||||
# this file is ALSO a real submodule of the same name — the moment anything,
|
||||
# anywhere in the process, does `import pageindex.page_index` (exactly what
|
||||
# `from pageindex.page_index import X` triggers), Python's import machinery
|
||||
# overwrites that package attribute with THIS module object, clobbering the
|
||||
# function binding. Afterwards `from pageindex import page_index; page_index(x)`
|
||||
# would raise "TypeError: 'module' object is not callable" — silently, and
|
||||
# depending entirely on whether this submodule happened to be imported yet.
|
||||
#
|
||||
# Fix: make this module itself callable, delegating to the real function, so
|
||||
# whichever object ends up sitting in the `pageindex.page_index` slot — the
|
||||
# function or this module — is callable either way. Both `from pageindex.page_index
|
||||
# import page_index_main` (module attribute access) and
|
||||
# `from pageindex import page_index; page_index(x)` (call) keep working
|
||||
# regardless of import order.
|
||||
class _CallableModule(types.ModuleType):
|
||||
def __call__(self, *args, **kwargs):
|
||||
return page_index(*args, **kwargs)
|
||||
|
||||
|
||||
sys.modules[__name__].__class__ = _CallableModule
|
||||
|
|
|
|||
|
|
@ -12,7 +12,12 @@ class MarkdownParser:
|
|||
path = Path(file_path)
|
||||
model = kwargs.get("model")
|
||||
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
# utf-8-sig strips a leading BOM if present (common from Windows
|
||||
# editors/exporters) and is otherwise identical to plain utf-8. Without
|
||||
# it, a BOM-prefixed first line fails the header regex below (the BOM
|
||||
# isn't whitespace, so .strip() doesn't remove it), misclassifying the
|
||||
# document's first heading as unrecognized preamble text.
|
||||
with open(path, "r", encoding="utf-8-sig") as f:
|
||||
content = f.read()
|
||||
|
||||
lines = content.split("\n")
|
||||
|
|
@ -23,7 +28,10 @@ class MarkdownParser:
|
|||
|
||||
def _extract_headers(self, lines: list[str]) -> list[dict]:
|
||||
header_pattern = r"^(#{1,6})\s+(.+)$"
|
||||
code_block_pattern = r"^```"
|
||||
# CommonMark allows both backtick and tilde fences; only recognizing
|
||||
# backticks let a '#'-prefixed line inside a ~~~-fenced block (e.g. a
|
||||
# shell comment in a code sample) be misparsed as a real heading.
|
||||
code_block_pattern = r"^(?:```|~~~)"
|
||||
headers = []
|
||||
in_code_block = False
|
||||
|
||||
|
|
|
|||
|
|
@ -77,7 +77,13 @@ class PdfParser:
|
|||
|
||||
try:
|
||||
pix = pymupdf.Pixmap(image_bytes)
|
||||
if pix.n > 4:
|
||||
# n includes the alpha channel, so a plain RGBA pixmap also
|
||||
# has n==4 — subtract alpha before comparing. Without this,
|
||||
# a CMYK image with no alpha (n==4, same as RGBA) skips the
|
||||
# RGB conversion, and pix.save() as .png then raises
|
||||
# "unsupported colorspace for 'png'", silently dropping the
|
||||
# image via the bare except below.
|
||||
if pix.n - pix.alpha >= 4:
|
||||
pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
|
||||
filename = f"p{page_num}_img{img_idx}.png"
|
||||
save_path = images_path / filename
|
||||
|
|
|
|||
|
|
@ -1,27 +1,24 @@
|
|||
import json
|
||||
import PyPDF2
|
||||
|
||||
try:
|
||||
from .index.utils import get_number_of_pages, remove_fields, get_md_page_content
|
||||
from .index.utils import (
|
||||
get_number_of_pages, remove_fields, get_md_page_content,
|
||||
parse_pages, get_pdf_page_content,
|
||||
)
|
||||
except ImportError:
|
||||
from index.utils import get_number_of_pages, remove_fields, get_md_page_content
|
||||
from index.utils import (
|
||||
get_number_of_pages, remove_fields, get_md_page_content,
|
||||
parse_pages, get_pdf_page_content,
|
||||
)
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _parse_pages(pages: str) -> list[int]:
|
||||
"""Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints."""
|
||||
result = []
|
||||
for part in pages.split(','):
|
||||
part = part.strip()
|
||||
if '-' in part:
|
||||
start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip())
|
||||
if start > end:
|
||||
raise ValueError(f"Invalid range '{part}': start must be <= end")
|
||||
result.extend(range(start, end + 1))
|
||||
else:
|
||||
result.append(int(part))
|
||||
return sorted(set(result))
|
||||
"""Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.
|
||||
Delegates to the canonical implementation so the two never drift again —
|
||||
this one used to lack the p>=1 filter and the 1000-page DoS cap."""
|
||||
return parse_pages(pages)
|
||||
|
||||
|
||||
def _count_pages(doc_info: dict) -> int:
|
||||
|
|
@ -34,7 +31,8 @@ def _count_pages(doc_info: dict) -> int:
|
|||
|
||||
|
||||
def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
|
||||
"""Extract text for specific PDF pages (1-indexed). Prefer cached pages, fallback to PDF."""
|
||||
"""Extract text for specific PDF pages (1-indexed). Prefer cached pages,
|
||||
else delegate the file-read fallback to the canonical implementation."""
|
||||
cached_pages = doc_info.get('pages')
|
||||
if cached_pages:
|
||||
page_map = {p['page']: p['content'] for p in cached_pages}
|
||||
|
|
@ -42,15 +40,7 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
|
|||
{'page': p, 'content': page_map[p]}
|
||||
for p in page_nums if p in page_map
|
||||
]
|
||||
path = doc_info['path']
|
||||
with open(path, 'rb') as f:
|
||||
pdf_reader = PyPDF2.PdfReader(f)
|
||||
total = len(pdf_reader.pages)
|
||||
valid_pages = [p for p in page_nums if 1 <= p <= total]
|
||||
return [
|
||||
{'page': p, 'content': pdf_reader.pages[p - 1].extract_text() or ''}
|
||||
for p in valid_pages
|
||||
]
|
||||
return get_pdf_page_content(doc_info['path'], page_nums)
|
||||
|
||||
|
||||
def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,22 @@
|
|||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from ..errors import CollectionAlreadyExistsError, PageIndexError
|
||||
|
||||
# Mirrors LocalBackend's own collection-name rule. SQLiteStorage enforces this
|
||||
# itself (not just relying on LocalBackend's pre-check or the schema's CHECK
|
||||
# constraint below) because it's a public StorageEngine that can be used
|
||||
# directly, bypassing LocalBackend entirely.
|
||||
_COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
|
||||
|
||||
|
||||
def _validate_collection_name(name: str) -> None:
|
||||
if not _COLLECTION_NAME_RE.match(name):
|
||||
raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].")
|
||||
|
||||
|
||||
class SQLiteStorage:
|
||||
def __init__(self, db_path: str):
|
||||
|
|
@ -48,7 +62,17 @@ class SQLiteStorage:
|
|||
conn.execute("PRAGMA user_version = 1")
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS collections (
|
||||
name TEXT PRIMARY KEY CHECK(length(name) <= 128 AND name GLOB '[a-zA-Z0-9_-]*'),
|
||||
-- GLOB '*' is "any characters", not a regex quantifier over the
|
||||
-- preceding class — '[a-zA-Z0-9_-]*' alone only constrains the
|
||||
-- FIRST character. The second GLOB (NOT ... '*[^...]*') checks
|
||||
-- every remaining character too, so this is real defense-in-depth
|
||||
-- for direct SQLiteStorage use (bypassing _validate_collection_name
|
||||
-- above), not just a first-character gate.
|
||||
name TEXT PRIMARY KEY CHECK(
|
||||
length(name) BETWEEN 1 AND 128
|
||||
AND name GLOB '[a-zA-Z0-9_-]*'
|
||||
AND name NOT GLOB '*[^a-zA-Z0-9_-]*'
|
||||
),
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
|
|
@ -70,12 +94,17 @@ class SQLiteStorage:
|
|||
conn.commit()
|
||||
|
||||
def create_collection(self, name: str) -> None:
|
||||
_validate_collection_name(name)
|
||||
with self._write_lock:
|
||||
conn = self._get_conn()
|
||||
conn.execute("INSERT INTO collections (name) VALUES (?)", (name,))
|
||||
try:
|
||||
conn.execute("INSERT INTO collections (name) VALUES (?)", (name,))
|
||||
except sqlite3.IntegrityError as e:
|
||||
raise CollectionAlreadyExistsError(f"Collection '{name}' already exists") from e
|
||||
conn.commit()
|
||||
|
||||
def get_or_create_collection(self, name: str) -> None:
|
||||
_validate_collection_name(name)
|
||||
with self._write_lock:
|
||||
conn = self._get_conn()
|
||||
conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,))
|
||||
|
|
|
|||
|
|
@ -14,13 +14,23 @@ class DocumentInfo(TypedDict):
|
|||
doc_type: str
|
||||
|
||||
|
||||
class DocumentDetail(DocumentInfo, total=False):
|
||||
class _DocumentDetailRequired(DocumentInfo):
|
||||
"""``structure`` is always present — split into its own (default
|
||||
total=True) base so the total=False below only applies to the genuinely
|
||||
optional, backend-specific fields below. A single
|
||||
``class DocumentDetail(DocumentInfo, total=False): structure: ...`` would
|
||||
incorrectly mark ``structure`` optional too, since total=False applies to
|
||||
the whole class body, not just the fields declared after it.
|
||||
"""
|
||||
structure: list[dict[str, Any]]
|
||||
|
||||
|
||||
class DocumentDetail(_DocumentDetailRequired, total=False):
|
||||
"""A document with its tree, as returned by ``get_document()``.
|
||||
|
||||
``structure`` is always present; ``file_path`` is local-only and
|
||||
``status`` is cloud-only, hence total=False.
|
||||
``status`` is cloud-only, hence total=False for those two only.
|
||||
"""
|
||||
structure: list[dict[str, Any]]
|
||||
file_path: str # local backend only
|
||||
status: str # cloud backend only
|
||||
|
||||
|
|
|
|||
|
|
@ -37,8 +37,8 @@ if __name__ == "__main__":
|
|||
help='Add raw text to nodes (off by default). Bare flag or yes/no')
|
||||
|
||||
# Markdown specific arguments
|
||||
parser.add_argument('--if-thinning', type=str, default='no',
|
||||
help='Whether to apply tree thinning for markdown (markdown only)')
|
||||
parser.add_argument('--if-thinning', nargs='?', const=True, type=_cli_bool, default=None,
|
||||
help='Apply tree thinning (off by default, markdown only). Bare flag or yes/no')
|
||||
parser.add_argument('--thinning-threshold', type=int, default=5000,
|
||||
help='Minimum token threshold for thinning (markdown only)')
|
||||
parser.add_argument('--summary-token-threshold', type=int, default=200,
|
||||
|
|
@ -101,7 +101,7 @@ if __name__ == "__main__":
|
|||
|
||||
toc_with_page_number = asyncio.run(md_to_tree(
|
||||
md_path=args.md_path,
|
||||
if_thinning=args.if_thinning.lower() == 'yes',
|
||||
if_thinning=bool(args.if_thinning),
|
||||
min_token_threshold=args.thinning_threshold,
|
||||
if_add_node_summary=opt.if_add_node_summary,
|
||||
summary_token_threshold=args.summary_token_threshold,
|
||||
|
|
|
|||
|
|
@ -3,10 +3,15 @@ now deprecation shims over the canonical pageindex.index.* modules. These
|
|||
tests pin the compatibility contract."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import subprocess
|
||||
import sys
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def test_plain_import_pageindex_does_not_warn():
|
||||
# `import pageindex` must not route through the deprecation shims.
|
||||
|
|
@ -116,3 +121,27 @@ def test_md_to_tree_coerces_legacy_yes_no_strings(tmp_path):
|
|||
|
||||
assert not _has_summary(result["structure"])
|
||||
assert all("node_id" in n for n in result["structure"])
|
||||
|
||||
|
||||
def test_page_index_stays_callable_after_the_submodule_is_imported():
|
||||
"""pageindex/__init__.py binds the FUNCTION `page_index` as the package
|
||||
attribute, but pageindex/page_index.py is ALSO a real submodule of the
|
||||
same name — importing that submodule anywhere clobbers the package
|
||||
attribute with the module object (Python's import machinery does this
|
||||
unconditionally). Must run in a fresh subprocess: the effect depends on
|
||||
import order, so it can't be reliably observed against an
|
||||
already-imported pageindex in this test process."""
|
||||
script = (
|
||||
"import warnings; warnings.simplefilter('ignore')\n"
|
||||
"import pageindex.page_index\n" # the clobbering import
|
||||
"from pageindex import page_index\n"
|
||||
"assert callable(page_index), f'page_index is not callable: {type(page_index)}'\n"
|
||||
"from pageindex.page_index import page_index_main\n" # old multi-symbol import still works
|
||||
"assert callable(page_index_main)\n"
|
||||
"print('OK')\n"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script], capture_output=True, text=True, cwd=str(_REPO_ROOT),
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert "OK" in result.stdout
|
||||
|
|
|
|||
|
|
@ -35,6 +35,27 @@ def test_unsupported_file_type_raises(backend, tmp_path):
|
|||
backend.add_document("papers", str(bad_file))
|
||||
|
||||
|
||||
def test_add_document_on_empty_markdown_file_does_not_crash(tmp_path):
|
||||
"""An empty/whitespace-only .md file used to route into the PDF-oriented
|
||||
TOC-detection pipeline (no node ever has 'level' set), wasting an LLM call
|
||||
and then raising IndexingError. Must complete instantly with zero LLM
|
||||
calls when summary/description are off."""
|
||||
from pageindex.config import IndexConfig
|
||||
|
||||
storage = SQLiteStorage(str(tmp_path / "test.db"))
|
||||
backend = LocalBackend(
|
||||
storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o",
|
||||
index_config=IndexConfig(if_add_node_summary=False, if_add_doc_description=False),
|
||||
)
|
||||
backend.get_or_create_collection("papers")
|
||||
empty_md = tmp_path / "empty.md"
|
||||
empty_md.write_text(" \n\n \n")
|
||||
|
||||
doc_id = backend.add_document("papers", str(empty_md)) # must not raise
|
||||
|
||||
assert backend.get_document_structure("papers", doc_id) == []
|
||||
|
||||
|
||||
def test_register_custom_parser(backend):
|
||||
from pageindex.parser.protocol import ParsedDocument, ContentNode
|
||||
|
||||
|
|
|
|||
|
|
@ -71,3 +71,32 @@ def test_headerless_file_yields_single_node(tmp_path):
|
|||
assert len(result.nodes) == 1
|
||||
assert result.nodes[0].title == "plain"
|
||||
assert "No headings at all" in result.nodes[0].content
|
||||
|
||||
|
||||
def test_utf8_bom_does_not_break_the_first_header(tmp_path):
|
||||
"""A leading BOM (common from Windows editors/exporters) isn't
|
||||
whitespace, so .strip() doesn't remove it — without utf-8-sig decoding,
|
||||
the header regex fails to match the BOM-prefixed first line, and it gets
|
||||
misclassified as unrecognized preamble text instead of a real heading."""
|
||||
md = tmp_path / "bom.md"
|
||||
md.write_bytes(b"\xef\xbb\xbf# First Header\nbody text\n")
|
||||
result = MarkdownParser().parse(str(md))
|
||||
assert len(result.nodes) == 1
|
||||
assert result.nodes[0].title == "First Header"
|
||||
assert result.nodes[0].level == 1
|
||||
|
||||
|
||||
def test_tilde_fenced_code_blocks_are_recognized(tmp_path):
|
||||
"""CommonMark allows both backtick and tilde code fences. Only
|
||||
recognizing backticks let a '#'-prefixed line inside a ~~~-fenced block
|
||||
(e.g. a shell comment in a sample) be misparsed as a real heading."""
|
||||
md = tmp_path / "tilde.md"
|
||||
md.write_text(
|
||||
"# Real Header\nintro\n"
|
||||
"~~~\n# not a real header, just a comment\n~~~\n"
|
||||
"## Real Sub\nmore\n"
|
||||
)
|
||||
result = MarkdownParser().parse(str(md))
|
||||
titles = [n.title for n in result.nodes]
|
||||
assert titles == ["Real Header", "Real Sub"]
|
||||
assert "not a real header" not in " ".join(n.title for n in result.nodes)
|
||||
|
|
|
|||
|
|
@ -34,3 +34,43 @@ def test_retrieve_md_page_content_returns_only_requested_lines():
|
|||
|
||||
out = _get_md_page_content({"structure": _md_structure()}, [5, 100])
|
||||
assert [r["page"] for r in out] == [5, 100]
|
||||
|
||||
|
||||
def test_retrieve_parse_pages_delegates_to_canonical_and_enforces_dos_cap():
|
||||
"""retrieve._parse_pages used to be an independent copy that lacked the
|
||||
canonical parse_pages' p>=1 filter and 1000-page cap — a caller of the
|
||||
legacy pageindex.get_page_content could bypass the DoS guard the SDK path
|
||||
enforces. Now it's a one-line delegate, so they can't drift again."""
|
||||
from pageindex.retrieve import _parse_pages
|
||||
from pageindex.index.utils import parse_pages
|
||||
import pytest
|
||||
|
||||
assert _parse_pages("5-7") == parse_pages("5-7") == [5, 6, 7]
|
||||
with pytest.raises(ValueError, match="too large"):
|
||||
_parse_pages("1-99999999")
|
||||
|
||||
|
||||
def test_retrieve_get_pdf_page_content_falls_back_to_canonical(tmp_path, monkeypatch):
|
||||
"""When no cached 'pages' are present, the file-read fallback must
|
||||
delegate to the canonical get_pdf_page_content instead of re-implementing
|
||||
PDF text extraction inline (a second, independently-maintained copy)."""
|
||||
from pageindex.retrieve import _get_pdf_page_content
|
||||
import pageindex.retrieve as retrieve_mod
|
||||
|
||||
calls = []
|
||||
monkeypatch.setattr(
|
||||
retrieve_mod, "get_pdf_page_content",
|
||||
lambda path, page_nums: calls.append((path, page_nums)) or [{"page": 1, "content": "x"}],
|
||||
)
|
||||
result = _get_pdf_page_content({"path": "/fake/doc.pdf"}, [1])
|
||||
assert calls == [("/fake/doc.pdf", [1])]
|
||||
assert result == [{"page": 1, "content": "x"}]
|
||||
|
||||
|
||||
def test_retrieve_get_pdf_page_content_prefers_cache_over_file():
|
||||
from pageindex.retrieve import _get_pdf_page_content
|
||||
|
||||
doc_info = {"path": "/should/not/be/opened.pdf",
|
||||
"pages": [{"page": 1, "content": "cached one"}, {"page": 2, "content": "cached two"}]}
|
||||
result = _get_pdf_page_content(doc_info, [2])
|
||||
assert result == [{"page": 2, "content": "cached two"}]
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import pymupdf
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from pageindex.parser.pdf import PdfParser
|
||||
|
|
@ -29,6 +30,24 @@ def test_parse_nodes_are_flat_without_level():
|
|||
assert node.level is None
|
||||
|
||||
|
||||
def test_cmyk_pixmap_without_alpha_is_saveable_as_png(tmp_path):
|
||||
"""A CMYK image with no alpha has n==4 -- same as RGBA -- so `pix.n > 4`
|
||||
wrongly skips the RGB conversion PNG needs, and pix.save() raises
|
||||
'unsupported colorspace for png', silently dropping the image via the
|
||||
extractor's bare except. The fix (`pix.n - pix.alpha >= 4`) must convert
|
||||
CMYK (4-0=4) while leaving RGBA (4-1=3) untouched."""
|
||||
cmyk = pymupdf.Pixmap(pymupdf.csCMYK, pymupdf.Rect(0, 0, 10, 10))
|
||||
assert cmyk.n == 4 and cmyk.alpha == 0
|
||||
assert cmyk.n - cmyk.alpha >= 4 # the fixed condition: must convert
|
||||
converted = pymupdf.Pixmap(pymupdf.csRGB, cmyk)
|
||||
converted.save(str(tmp_path / "cmyk.png")) # must not raise
|
||||
|
||||
rgba = pymupdf.Pixmap(pymupdf.Pixmap(pymupdf.csRGB, pymupdf.Rect(0, 0, 10, 10)), 1)
|
||||
assert rgba.n == 4 and rgba.alpha == 1
|
||||
assert not (rgba.n - rgba.alpha >= 4) # unchanged: RGBA needs no conversion
|
||||
rgba.save(str(tmp_path / "rgba.png")) # already saveable as-is
|
||||
|
||||
|
||||
def test_image_paths_are_absolute(tmp_path):
|
||||
"""Image references must be absolute so they resolve regardless of cwd
|
||||
(cwd-relative paths broke after the query ran from another directory)."""
|
||||
|
|
|
|||
|
|
@ -25,6 +25,22 @@ def test_detect_strategy_without_level():
|
|||
assert detect_strategy(nodes) == "content_based"
|
||||
|
||||
|
||||
def test_detect_strategy_empty_nodes_is_level_based():
|
||||
"""An empty node list (e.g. an empty/whitespace-only source file) must
|
||||
route to level_based, whose build_tree_from_levels([]) returns an empty
|
||||
structure with zero LLM calls — not content_based, whose TOC-detection
|
||||
pipeline needs real page content and wastes an LLM call before failing."""
|
||||
assert detect_strategy([]) == "level_based"
|
||||
|
||||
|
||||
def test_build_index_on_empty_document_makes_no_llm_calls():
|
||||
from pageindex.config import IndexConfig
|
||||
parsed = ParsedDocument(doc_name="empty", nodes=[])
|
||||
opt = IndexConfig(if_add_node_summary=False, if_add_doc_description=False)
|
||||
result = build_index(parsed, opt=opt)
|
||||
assert result == {"doc_name": "empty", "structure": []}
|
||||
|
||||
|
||||
def test_build_tree_from_levels():
|
||||
nodes = [
|
||||
ContentNode(content="ch1 text", tokens=10, title="Chapter 1", index=1, level=1),
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
2d46d68..8f536cb): the Markdown text-stripping fix's fallout, plus the other
|
||||
directly-fixable findings from that pass."""
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -122,6 +123,20 @@ def test_cloud_delete_collection_still_clears_real_folder_id(monkeypatch):
|
|||
|
||||
|
||||
# ── #7: remove_structure_text is skipped when text was never added ───────────
|
||||
def _mock_content_based_pipeline(monkeypatch, structure):
|
||||
"""content_based's real path (_content_based_pipeline) drives real LLM
|
||||
calls (TOC detection etc.) regardless of if_add_node_summary — a prior
|
||||
version of these two tests didn't mock this out, fell through to it, and
|
||||
made real network calls (with a dummy key: 10 retries before failing;
|
||||
with a real key: real billable requests) on every run."""
|
||||
from pageindex.index import pipeline
|
||||
|
||||
async def fake(page_list, opt):
|
||||
return structure
|
||||
|
||||
monkeypatch.setattr(pipeline, "_content_based_pipeline", fake)
|
||||
|
||||
|
||||
def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch):
|
||||
from pageindex.index import pipeline
|
||||
from pageindex.parser.protocol import ContentNode, ParsedDocument
|
||||
|
|
@ -131,6 +146,7 @@ def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch):
|
|||
# ...` inside the function body), so patch it on the utils module itself.
|
||||
import pageindex.index.utils as utils_mod
|
||||
monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
|
||||
_mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}])
|
||||
|
||||
nodes = [ContentNode(content="page one text", tokens=5, index=1)]
|
||||
parsed = ParsedDocument(doc_name="d", nodes=nodes)
|
||||
|
|
@ -147,6 +163,14 @@ def test_build_index_still_strips_text_when_summary_added_it(monkeypatch):
|
|||
calls = []
|
||||
import pageindex.index.utils as utils_mod
|
||||
monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s)
|
||||
_mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}])
|
||||
# Summary generation itself would otherwise make a real LLM call.
|
||||
monkeypatch.setattr(
|
||||
utils_mod, "generate_summaries_for_structure",
|
||||
AsyncMock(side_effect=lambda structure, model=None: [
|
||||
n.__setitem__("summary", "fake") for n in structure
|
||||
]),
|
||||
)
|
||||
|
||||
nodes = [ContentNode(content="page one text", tokens=5, index=1)]
|
||||
parsed = ParsedDocument(doc_name="d", nodes=nodes)
|
||||
|
|
|
|||
|
|
@ -19,6 +19,60 @@ def test_delete_collection(storage):
|
|||
storage.delete_collection("papers")
|
||||
assert "papers" not in storage.list_collections()
|
||||
|
||||
|
||||
def test_create_duplicate_collection_raises_pageindex_error(storage):
|
||||
"""A raw sqlite3.IntegrityError leaking out breaks `except PageIndexError`
|
||||
catch-alls; must be translated to a proper SDK exception."""
|
||||
from pageindex.errors import CollectionAlreadyExistsError, PageIndexError
|
||||
storage.create_collection("papers")
|
||||
with pytest.raises(CollectionAlreadyExistsError):
|
||||
storage.create_collection("papers")
|
||||
# also catchable via the SDK's generic base class
|
||||
storage.create_collection("other")
|
||||
with pytest.raises(PageIndexError):
|
||||
storage.create_collection("other")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_name", [
|
||||
"a/../../etc/passwd", "/etc/passwd", "a$(whoami)", ".hidden",
|
||||
"a b", "válid", "", "a" * 129,
|
||||
])
|
||||
def test_create_collection_rejects_invalid_names_at_the_python_layer(storage, bad_name):
|
||||
"""SQLiteStorage must validate collection names itself — it's a public
|
||||
StorageEngine that can be used directly, bypassing LocalBackend's own
|
||||
regex check entirely."""
|
||||
from pageindex.errors import PageIndexError
|
||||
with pytest.raises(PageIndexError):
|
||||
storage.create_collection(bad_name)
|
||||
|
||||
|
||||
def test_sql_check_constraint_also_rejects_invalid_names_directly(storage):
|
||||
"""Defense-in-depth: even bypassing SQLiteStorage's own Python validation
|
||||
and inserting via raw SQL, the schema's CHECK constraint must reject a
|
||||
name that isn't ENTIRELY [a-zA-Z0-9_-] — not just its first character
|
||||
(GLOB '*' is a wildcard, not a regex quantifier over the preceding class,
|
||||
so 'name GLOB [a-zA-Z0-9_-]*' alone only constrains the first character)."""
|
||||
import sqlite3
|
||||
conn = storage._get_conn()
|
||||
with pytest.raises(sqlite3.IntegrityError):
|
||||
conn.execute("INSERT INTO collections (name) VALUES (?)", ("a/../../etc/passwd",))
|
||||
|
||||
|
||||
def test_malicious_collection_name_rejected_through_local_backend_too(tmp_path):
|
||||
"""End-to-end via the normal LocalBackend entry point: a path-traversal-
|
||||
shaped collection name must never reach add_document's
|
||||
files_dir / collection path construction. Three independent layers now
|
||||
reject it (LocalBackend's own regex, SQLiteStorage's regex, and the SQL
|
||||
CHECK constraint) — this pins the outermost one."""
|
||||
from pageindex.backend.local import LocalBackend
|
||||
from pageindex.errors import PageIndexError
|
||||
|
||||
storage = SQLiteStorage(str(tmp_path / "t.db"))
|
||||
backend = LocalBackend(storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o")
|
||||
with pytest.raises(PageIndexError):
|
||||
backend.create_collection("a/../../escape_me")
|
||||
assert not (tmp_path / "escape_me").exists()
|
||||
|
||||
def test_save_and_get_document(storage):
|
||||
storage.create_collection("papers")
|
||||
doc = {
|
||||
|
|
|
|||
20
tests/test_types.py
Normal file
20
tests/test_types.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from pageindex.types import DocumentDetail, DocumentInfo, PageContent
|
||||
|
||||
|
||||
def test_document_detail_structure_field_is_required():
|
||||
"""structure is always populated by both LocalBackend.get_document and
|
||||
CloudBackend.get_document — must be a required key, not optional, or
|
||||
type checkers/tooling built on this TypedDict wrongly treat a
|
||||
DocumentDetail missing 'structure' as valid."""
|
||||
assert "structure" in DocumentDetail.__required_keys__
|
||||
assert "structure" not in DocumentDetail.__optional_keys__
|
||||
|
||||
|
||||
def test_document_detail_backend_specific_fields_stay_optional():
|
||||
assert "file_path" in DocumentDetail.__optional_keys__
|
||||
assert "status" in DocumentDetail.__optional_keys__
|
||||
|
||||
|
||||
def test_document_detail_inherits_document_info_as_required():
|
||||
for key in ("doc_id", "doc_name", "doc_description", "doc_type"):
|
||||
assert key in DocumentDetail.__required_keys__
|
||||
Loading…
Add table
Add a link
Reference in a new issue