Blame for document_assistant.py

"""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."""


@dataclass(frozen=True)
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))