yau-plant-assistant/api/convert.py

"""Uploaded file -> Markdown, before anything is chunked or embedded.

WHY THERE IS A CONVERSION STEP AT ALL. The raw PDF is never what the assistant
reads - it never was. Retrieval only ever sees `doc_chunks.chunk_text`. What
this module adds is that the extracted text becomes a thing a person can LOOK
AT before approving it. If a table comes out as garbage, the reviewer sees the
garbage and rejects the document, instead of an operator discovering it months
later inside a citation.

WHAT THESE CONVERTERS DO AND DO NOT DO. They extract text, not layout. A
multi-column page interleaves. A merged-cell spreadsheet flattens. A scanned
page yields nothing at all, and is refused rather than stored empty.

That limitation is acceptable HERE, and only here, because a person reads the
converted text before it can be cited - the same safety net the design already
required for the document header. It would not be acceptable in a pipeline that
published without review.

REPLACING THEM. Everything below the `convert()` boundary is swappable: return
Markdown from bytes, raise ConversionError when you cannot. Nothing outside this
module knows which library did the work, so a better converter is a change to
this file alone. If the real documents turn out to be scans, that is the
conversation to have - OCR is the missing capability, and nothing else here
changes.
"""

from __future__ import annotations

import io
import re
from dataclasses import dataclass

# A scanned page yields a handful of stray characters, not nothing - so an
# emptiness check has to have a floor above zero. Below this, across the whole
# document, we refuse rather than store a blank document that a reviewer might
# approve without noticing there is nothing in it.
_MIN_CHARS = 200

SUPPORTED = {".pdf", ".docx", ".xlsx", ".xlsm", ".md", ".txt"}


class ConversionError(Exception):
    """Conversion failed, or produced something not worth reviewing."""


@dataclass(frozen=True)
class Converted:
    markdown: str
    converter: str          # recorded on the upload row - which code produced this
    page_count: int | None


def _clean(text: str) -> str:
    """Collapse the whitespace damage that text extraction always leaves.

    Not cosmetic: runs of blank lines and trailing spaces change where a chunker
    splits, so the same document converted twice should look the same.
    """
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"[ \t]+\n", "\n", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def _from_pdf(data: bytes) -> Converted:
    from pypdf import PdfReader

    try:
        reader = PdfReader(io.BytesIO(data))
    except Exception as exc:
        raise ConversionError(f"not a readable PDF: {exc}") from exc

    if reader.is_encrypted:
        # Refuse rather than guess at an empty password. A locked document that
        # silently converts to nothing is the worst outcome here.
        raise ConversionError(
            "this PDF is encrypted - remove the protection and upload it again"
        )

    parts: list[str] = []
    for number, page in enumerate(reader.pages, start=1):
        try:
            text = page.extract_text() or ""
        except Exception:
            # One bad page must not lose the other ninety. Mark it so the
            # reviewer can see exactly what is missing.
            text = ""
            parts.append(f"\n\n<!-- page {number}: could not be extracted -->\n")
        if text.strip():
            parts.append(f"\n\n<!-- page {number} -->\n\n{text}")

    return Converted(
        markdown=_clean("".join(parts)),
        converter="pypdf",
        page_count=len(reader.pages),
    )


def _from_docx(data: bytes) -> Converted:
    import docx

    try:
        document = docx.Document(io.BytesIO(data))
    except Exception as exc:
        raise ConversionError(f"not a readable .docx: {exc}") from exc

    lines: list[str] = []
    for para in document.paragraphs:
        text = para.text.strip()
        if not text:
            continue
        # Word's built-in heading styles are the one piece of structure that
        # survives reliably, and headings matter: section_title is what a
        # citation falls back to.
        style = (para.style.name or "").lower() if para.style else ""
        if style.startswith("heading"):
            level = "".join(c for c in style if c.isdigit()) or "1"
            lines.append(f"{'#' * min(int(level), 6)} {text}")
        elif style.startswith("title"):
            lines.append(f"# {text}")
        else:
            lines.append(text)

    for index, table in enumerate(document.tables, start=1):
        lines.append(f"\n<!-- table {index} -->")
        for row in table.rows:
            cells = [c.text.strip().replace("|", "\\|") for c in row.cells]
            lines.append("| " + " | ".join(cells) + " |")

    return Converted(
        markdown=_clean("\n\n".join(lines)),
        converter="python-docx",
        page_count=None,
    )


def _from_xlsx(data: bytes) -> Converted:
    from openpyxl import load_workbook

    try:
        # read_only keeps a large workbook from being held in memory twice, and
        # data_only takes the cached VALUE of a formula rather than the formula
        # text. A cell reading "=SUM(B2:B9)" is not something to embed - and if
        # the workbook was never opened in Excel there is no cached value, so
        # that cell comes through empty. Say so in the output rather than
        # letting it look like a blank cell.
        book = load_workbook(io.BytesIO(data), read_only=True, data_only=True)
    except Exception as exc:
        raise ConversionError(f"not a readable spreadsheet: {exc}") from exc

    lines: list[str] = []
    for sheet in book.worksheets:
        lines.append(f"## {sheet.title}")
        for row in sheet.iter_rows(values_only=True):
            cells = ["" if v is None else str(v).strip().replace("|", "\\|")
                     for v in row]
            if not any(cells):
                continue
            lines.append("| " + " | ".join(cells) + " |")
        lines.append("")
    book.close()

    return Converted(
        markdown=_clean("\n".join(lines)),
        converter="openpyxl",
        page_count=None,
    )


def _from_text(data: bytes) -> Converted:
    try:
        text = data.decode("utf-8")
    except UnicodeDecodeError:
        text = data.decode("latin-1")
    return Converted(markdown=_clean(text), converter="passthrough", page_count=None)


_CONVERTERS = {
    ".pdf": _from_pdf,
    ".docx": _from_docx,
    ".xlsx": _from_xlsx,
    ".xlsm": _from_xlsx,
    ".md": _from_text,
    ".txt": _from_text,
}


def convert(filename: str, data: bytes) -> Converted:
    """Convert an uploaded file to Markdown, or raise ConversionError.

    Dispatch is on the extension, which is a claim the uploader made about the
    file. The libraries below all fail loudly on a mismatch, so a .docx renamed
    to .pdf raises rather than producing plausible rubbish.
    """
    suffix = "." + filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    handler = _CONVERTERS.get(suffix)
    if handler is None:
        raise ConversionError(
            f"{suffix or 'no extension'} is not supported - "
            f"expected one of {', '.join(sorted(SUPPORTED))}"
        )

    result = handler(data)

    if len(result.markdown) < _MIN_CHARS:
        # The common cause by far is a scanned document: a picture of text,
        # which text extraction cannot read. Name that explicitly, because
        # "conversion produced nothing" sends somebody looking for a bug in the
        # upload instead of at the file.
        raise ConversionError(
            f"only {len(result.markdown)} characters of text came out of this "
            f"file. If it is a scan or a photograph, no text can be extracted "
            f"from it - it needs OCR, which this converter does not do. "
            f"Nothing has been added to the library."
        )

    return result