chatkit/backend/app/parsers/base.py

"""Parser base types and shared file utilities."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Protocol

TEXT_EXTENSIONS = {".txt", ".csv", ".tsv", ".json", ".xml", ".md", ".html", ".htm", ".log"}

IMAGE_MIME = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".webp": "image/webp",
    ".tiff": "image/tiff",
    ".tif": "image/tiff",
}


@dataclass(frozen=True)
class ParserFile:
    """A file uploaded for parsing."""

    filename: str
    content: bytes


@dataclass(frozen=True)
class ParserResult:
    """Structured output of a parser run."""

    data: dict
    parser_name: str
    notes: tuple[str, ...] = field(default_factory=tuple)


class AnswerSheetParser(Protocol):
    """Protocol for answer-sheet parsers.

    Implementations should be stateless — one instance per process is fine.
    """

    name: str
    display_name: str
    description: str

    def can_handle(self, file_bytes: bytes, filename: str, data_type: str) -> bool: ...

    async def parse(
        self,
        files: list[ParserFile],
        data_type: str,
        description: str = "",
        model: str = "gpt-5.4",
    ) -> ParserResult: ...


# ---------------------------------------------------------------------------
# Shared file helpers
# ---------------------------------------------------------------------------


def get_extension(filename: str) -> str:
    return Path(filename).suffix.lower()


def is_pdf(filename: str) -> bool:
    return get_extension(filename) == ".pdf"


def is_image(filename: str) -> bool:
    return get_extension(filename) in IMAGE_MIME


def get_image_mime(filename: str) -> Optional[str]:
    return IMAGE_MIME.get(get_extension(filename))


def try_read_as_text(file_bytes: bytes, filename: str) -> Optional[str]:
    """Decode text files using a list of likely encodings."""
    if get_extension(filename) not in TEXT_EXTENSIONS:
        return None
    for encoding in ("utf-8", "utf-8-sig", "big5", "gb2312", "shift_jis", "latin-1"):
        try:
            return file_bytes.decode(encoding)
        except (UnicodeDecodeError, ValueError):
            continue
    return None


def pdf_has_text_layer(pdf_bytes: bytes) -> bool:
    """True if the PDF has an extractable text layer (i.e., not an image scan)."""
    try:
        import fitz  # local import so parsers without pymupdf don't die
    except ImportError:
        return False
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        try:
            for page in doc:
                text = page.get_text("text").strip()
                if text:
                    return True
            return False
        finally:
            doc.close()
    except Exception:
        return False