Spaces:
Running
Running
| """Document-ingestion and prompt helpers for the uploaded-file assistant.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import hashlib | |
| import re | |
| from typing import Callable, Iterable, Mapping, Sequence | |
| import numpy as np | |
| try: | |
| import pymupdf as fitz | |
| except ImportError: # pragma: no cover - compatibility with older PyMuPDF releases | |
| import fitz | |
| PAGE_MARKER_PATTERN = re.compile(r"(?=^--- \[Page \d+ / )", re.MULTILINE) | |
| WORD_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{2,}") | |
| QUERY_STOP_WORDS = { | |
| "about", "all", "and", "answer", "document", "file", "from", "have", | |
| "noting", "page", "provide", "sheet", "that", "the", "this", "what", | |
| "when", "where", "which", "with", | |
| } | |
| class DocumentExtractionError(ValueError): | |
| """A safe, user-facing PDF extraction failure.""" | |
| class ExtractionResult: | |
| text: str | |
| page_count: int | |
| scanned_page_count: int | |
| warnings: tuple[str, ...] | |
| def uploaded_file_signature(filename: str, data: bytes) -> str: | |
| """Return a stable signature that changes when either name or bytes change.""" | |
| digest = hashlib.sha256() | |
| digest.update(filename.encode("utf-8", errors="replace")) | |
| digest.update(b"\0") | |
| digest.update(data) | |
| return digest.hexdigest() | |
| def _normalise_ocr_results(results: Iterable[object]) -> str: | |
| lines = [] | |
| for item in results: | |
| if isinstance(item, str): | |
| text = item | |
| elif isinstance(item, (list, tuple)) and len(item) >= 2: | |
| text = item[1] | |
| else: | |
| continue | |
| text = str(text).strip() | |
| if text: | |
| lines.append(text) | |
| return "\n".join(lines) | |
| def _ocr_page(page, ocr_reader, dpi: int) -> str: | |
| """Render a PyMuPDF page and run EasyOCR without requiring Poppler.""" | |
| scale = max(dpi, 72) / 72 | |
| pixmap = page.get_pixmap(matrix=fitz.Matrix(scale, scale), alpha=False) | |
| channels = pixmap.n | |
| image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape( | |
| pixmap.height, | |
| pixmap.width, | |
| channels, | |
| ) | |
| results = ocr_reader.readtext(image, detail=0, paragraph=True) | |
| return _normalise_ocr_results(results) | |
| def extract_pdf_text( | |
| pdf_data: bytes, | |
| *, | |
| ocr_reader_factory: Callable[[], object] | None = None, | |
| native_text_min_words: int = 25, | |
| ocr_dpi: int = 200, | |
| ) -> ExtractionResult: | |
| """Extract page-labelled text with lazy OCR and guaranteed temp cleanup.""" | |
| if not pdf_data: | |
| raise DocumentExtractionError("The uploaded PDF is empty.") | |
| try: | |
| document = fitz.open(stream=pdf_data, filetype="pdf") | |
| except Exception as exc: | |
| raise DocumentExtractionError( | |
| "The uploaded file is not a readable PDF or is damaged." | |
| ) from exc | |
| with document: | |
| if document.needs_pass: | |
| raise DocumentExtractionError( | |
| "This PDF is password-protected. Upload an unlocked copy." | |
| ) | |
| page_count = len(document) | |
| if page_count == 0: | |
| raise DocumentExtractionError("The uploaded PDF has no pages.") | |
| extracted_pages: list[str | None] = [None] * page_count | |
| native_fallbacks: dict[int, str] = {} | |
| scanned_indices: list[int] = [] | |
| for page_number in range(page_count): | |
| page = document[page_number] | |
| native_text = (page.get_text("text") or "").strip() | |
| marker = f"--- [Page {page_number + 1} / Note Page] ---" | |
| if len(native_text.split()) >= native_text_min_words: | |
| extracted_pages[page_number] = f"{marker}\n{native_text}" | |
| else: | |
| native_fallbacks[page_number] = native_text | |
| scanned_indices.append(page_number) | |
| warnings: list[str] = [] | |
| ocr_reader = None | |
| if scanned_indices and ocr_reader_factory is not None: | |
| try: | |
| ocr_reader = ocr_reader_factory() | |
| except Exception: | |
| warnings.append( | |
| "OCR could not be initialized. Short native text was preserved " | |
| "where available, but some scanned pages may be unreadable." | |
| ) | |
| unreadable_pages: list[int] = [] | |
| for page_number in scanned_indices: | |
| ocr_text = "" | |
| if ocr_reader is not None: | |
| try: | |
| ocr_text = _ocr_page(document[page_number], ocr_reader, ocr_dpi) | |
| except Exception: | |
| unreadable_pages.append(page_number + 1) | |
| fallback_text = native_fallbacks.get(page_number, "") | |
| content = ocr_text or fallback_text | |
| page_type = "Scanned Page" if ocr_text else "Short/Scanned Page" | |
| if not content: | |
| if page_number + 1 not in unreadable_pages: | |
| unreadable_pages.append(page_number + 1) | |
| content = "[No readable text could be extracted from this page.]" | |
| extracted_pages[page_number] = ( | |
| f"--- [Page {page_number + 1} / {page_type}] ---\n{content}" | |
| ) | |
| if unreadable_pages: | |
| page_list = ", ".join(str(number) for number in unreadable_pages) | |
| warnings.append(f"No readable text was found on page(s): {page_list}.") | |
| full_text = "\n\n".join(page for page in extracted_pages if page) | |
| return ExtractionResult( | |
| text=full_text, | |
| page_count=page_count, | |
| scanned_page_count=len(scanned_indices), | |
| warnings=tuple(warnings), | |
| ) | |
| def _truncate_middle(text: str, max_chars: int) -> str: | |
| if len(text) <= max_chars: | |
| return text | |
| separator = "\n\n[... content omitted to fit the model context ...]\n\n" | |
| available = max(0, max_chars - len(separator)) | |
| head = available // 2 | |
| tail = available - head | |
| return f"{text[:head]}{separator}{text[-tail:] if tail else ''}" | |
| def select_document_context( | |
| full_text: str, | |
| query: str, | |
| max_chars: int, | |
| ) -> tuple[str, bool]: | |
| """Bound document input while retaining boundary and query-relevant pages.""" | |
| max_chars = max(2_000, int(max_chars)) | |
| if len(full_text) <= max_chars: | |
| return full_text, False | |
| page_blocks = [block.strip() for block in PAGE_MARKER_PATTERN.split(full_text) if block.strip()] | |
| if len(page_blocks) <= 1: | |
| return _truncate_middle(full_text, max_chars), True | |
| query_terms = { | |
| term.lower() | |
| for term in WORD_PATTERN.findall(query) | |
| if term.lower() not in QUERY_STOP_WORDS | |
| } | |
| scored_indices = [] | |
| for index, block in enumerate(page_blocks): | |
| lowered = block.lower() | |
| relevance = sum(lowered.count(term) for term in query_terms) | |
| scored_indices.append((relevance, index)) | |
| candidates = sorted(scored_indices, key=lambda item: (-item[0], item[1])) | |
| relevant_indices = [index for score, index in candidates if score > 0][:3] | |
| mandatory_indices = {0, len(page_blocks) - 1, *relevant_indices} | |
| # With broad prompts, evenly distributed candidates give summaries better coverage. | |
| evenly_spaced = sorted( | |
| { | |
| round(position * (len(page_blocks) - 1) / 7) | |
| for position in range(1, 7) | |
| } | |
| ) | |
| candidate_indices = [index for _, index in candidates] + evenly_spaced | |
| notice = ( | |
| "[Document context was limited for this request. The first, last, and " | |
| "most query-relevant pages are included below.]\n\n" | |
| ) | |
| budget = max_chars - len(notice) | |
| separator_cost = 2 * max(0, len(mandatory_indices) - 1) | |
| per_mandatory_budget = max( | |
| 200, | |
| (budget - separator_cost) // max(1, len(mandatory_indices)), | |
| ) | |
| rendered_blocks = { | |
| index: _truncate_middle(page_blocks[index], per_mandatory_budget) | |
| for index in mandatory_indices | |
| } | |
| used_chars = sum(len(block) + 2 for block in rendered_blocks.values()) | |
| for index in candidate_indices: | |
| if index in rendered_blocks: | |
| continue | |
| block_size = len(page_blocks[index]) + 2 | |
| if used_chars + block_size <= budget: | |
| rendered_blocks[index] = page_blocks[index] | |
| used_chars += block_size | |
| ordered_blocks = [rendered_blocks[index] for index in sorted(rendered_blocks)] | |
| context = notice + "\n\n".join(ordered_blocks) | |
| return _truncate_middle(context, max_chars), True | |
| def format_conversation_history( | |
| history: Sequence[Mapping[str, object]], | |
| *, | |
| max_chars: int, | |
| max_messages: int = 8, | |
| ) -> str: | |
| """Return the most recent bounded conversation turns for follow-up questions.""" | |
| formatted: list[str] = [] | |
| for message in history[-max_messages:]: | |
| role = "User" if message.get("role") == "user" else "Assistant" | |
| content = str(message.get("content", "")).strip() | |
| if content: | |
| formatted.append(f"{role}: {content}") | |
| return _truncate_middle("\n\n".join(formatted), max(500, max_chars)) | |