ss-tools/backend/src/agent/document_parser.py

# backend/src/agent/document_parser.py
# #region AgentChat.Document.Parser [C:3] [TYPE Module] [SEMANTICS agent-chat,document,parser]
# @defgroup AgentChat Parse PDF and XLSX files into text/structured data.
# @RELATION DEPENDS_ON -> [EXT:pdfplumber]
# @RELATION DEPENDS_ON -> [EXT:openpyxl]
# @PRE File exists, valid format, ≤10MB.
# @POST Returns extracted text (PDF) or structured dict (XLSX).

from pathlib import Path


class ParseError(Exception):
    """Raised when document parsing fails."""


def parse_pdf(file_path: str) -> str:
    """Extract text from PDF using pdfplumber (primary) with PyPDF2 fallback."""
    try:
        import pdfplumber
    except ImportError:
        raise ParseError("pdfplumber not installed") from None

    try:
        with pdfplumber.open(file_path) as pdf:
            pages = []
            for page in pdf.pages:
                text = page.extract_text()
                if text:
                    pages.append(text)
            return "\n\n".join(pages) if pages else ""
    except Exception as e:
        # Fallback to PyPDF2
        try:
            import PyPDF2
            with open(file_path, "rb") as f:
                reader = PyPDF2.PdfReader(f)
                return "\n\n".join(p.extract_text() for p in reader.pages if p.extract_text())
        except Exception:
            raise ParseError(f"Failed to parse PDF: {e}") from None


def parse_xlsx(file_path: str) -> str:
    """Extract structured data from XLSX — sheet names + cell data."""
    try:
        import openpyxl
    except ImportError:
        raise ParseError("openpyxl not installed") from None

    try:
        wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
        parts = []
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            rows = []
            for row in ws.iter_rows(values_only=True):
                cells = [str(c) if c is not None else "" for c in row]
                rows.append("\t".join(cells))
            parts.append(f"=== Sheet: {sheet_name} ===\n" + "\n".join(rows))
        return "\n\n".join(parts)
    except Exception as e:
        raise ParseError(f"Failed to parse XLSX: {e}") from e


def _detect_format_by_magic(path: str) -> str | None:
    """Detect file format by reading magic bytes. Returns extension or None."""
    try:
        with open(path, "rb") as f:
            header = f.read(8)
    except OSError:
        return None

    if header[:4] == b"%PDF":
        return ".pdf"
    if header[:4] == b"PK\x03\x04":
        # ZIP-based: XLSX, DOCX, etc. Try XLSX first.
        return ".xlsx"
    if header[:1] in (b"{", b"["):
        return ".json"
    # CSV often starts with text characters; safe fallback
    return None


def parse_upload(file_data) -> str:
    """Parse an uploaded file based on its extension with magic-byte fallback.

    Args:
        file_data: str (file path) or dict with "name" and "path"/"file_path" keys.
    """
    if isinstance(file_data, str):
        path = file_data
        name = Path(path).name
    else:
        name = file_data.get("name") or file_data.get("orig_name", "")
        path = file_data.get("path") or file_data.get("file_path", "")
        # Gradio sometimes sends files without 'name' key — fall back to path stem
        if not name and path:
            name = Path(path).name
    ext = Path(name).suffix.lower()
    # Double fallback: if name has no extension, try the physical file path
    if not ext and path:
        ext = Path(path).suffix.lower()

    # Triple fallback: detect by magic bytes (Gradio ChatInterface loses filename)
    if not ext and path:
        ext = _detect_format_by_magic(path)

    if ext == ".pdf":
        return parse_pdf(path)
    elif ext in (".xlsx", ".xls"):
        return parse_xlsx(path)
    elif ext in (".json", ".csv", ".txt"):
        with open(path, encoding="utf-8", errors="replace") as f:
            return f.read(100_000)  # truncate at ~100k chars
    elif ext is None:
        # Magic bytes detection returned None — unknown binary, try as text
        try:
            with open(path, encoding="utf-8", errors="replace") as f:
                return f.read(100_000)
        except Exception as e:
            raise ParseError(
                f"Could not detect file format for '{name}'. "
                f"Supported: PDF, XLSX, JSON, CSV, TXT"
            ) from e
    else:
        raise ParseError(
            f"Unsupported format: '{ext}' (file: {name}). "
            f"Supported: PDF, XLSX, JSON, CSV, TXT"
        )
# #endregion AgentChat.Document.Parser