diff options
Diffstat (limited to 'src/ollama_web/tools/helper/pdf.py')
| -rw-r--r-- | src/ollama_web/tools/helper/pdf.py | 155 |
1 files changed, 97 insertions, 58 deletions
diff --git a/src/ollama_web/tools/helper/pdf.py b/src/ollama_web/tools/helper/pdf.py index 7b9fdac..243a837 100644 --- a/src/ollama_web/tools/helper/pdf.py +++ b/src/ollama_web/tools/helper/pdf.py @@ -1,4 +1,4 @@ -"""PDF text extraction helpers.""" +"""Lossless PDF text extraction helpers.""" from __future__ import annotations @@ -7,78 +7,117 @@ import io from pypdf import PdfReader from ...config import settings +from .pagination import PageResult -def extract_pdf_text( +def extract_pdf_page( data: bytes, - max_chars: int | None = None, - page_offset: int = 0, -) -> tuple[str, bool, int]: - """Extract text from PDF bytes, starting at ``page_offset``. - - Args: - data: Raw PDF bytes. - max_chars: Maximum characters to return. Defaults to - ``settings.fetch_max_chars``. - page_offset: Zero-based index of the first page to extract. Pages before - this index are skipped. Used by the pagination layer to resume. - - Returns: - A tuple ``(text, has_more, next_page_offset)`` where: - - ``text`` is the extracted text for this page (truncated to - ``max_chars``). - - ``has_more`` is True when there are remaining pages after the - extracted slice. - - ``next_page_offset`` is the page index to resume from on the next - call (only meaningful when ``has_more`` is True). - """ + *, + max_chars: int, + page_index: int = 0, + page_char_offset: int = 0, + max_size_mb: int | None = None, +) -> PageResult: + """Extract a PDF slice using a page-and-character continuation cursor.""" + start = {"page_index": page_index, "page_char_offset": page_char_offset} if not data: - return "Empty PDF data.", False, 0 - - limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars - - if len(data) > settings.max_fetch_mb * 1024 * 1024: - return "PDF exceeds the configured size limit.", False, 0 + return PageResult("Empty PDF data.", start, start, False) + size_limit_mb = settings.max_fetch_mb if max_size_mb is None else max_size_mb + if len(data) > size_limit_mb * 1024 * 1024: + return PageResult("PDF exceeds the configured size limit.", start, start, False) + if page_index < 0 or page_char_offset < 0: + return PageResult("Invalid PDF continuation cursor.", start, start, False) try: reader = PdfReader(stream=io.BytesIO(data)) except Exception as exc: # noqa: BLE001 - return f"Failed to parse PDF: {exc}", False, 0 + return PageResult(f"Failed to parse PDF: {exc}", start, start, False) total_pages = len(reader.pages) - if page_offset < 0: - page_offset = 0 - if page_offset >= total_pages: - return "No text could be extracted from the PDF.", False, 0 - - # Cap the number of pages processed per call to max_pdf_pages. - end = min(page_offset + settings.max_pdf_pages, total_pages) + if page_index > total_pages or (page_index == total_pages and page_char_offset): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + if page_index == total_pages: + return PageResult("", start, start, False) parts: list[str] = [] - collected = 0 - last_extracted_page = page_offset - for page_index in range(page_offset, end): + warnings: list[str] = [] + remaining = max(1, max_chars) + current_page = page_index + current_char = page_char_offset + pages_processed = 0 + + while current_page < total_pages and pages_processed < settings.max_pdf_pages: try: - text = reader.pages[page_index].extract_text() - except Exception: # noqa: BLE001 + page_text = reader.pages[current_page].extract_text() or "" + except Exception as exc: # noqa: BLE001 + detail = str(exc).replace("\n", " ")[:120] + warnings.append(f"PDF page {current_page + 1} could not be extracted: {detail}") + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 continue - if text: - parts.append(text) - collected += len(text) - last_extracted_page = page_index - # Stop as soon as we reach the per-call character budget so the page - # boundary for the next token is accurate. - if collected >= limit: + + if current_char > len(page_text): + return PageResult("Invalid PDF continuation cursor.", start, start, False) + + available = page_text[current_char:] + take = min(remaining, len(available)) + if take: + parts.append(available[:take]) + current_char += take + remaining -= take + + if current_char < len(page_text): + break + + if current_page + 1 < total_pages: + parts.append("\n\n") + current_page += 1 + current_char = 0 + pages_processed += 1 + if remaining <= 0: break - full = "\n\n".join(parts) - if len(full) > limit: - full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]" + next_cursor = {"page_index": current_page, "page_char_offset": current_char} + has_more = current_page < total_pages + text = "".join(parts) + if not text and not warnings and not has_more: + text = "No text could be extracted from the PDF." + return PageResult(text, start, next_cursor, has_more, tuple(warnings)) + + +def legacy_pdf_offset_to_cursor(data: bytes, flat_offset: int) -> tuple[int, int]: + """Translate a v1 ``\n\n``-joined character offset to the v2 PDF cursor.""" + if flat_offset < 0: + raise ValueError("Legacy PDF offset must not be negative.") + reader = PdfReader(stream=io.BytesIO(data)) + remaining = flat_offset + for page_index, page in enumerate(reader.pages): + try: + page_text = page.extract_text() or "" + except Exception: # noqa: BLE001 + page_text = "" + if remaining <= len(page_text): + return page_index, remaining + remaining -= len(page_text) + if page_index + 1 < len(reader.pages): + if remaining <= 2: + return page_index + 1, 0 + remaining -= 2 + if remaining == 0: + return len(reader.pages), 0 + raise ValueError("Legacy PDF offset is out of range.") - # has_more when there are pages beyond the last one we extracted. - has_more = (last_extracted_page + 1) < total_pages - next_page_offset = last_extracted_page + 1 if has_more else 0 - if not full: - return "No text could be extracted from the PDF.", has_more, next_page_offset - return full, has_more, next_page_offset +def extract_pdf_text( + data: bytes, + max_chars: int | None = None, + page_offset: int = 0, +) -> tuple[str, bool, int]: + """Backward-compatible wrapper for page-only internal callers.""" + limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars + page = extract_pdf_page(data, max_chars=limit, page_index=page_offset) + text = "".join(f"[warning] {warning}\n" for warning in page.warnings) + page.text + return text, page.has_more, page.next_cursor["page_index"] if page.has_more else 0 |