gauravmeena0708/epfo-circulars

1"""Document-ingestion and prompt helpers for the uploaded-file assistant."""2 3from __future__ import annotations4 5from dataclasses import dataclass6import hashlib7import re8from typing import Callable, Iterable, Mapping, Sequence9 10import numpy as np11 12try:13    import pymupdf as fitz14except ImportError:  # pragma: no cover - compatibility with older PyMuPDF releases15    import fitz16 17 18PAGE_MARKER_PATTERN = re.compile(r"(?=^--- \[Page \d+ / )", re.MULTILINE)19WORD_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{2,}")20QUERY_STOP_WORDS = {21    "about", "all", "and", "answer", "document", "file", "from", "have",22    "noting", "page", "provide", "sheet", "that", "the", "this", "what",23    "when", "where", "which", "with",24}25 26 27class DocumentExtractionError(ValueError):28    """A safe, user-facing PDF extraction failure."""29 30 31@dataclass(frozen=True)32class ExtractionResult:33    text: str34    page_count: int35    scanned_page_count: int36    warnings: tuple[str, ...]37 38 39def uploaded_file_signature(filename: str, data: bytes) -> str:40    """Return a stable signature that changes when either name or bytes change."""41    digest = hashlib.sha256()42    digest.update(filename.encode("utf-8", errors="replace"))43    digest.update(b"\0")44    digest.update(data)45    return digest.hexdigest()46 47 48def _normalise_ocr_results(results: Iterable[object]) -> str:49    lines = []50    for item in results:51        if isinstance(item, str):52            text = item53        elif isinstance(item, (list, tuple)) and len(item) >= 2:54            text = item[1]55        else:56            continue57        text = str(text).strip()58        if text:59            lines.append(text)60    return "\n".join(lines)61 62 63def _ocr_page(page, ocr_reader, dpi: int) -> str:64    """Render a PyMuPDF page and run EasyOCR without requiring Poppler."""65    scale = max(dpi, 72) / 7266    pixmap = page.get_pixmap(matrix=fitz.Matrix(scale, scale), alpha=False)67    channels = pixmap.n68    image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(69        pixmap.height,70        pixmap.width,71        channels,72    )73    results = ocr_reader.readtext(image, detail=0, paragraph=True)74    return _normalise_ocr_results(results)75 76 77def extract_pdf_text(78    pdf_data: bytes,79    *,80    ocr_reader_factory: Callable[[], object] | None = None,81    native_text_min_words: int = 25,82    ocr_dpi: int = 200,83) -> ExtractionResult:84    """Extract page-labelled text with lazy OCR and guaranteed temp cleanup."""85    if not pdf_data:86        raise DocumentExtractionError("The uploaded PDF is empty.")87 88    try:89        document = fitz.open(stream=pdf_data, filetype="pdf")90    except Exception as exc:91        raise DocumentExtractionError(92            "The uploaded file is not a readable PDF or is damaged."93        ) from exc94 95    with document:96        if document.needs_pass:97            raise DocumentExtractionError(98                "This PDF is password-protected. Upload an unlocked copy."99            )100 101        page_count = len(document)102        if page_count == 0:103            raise DocumentExtractionError("The uploaded PDF has no pages.")104 105        extracted_pages: list[str | None] = [None] * page_count106        native_fallbacks: dict[int, str] = {}107        scanned_indices: list[int] = []108 109        for page_number in range(page_count):110            page = document[page_number]111            native_text = (page.get_text("text") or "").strip()112            marker = f"--- [Page {page_number + 1} / Note Page] ---"113            if len(native_text.split()) >= native_text_min_words:114                extracted_pages[page_number] = f"{marker}\n{native_text}"115            else:116                native_fallbacks[page_number] = native_text117                scanned_indices.append(page_number)118 119        warnings: list[str] = []120        ocr_reader = None121        if scanned_indices and ocr_reader_factory is not None:122            try:123                ocr_reader = ocr_reader_factory()124            except Exception:125                warnings.append(126                    "OCR could not be initialized. Short native text was preserved "127                    "where available, but some scanned pages may be unreadable."128                )129 130        unreadable_pages: list[int] = []131        for page_number in scanned_indices:132            ocr_text = ""133            if ocr_reader is not None:134                try:135                    ocr_text = _ocr_page(document[page_number], ocr_reader, ocr_dpi)136                except Exception:137                    unreadable_pages.append(page_number + 1)138 139            fallback_text = native_fallbacks.get(page_number, "")140            content = ocr_text or fallback_text141            page_type = "Scanned Page" if ocr_text else "Short/Scanned Page"142            if not content:143                if page_number + 1 not in unreadable_pages:144                    unreadable_pages.append(page_number + 1)145                content = "[No readable text could be extracted from this page.]"146            extracted_pages[page_number] = (147                f"--- [Page {page_number + 1} / {page_type}] ---\n{content}"148            )149 150        if unreadable_pages:151            page_list = ", ".join(str(number) for number in unreadable_pages)152            warnings.append(f"No readable text was found on page(s): {page_list}.")153 154        full_text = "\n\n".join(page for page in extracted_pages if page)155        return ExtractionResult(156            text=full_text,157            page_count=page_count,158            scanned_page_count=len(scanned_indices),159            warnings=tuple(warnings),160        )161 162 163def _truncate_middle(text: str, max_chars: int) -> str:164    if len(text) <= max_chars:165        return text166    separator = "\n\n[... content omitted to fit the model context ...]\n\n"167    available = max(0, max_chars - len(separator))168    head = available // 2169    tail = available - head170    return f"{text[:head]}{separator}{text[-tail:] if tail else ''}"171 172 173def select_document_context(174    full_text: str,175    query: str,176    max_chars: int,177) -> tuple[str, bool]:178    """Bound document input while retaining boundary and query-relevant pages."""179    max_chars = max(2_000, int(max_chars))180    if len(full_text) <= max_chars:181        return full_text, False182 183    page_blocks = [block.strip() for block in PAGE_MARKER_PATTERN.split(full_text) if block.strip()]184    if len(page_blocks) <= 1:185        return _truncate_middle(full_text, max_chars), True186 187    query_terms = {188        term.lower()189        for term in WORD_PATTERN.findall(query)190        if term.lower() not in QUERY_STOP_WORDS191    }192 193    scored_indices = []194    for index, block in enumerate(page_blocks):195        lowered = block.lower()196        relevance = sum(lowered.count(term) for term in query_terms)197        scored_indices.append((relevance, index))198 199    candidates = sorted(scored_indices, key=lambda item: (-item[0], item[1]))200    relevant_indices = [index for score, index in candidates if score > 0][:3]201    mandatory_indices = {0, len(page_blocks) - 1, *relevant_indices}202    # With broad prompts, evenly distributed candidates give summaries better coverage.203    evenly_spaced = sorted(204        {205            round(position * (len(page_blocks) - 1) / 7)206            for position in range(1, 7)207        }208    )209    candidate_indices = [index for _, index in candidates] + evenly_spaced210 211    notice = (212        "[Document context was limited for this request. The first, last, and "213        "most query-relevant pages are included below.]\n\n"214    )215    budget = max_chars - len(notice)216    separator_cost = 2 * max(0, len(mandatory_indices) - 1)217    per_mandatory_budget = max(218        200,219        (budget - separator_cost) // max(1, len(mandatory_indices)),220    )221    rendered_blocks = {222        index: _truncate_middle(page_blocks[index], per_mandatory_budget)223        for index in mandatory_indices224    }225    used_chars = sum(len(block) + 2 for block in rendered_blocks.values())226 227    for index in candidate_indices:228        if index in rendered_blocks:229            continue230        block_size = len(page_blocks[index]) + 2231        if used_chars + block_size <= budget:232            rendered_blocks[index] = page_blocks[index]233            used_chars += block_size234 235    ordered_blocks = [rendered_blocks[index] for index in sorted(rendered_blocks)]236    context = notice + "\n\n".join(ordered_blocks)237    return _truncate_middle(context, max_chars), True238 239 240def format_conversation_history(241    history: Sequence[Mapping[str, object]],242    *,243    max_chars: int,244    max_messages: int = 8,245) -> str:246    """Return the most recent bounded conversation turns for follow-up questions."""247    formatted: list[str] = []248    for message in history[-max_messages:]:249        role = "User" if message.get("role") == "user" else "Assistant"250        content = str(message.get("content", "")).strip()251        if content:252            formatted.append(f"{role}: {content}")253    return _truncate_middle("\n\n".join(formatted), max(500, max_chars))254