Merge pull request #775 from cbcoutinho/feat/chunk-bbox-payload

feat(vector): replace inline PDF page images with chunk_bbox (Deck #76)
This commit is contained in:
Chris Coutinho
2026-05-08 23:05:53 +02:00
committed by GitHub
9 changed files with 544 additions and 81 deletions
+1
View File
@@ -7,6 +7,7 @@ description: |
in this repo's automated PR reviews. Use when the user is about to push, says "ready in this repo's automated PR reviews. Use when the user is about to push, says "ready
to push", "review my work", "check before PR", or invokes /pre-push-review. to push", "review my work", "check before PR", or invokes /pre-push-review.
Report-only — does not modify code. Report-only — does not modify code.
model: sonnet
allowed-tools: allowed-tools:
- Bash - Bash
- Read - Read
+9 -9
View File
@@ -553,9 +553,10 @@ async def get_chunk_context(request: Request) -> JSONResponse:
status_code=404, status_code=404,
) )
# For PDF files, also fetch the highlighted page image from Qdrant if available # For PDF files, also fetch the chunk's bounding box from Qdrant if
# This is useful for clients that want to show a pre-rendered image # available so the client can overlay a highlight on top of a
highlighted_page_image = None # render-on-demand page image (Deck #76).
chunk_bbox = None
page_number = chunk_context.page_number page_number = chunk_context.page_number
if doc_type == "file": if doc_type == "file":
@@ -563,7 +564,6 @@ async def get_chunk_context(request: Request) -> JSONResponse:
settings = get_settings() settings = get_settings()
qdrant_client = await get_qdrant_client() qdrant_client = await get_qdrant_client()
# Query for this specific chunk's highlighted image
points_response = await qdrant_client.scroll( points_response = await qdrant_client.scroll(
collection_name=settings.get_collection_name(), collection_name=settings.get_collection_name(),
scroll_filter=Filter( scroll_filter=Filter(
@@ -585,19 +585,19 @@ async def get_chunk_context(request: Request) -> JSONResponse:
), ),
limit=1, limit=1,
with_vectors=False, with_vectors=False,
with_payload=["highlighted_page_image", "page_number"], with_payload=["chunk_bbox", "page_number"],
) )
if points_response[0]: if points_response[0]:
payload = points_response[0][0].payload payload = points_response[0][0].payload
if payload: if payload:
highlighted_page_image = payload.get("highlighted_page_image") chunk_bbox = payload.get("chunk_bbox")
# Trust Qdrant page number if available (might be more accurate than context expansion logic) # Trust Qdrant page number if available (might be more accurate than context expansion logic)
if payload.get("page_number") is not None: if payload.get("page_number") is not None:
page_number = payload.get("page_number") page_number = payload.get("page_number")
except Exception as e: except Exception as e:
logger.warning(f"Failed to fetch highlighted image: {e}") logger.warning(f"Failed to fetch chunk bbox: {e}")
# Build response # Build response
response_data = { response_data = {
@@ -612,8 +612,8 @@ async def get_chunk_context(request: Request) -> JSONResponse:
"total_chunks": chunk_context.total_chunks, "total_chunks": chunk_context.total_chunks,
} }
if highlighted_page_image: if chunk_bbox:
response_data["highlighted_page_image"] = highlighted_page_image response_data["chunk_bbox"] = chunk_bbox
return JSONResponse(response_data) return JSONResponse(response_data)
+2 -1
View File
@@ -125,12 +125,13 @@ from nextcloud_mcp_server.server import (
) )
from nextcloud_mcp_server.server.auth_tools import register_auth_tools from nextcloud_mcp_server.server.auth_tools import register_auth_tools
from nextcloud_mcp_server.server.oauth_tools import register_oauth_tools from nextcloud_mcp_server.server.oauth_tools import register_oauth_tools
from nextcloud_mcp_server.vector import processor_task, scanner_task
from nextcloud_mcp_server.vector.oauth_sync import ( from nextcloud_mcp_server.vector.oauth_sync import (
oauth_processor_task, oauth_processor_task,
user_manager_task, user_manager_task,
) )
from nextcloud_mcp_server.vector.processor import processor_task
from nextcloud_mcp_server.vector.qdrant_client import get_qdrant_client from nextcloud_mcp_server.vector.qdrant_client import get_qdrant_client
from nextcloud_mcp_server.vector.scanner import scanner_task
from nextcloud_mcp_server.vector.webhook_receiver import handle_nextcloud_webhook from nextcloud_mcp_server.vector.webhook_receiver import handle_nextcloud_webhook
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
+12 -14
View File
@@ -604,8 +604,10 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse:
f"after_len={len(chunk_context.after_context)}" f"after_len={len(chunk_context.after_context)}"
) )
# For PDF files, also fetch the highlighted page image from Qdrant # For PDF files, also fetch the chunk bbox from Qdrant so the client
highlighted_page_image = None # can overlay a highlight on top of a render-on-demand page image
# (Deck #76).
chunk_bbox = None
page_number = None page_number = None
if doc_type == "file": if doc_type == "file":
try: try:
@@ -613,7 +615,6 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse:
qdrant_client = await get_qdrant_client() qdrant_client = await get_qdrant_client()
username = request.user.display_name username = request.user.display_name
# Query for this specific chunk's highlighted image
points_response = await qdrant_client.scroll( points_response = await qdrant_client.scroll(
collection_name=settings.get_collection_name(), collection_name=settings.get_collection_name(),
scroll_filter=Filter( scroll_filter=Filter(
@@ -635,22 +636,20 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse:
), ),
limit=1, limit=1,
with_vectors=False, with_vectors=False,
with_payload=["highlighted_page_image", "page_number"], with_payload=["chunk_bbox", "page_number"],
) )
points = points_response[0] points = points_response[0]
if points and points[0].payload: if points and points[0].payload:
highlighted_page_image = points[0].payload.get( chunk_bbox = points[0].payload.get("chunk_bbox")
"highlighted_page_image"
)
page_number = points[0].payload.get("page_number") page_number = points[0].payload.get("page_number")
if highlighted_page_image: if chunk_bbox:
logger.info( logger.info(
f"Found highlighted image for chunk: " f"Found chunk bbox: page={page_number}, "
f"page={page_number}, image_size={len(highlighted_page_image)}" f"rects={len(chunk_bbox)}"
) )
except Exception as e: except Exception as e:
logger.warning(f"Failed to fetch highlighted image: {e}") logger.warning(f"Failed to fetch chunk bbox: {e}")
# Return response compatible with frontend expectations # Return response compatible with frontend expectations
response_data: dict = { response_data: dict = {
@@ -662,9 +661,8 @@ async def chunk_context_endpoint(request: Request) -> JSONResponse:
"has_more_after": chunk_context.has_after_truncation, "has_more_after": chunk_context.has_after_truncation,
} }
# Add image data if available if chunk_bbox:
if highlighted_page_image: response_data["chunk_bbox"] = chunk_bbox
response_data["highlighted_page_image"] = highlighted_page_image
response_data["page_number"] = page_number response_data["page_number"] = page_number
return JSONResponse(response_data) return JSONResponse(response_data)
@@ -691,6 +691,122 @@ class PDFHighlighter:
f"Failed to delete temp directory {temp_pdf_path.parent}: {e}" f"Failed to delete temp directory {temp_pdf_path.parent}: {e}"
) )
@staticmethod
def compute_chunk_bboxes_batch(
pdf_bytes: bytes,
chunks: list[tuple[int, int, int, int | None, str]],
page_boundaries: list[dict],
full_text: str,
) -> dict[int, tuple[list[tuple[float, float, float, float]], int]]:
"""Compute normalized bounding boxes for chunks without rendering.
Lightweight alternative to highlight_chunks_batch — opens the PDF,
locates each chunk on its assigned page using the same text-search
path as the highlighter (`_find_chunk_bbox`), and returns
page-normalized rectangles. Skips the get_pixmap + PIL pipeline
entirely, so no PNG bytes are produced.
Args:
pdf_bytes: PDF file bytes.
chunks: List of (chunk_index, start_offset, end_offset,
stored_page_number, chunk_text). chunk_index is the dict key.
page_boundaries: Pre-computed page boundaries from the document
processor; each entry is {"page", "start_offset", "end_offset"}.
full_text: Full document text (for cross-page chunk handling).
Returns:
dict mapping chunk_index to (normalized_bboxes, page_number).
Each bbox is (x0, y0, x1, y1) in [0, 1] relative to page width
and height, top-left origin. Chunks whose bbox cannot be located
are omitted from the result.
"""
results: dict[int, tuple[list[tuple[float, float, float, float]], int]] = {}
if not chunks:
return results
temp_pdf_path = None
doc = None
try:
temp_dir = Path(tempfile.mkdtemp(prefix="pdf_bbox_batch_"))
temp_pdf_path = temp_dir / "pdf.pdf"
temp_pdf_path.write_bytes(pdf_bytes)
doc = pymupdf.open(temp_pdf_path)
for (
chunk_index,
start_offset,
end_offset,
_,
_,
) in chunks:
chunk_page_info = PDFHighlighter.find_chunk_page(
start_offset, end_offset, page_boundaries
)
if not chunk_page_info:
logger.debug("Chunk %s: not found on any page", chunk_index)
continue
page_num = chunk_page_info["page_num"]
page_boundary = next(
(b for b in page_boundaries if b["page"] == page_num), None
)
if page_boundary is None:
logger.debug(
"Chunk %s: page %s not found in boundaries",
chunk_index,
page_num,
)
continue
page_text_length = (
page_boundary["end_offset"] - page_boundary["start_offset"]
)
# Page-relative slice (handles chunks that span page boundaries)
chunk_start_on_page = max(start_offset, page_boundary["start_offset"])
chunk_end_on_page = min(end_offset, page_boundary["end_offset"])
page_relative_text = full_text[chunk_start_on_page:chunk_end_on_page]
page = doc[page_num - 1]
bbox = PDFHighlighter._find_chunk_bbox(
page,
page_relative_text,
chunk_page_info["page_relative_start"],
chunk_page_info["page_relative_end"],
page_text_length,
)
if bbox is None:
continue
page_rect = page.rect
w = page_rect.width or 1.0
h = page_rect.height or 1.0
normalized = (
bbox[0] / w,
bbox[1] / h,
bbox[2] / w,
bbox[3] / h,
)
results[chunk_index] = ([normalized], page_num)
logger.info(f"Computed bboxes for {len(results)}/{len(chunks)} chunks")
return results
except Exception as e:
logger.error(f"Error computing chunk bboxes: {e}", exc_info=True)
return results
finally:
if doc is not None:
doc.close()
if temp_pdf_path and temp_pdf_path.parent.exists():
try:
shutil.rmtree(temp_pdf_path.parent)
except Exception as e:
logger.warning(f"Failed to clean up temp dir: {e}")
@staticmethod @staticmethod
def highlight_chunks_batch( def highlight_chunks_batch(
pdf_bytes: bytes, pdf_bytes: bytes,
+9 -8
View File
@@ -1,16 +1,17 @@
"""Vector database and background sync package.""" """Vector database and background sync package.
`processor` and `scanner` are intentionally NOT re-exported from this
package init: they transitively import `server.semantic` ->
`search.bm25_hybrid`, which forms an import cycle with
`search.algorithms` -> `vector.placeholder` -> `vector/__init__`.
Consumers that need those symbols import them from their submodules
directly (e.g. `from nextcloud_mcp_server.vector.processor import ...`).
"""
from .document_chunker import DocumentChunker from .document_chunker import DocumentChunker
from .processor import process_document, processor_task
from .qdrant_client import get_qdrant_client from .qdrant_client import get_qdrant_client
from .scanner import DocumentTask, scan_user_documents, scanner_task
__all__ = [ __all__ = [
"get_qdrant_client", "get_qdrant_client",
"DocumentChunker", "DocumentChunker",
"scanner_task",
"scan_user_documents",
"DocumentTask",
"processor_task",
"process_document",
] ]
+26 -49
View File
@@ -3,7 +3,6 @@
Processes documents from stream: fetches content, generates embeddings, stores in Qdrant. Processes documents from stream: fetches content, generates embeddings, stores in Qdrant.
""" """
import base64
import logging import logging
import time import time
import uuid import uuid
@@ -538,7 +537,11 @@ async def _index_document(
# Initialize results containers # Initialize results containers
dense_embeddings: list = [] dense_embeddings: list = []
sparse_embeddings: list = [] sparse_embeddings: list = []
chunk_images: dict[int, dict] = {} # chunk_index -> list[(x0, y0, x1, y1)] of normalized rectangles
# in [0, 1] relative to page width/height. The page is taken from
# `chunk.page_number` (offset-based) and stored as `page_number`
# in the Qdrant payload, so we don't carry an `actual_page_num` here.
chunk_bboxes: dict[int, list[tuple[float, float, float, float]]] = {}
# Determine if we need PDF highlighting # Determine if we need PDF highlighting
is_pdf = doc_task.doc_type == "file" and content_type == "application/pdf" is_pdf = doc_task.doc_type == "file" and content_type == "application/pdf"
@@ -570,8 +573,8 @@ async def _index_document(
sparse_embeddings = await bm25_service.encode_batch(chunk_texts) sparse_embeddings = await bm25_service.encode_batch(chunk_texts)
async def generate_highlights(): async def generate_highlights():
"""Generate highlighted page images for PDF chunks (CPU-bound).""" """Compute chunk bounding boxes for PDF chunks (CPU-bound, no rendering)."""
nonlocal chunk_images nonlocal chunk_bboxes
if not is_pdf: if not is_pdf:
return return
@@ -579,64 +582,42 @@ async def _index_document(
assert content_bytes is not None assert content_bytes is not None
with trace_operation( with trace_operation(
"vector_sync.generate_highlights", "vector_sync.compute_chunk_bboxes",
attributes={ attributes={
"vector_sync.chunk_count": len(chunks), "vector_sync.chunk_count": len(chunks),
"vector_sync.pdf_size": len(content_bytes), "vector_sync.pdf_size": len(content_bytes),
}, },
): ):
# Build chunk data for batch processing
# Format: (chunk_index, start_offset, end_offset, page_number, chunk_text)
chunk_data: list[tuple[int, int, int, int | None, str]] = [ chunk_data: list[tuple[int, int, int, int | None, str]] = [
(i, chunk.start_offset, chunk.end_offset, chunk.page_number, chunk.text) (i, chunk.start_offset, chunk.end_offset, chunk.page_number, chunk.text)
for i, chunk in enumerate(chunks) for i, chunk in enumerate(chunks)
if chunk.page_number is not None if chunk.page_number is not None
] ]
# Get pre-computed page boundaries from document processor
page_boundaries = file_metadata.get("page_boundaries") page_boundaries = file_metadata.get("page_boundaries")
if not page_boundaries: if not page_boundaries:
logger.warning("No page boundaries available, skipping highlighting") logger.warning(
"No page boundaries available, skipping bbox computation"
)
return return
# Type narrowing: page_boundaries is guaranteed to be list[dict] here
page_boundaries_list = cast(list[dict[str, Any]], page_boundaries) page_boundaries_list = cast(list[dict[str, Any]], page_boundaries)
logger.info( logger.info(f"Computing chunk bboxes for {len(chunk_data)} PDF chunks")
f"Batch generating highlighted page images for {len(chunk_data)} PDF chunks"
)
# Run CPU-bound highlighting in thread pool
# Pass pre-computed page boundaries and full text to avoid re-processing the PDF
batch_results = await anyio.to_thread.run_sync( # type: ignore[attr-defined] batch_results = await anyio.to_thread.run_sync( # type: ignore[attr-defined]
lambda: PDFHighlighter.highlight_chunks_batch( lambda: PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=content_bytes, pdf_bytes=content_bytes,
chunks=chunk_data, chunks=chunk_data,
page_boundaries=page_boundaries_list, page_boundaries=page_boundaries_list,
full_text=content, full_text=content,
color="yellow",
zoom=2.0,
) )
) )
# Convert results to storage format for chunk_index, (bboxes, _) in batch_results.items():
for chunk_index, ( chunk_bboxes[chunk_index] = bboxes
png_bytes,
actual_page_num,
highlight_count,
) in batch_results.items():
image_base64 = base64.b64encode(png_bytes).decode("utf-8")
chunk_images[chunk_index] = {
"image": image_base64,
"page": actual_page_num,
"highlights": highlight_count,
"size": len(png_bytes),
}
logger.info( logger.info(f"Computed bboxes for {len(chunk_bboxes)}/{len(chunks)} chunks")
f"Generated {len(chunk_images)}/{len(chunks)} highlighted page images "
f"(avg {sum(img['size'] for img in chunk_images.values()) // max(len(chunk_images), 1):,} bytes)"
)
# Run all embedding/highlighting operations in parallel # Run all embedding/highlighting operations in parallel
# - Dense embeddings: I/O bound (API call) # - Dense embeddings: I/O bound (API call)
@@ -752,16 +733,11 @@ async def _index_document(
if doc_task.doc_type == "deck_card" if doc_task.doc_type == "deck_card"
else {} else {}
), ),
# Highlighted page image (PDF only) # Chunk bbox (PDF only) — normalized rectangles in [0,1]
**( # relative to page width/height. Replaces the legacy
{ # `highlighted_page_image` (Deck #76). The page number
"highlighted_page_image": chunk_images[i]["image"], # comes from `page_number` (set above for PDF chunks).
"highlighted_page_number": chunk_images[i]["page"], **({"chunk_bbox": chunk_bboxes[i]} if i in chunk_bboxes else {}),
"highlight_count": chunk_images[i]["highlights"],
}
if i in chunk_images
else {}
),
}, },
) )
) )
@@ -780,15 +756,16 @@ async def _index_document(
f"Failed to delete placeholder for {doc_task.doc_type}_{doc_task.doc_id}: {e}" f"Failed to delete placeholder for {doc_task.doc_type}_{doc_task.doc_id}: {e}"
) )
# Upsert to Qdrant in batches to avoid timeout with large payloads # Upsert to Qdrant in batches. Now that we no longer embed PNG payloads,
# Each batch is limited to avoid WriteTimeout when sending large image payloads # per-point payloads are small (chunk text + small metadata), so we can
BATCH_SIZE = 10 # ~2MB per batch with images # safely use a larger batch size.
BATCH_SIZE = 100
with trace_operation( with trace_operation(
"vector_sync.qdrant_upsert", "vector_sync.qdrant_upsert",
attributes={ attributes={
"vector_sync.point_count": len(points), "vector_sync.point_count": len(points),
"vector_sync.collection": settings.get_collection_name(), "vector_sync.collection": settings.get_collection_name(),
"vector_sync.images_count": len(chunk_images), "vector_sync.bboxes_count": len(chunk_bboxes),
"vector_sync.batch_size": BATCH_SIZE, "vector_sync.batch_size": BATCH_SIZE,
}, },
): ):
+137
View File
@@ -0,0 +1,137 @@
#!/usr/bin/env python3
"""Purge legacy `highlighted_page_image` payloads from Qdrant (Deck #76).
Iterates all points in the configured Qdrant collection and deletes the
legacy payload keys `highlighted_page_image`, `highlighted_page_number`,
and `highlight_count`. This relieves disk pressure caused by inline
base64 PNGs that the new code path no longer writes.
Idempotent: deleting non-existent keys is a no-op, so re-runs are safe.
Usage:
uv run python scripts/purge_page_images.py [--dry-run] [--batch-size 256]
Connection settings (Qdrant URL/API key, collection name) are read from
the same `Settings` object the server uses.
"""
from __future__ import annotations
import argparse
import logging
import sys
from functools import partial
import anyio
from qdrant_client import AsyncQdrantClient
from nextcloud_mcp_server.config import get_settings
logger = logging.getLogger("purge_page_images")
LEGACY_FIELDS = [
"highlighted_page_image",
"highlighted_page_number",
"highlight_count",
]
async def purge(dry_run: bool, batch_size: int) -> None:
settings = get_settings()
if not settings.qdrant_url:
raise SystemExit(
"qdrant_url is not configured. Set QDRANT_URL (and QDRANT_API_KEY "
"if required) before running this script."
)
collection = settings.get_collection_name()
# AsyncQdrantClient doesn't implement __aenter__/__aexit__, so use
# try/finally to guarantee the underlying aiohttp session is closed.
client = AsyncQdrantClient(
url=settings.qdrant_url,
api_key=settings.qdrant_api_key,
timeout=60,
)
try:
next_offset = None
total_seen = 0
total_updated = 0
logger.info(
"Scanning collection %s; will delete keys %s%s",
collection,
LEGACY_FIELDS,
" (dry run)" if dry_run else "",
)
while True:
points, next_offset = await client.scroll(
collection_name=collection,
limit=batch_size,
offset=next_offset,
with_payload=False,
with_vectors=False,
)
if not points:
break
ids = [p.id for p in points]
total_seen += len(ids)
if not dry_run:
await client.delete_payload(
collection_name=collection,
keys=LEGACY_FIELDS,
points=ids,
)
total_updated += len(ids)
logger.info(
"Batch: ids=%d total_seen=%d total_updated=%d",
len(ids),
total_seen,
total_updated,
)
if next_offset is None:
break
logger.info(
"Done. total_seen=%d total_updated=%d%s",
total_seen,
total_updated,
" (dry run, no writes)" if dry_run else "",
)
finally:
await client.close()
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--dry-run",
action="store_true",
help="Scan only; do not write any changes.",
)
parser.add_argument(
"--batch-size",
type=int,
default=256,
help="Points per scroll/update batch (default: 256).",
)
parser.add_argument(
"-v", "--verbose", action="store_true", help="Enable debug logging."
)
args = parser.parse_args()
logging.basicConfig(
level=logging.DEBUG if args.verbose else logging.INFO,
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
anyio.run(partial(purge, dry_run=args.dry_run, batch_size=args.batch_size))
return 0
if __name__ == "__main__":
sys.exit(main())
@@ -0,0 +1,232 @@
"""Unit tests for PDFHighlighter.compute_chunk_bboxes_batch (Deck #76).
Replaces the legacy `highlight_chunks_batch`-+-base64 pipeline that inflated
Qdrant payloads with per-chunk PNG screenshots. The new path returns
normalized bounding boxes only.
"""
from __future__ import annotations
import pymupdf
import pytest
from nextcloud_mcp_server.search.pdf_highlighter import PDFHighlighter
def _make_pdf(pages: list[str]) -> bytes:
"""Build an in-memory PDF whose pages contain the given text."""
doc = pymupdf.open()
for body in pages:
page = doc.new_page(width=595, height=842) # A4
page.insert_text((50, 50), body)
pdf_bytes = doc.tobytes()
doc.close()
return pdf_bytes
def _page_boundaries(pages: list[str]) -> tuple[list[dict], str]:
"""Build (page_boundaries, full_text) compatible with the highlighter API."""
boundaries: list[dict] = []
cursor = 0
parts: list[str] = []
for i, body in enumerate(pages, start=1):
end = cursor + len(body)
boundaries.append({"page": i, "start_offset": cursor, "end_offset": end})
parts.append(body)
cursor = end
return boundaries, "".join(parts)
@pytest.mark.unit
def test_compute_chunk_bboxes_returns_normalized_rects():
"""Each returned bbox should be 4 floats in [0, 1] tagged with the page."""
pages = [
"Chapter 1: Introduction. Nextcloud is a self-hosted collaboration platform "
"covering installation, configuration and maintenance topics.",
"Chapter 2: Installation. Download the package, extract it to the web "
"server directory, and configure the database connection.",
]
pdf_bytes = _make_pdf(pages)
boundaries, full_text = _page_boundaries(pages)
chunks = [
(
0,
0,
len(pages[0]),
1,
"Chapter 1: Introduction. Nextcloud is a self-hosted collaboration platform.",
),
(
1,
len(pages[0]),
len(pages[0]) + len(pages[1]),
2,
"Chapter 2: Installation. Download the package.",
),
]
results = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=boundaries,
full_text=full_text,
)
assert set(results) == {0, 1}
bboxes_p1, page_p1 = results[0]
bboxes_p2, page_p2 = results[1]
assert page_p1 == 1
assert page_p2 == 2
for rects in (bboxes_p1, bboxes_p2):
assert len(rects) >= 1
for rect in rects:
assert len(rect) == 4
x0, y0, x1, y1 = rect
assert 0.0 <= x0 < x1 <= 1.0
assert 0.0 <= y0 < y1 <= 1.0
@pytest.mark.unit
def test_compute_chunk_bboxes_empty_input():
assert (
PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=b"",
chunks=[],
page_boundaries=[],
full_text="",
)
== {}
)
@pytest.mark.unit
def test_compute_chunk_bboxes_omits_when_offsets_out_of_range():
"""Chunks whose offsets fall outside every page boundary are omitted.
Verifies the docstring contract: *"Chunks whose bbox cannot be located
are omitted from the result."* (path: ``find_chunk_page`` returns None).
"""
pages = ["Page one body text content here for the test."]
pdf_bytes = _make_pdf(pages)
boundaries, full_text = _page_boundaries(pages)
# Offsets way beyond the document end — no page boundary matches.
out_of_range_start = len(full_text) + 1000
out_of_range_end = out_of_range_start + 50
chunks = [(0, out_of_range_start, out_of_range_end, 1, "irrelevant")]
results = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=boundaries,
full_text=full_text,
)
assert results == {}
@pytest.mark.unit
def test_compute_chunk_bboxes_omits_when_text_not_in_pdf():
"""Chunks whose page-relative text isn't on the page are omitted.
Verifies the second omission path: ``_find_chunk_bbox`` returns None
when the supplied text cannot be located on the rendered page.
"""
pages = ["Hello world."]
pdf_bytes = _make_pdf(pages)
# Build boundaries from the real text but pass a *different* full_text
# so the page-relative slice is content that does not exist in the PDF.
boundaries, _ = _page_boundaries(pages)
bogus_full_text = "Z" * len(pages[0])
chunks = [(0, 0, len(pages[0]), 1, "ignored")]
results = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=boundaries,
full_text=bogus_full_text,
)
assert results == {}
@pytest.mark.unit
@pytest.mark.parametrize("page_index", [0, 1])
def test_compute_chunk_bboxes_assigns_correct_page(page_index: int):
"""Verify the page number returned matches the page the chunk lives on."""
pages = [
"Page one talks about apples and oranges in detail.",
"Page two discusses bananas and grapes thoroughly.",
]
pdf_bytes = _make_pdf(pages)
boundaries, full_text = _page_boundaries(pages)
if page_index == 0:
chunk_text = "apples and oranges"
offsets = (0, len(pages[0]))
else:
chunk_text = "bananas and grapes"
offsets = (len(pages[0]), len(pages[0]) + len(pages[1]))
chunks = [(0, offsets[0], offsets[1], page_index + 1, chunk_text)]
results = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=boundaries,
full_text=full_text,
)
assert results, "expected a bbox for the chunk"
_, page_num = results[0]
assert page_num == page_index + 1
@pytest.mark.unit
def test_compute_chunk_bboxes_handles_unordered_page_boundaries():
"""Page lookup must match by ``page`` key, not by list position.
Regression guard: an earlier implementation indexed
``page_boundaries[page_num - 1]``, which silently produces a wrong
bbox if boundaries are passed out of order. Reverse the boundaries
and assert the result is identical to the in-order case.
"""
pages = [
"Page one talks about apples and oranges in detail.",
"Page two discusses bananas and grapes thoroughly.",
]
pdf_bytes = _make_pdf(pages)
boundaries, full_text = _page_boundaries(pages)
chunks = [
(0, 0, len(pages[0]), 1, "apples and oranges"),
(
1,
len(pages[0]),
len(pages[0]) + len(pages[1]),
2,
"bananas and grapes",
),
]
in_order = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=boundaries,
full_text=full_text,
)
reversed_order = PDFHighlighter.compute_chunk_bboxes_batch(
pdf_bytes=pdf_bytes,
chunks=chunks,
page_boundaries=list(reversed(boundaries)),
full_text=full_text,
)
assert in_order == reversed_order
assert reversed_order[0][1] == 1
assert reversed_order[1][1] == 2