app/ingestion.py

"""Document text extraction (with OCR fallback) and chunking."""
from __future__ import annotations

import io
import re
from dataclasses import dataclass
from typing import List, Tuple


@dataclass
class ExtractResult:
    text: str
    ocr_used: bool


def _ocr_image_bytes(data: bytes) -> str:
    """OCR a single image using pytesseract."""
    from PIL import Image
    import pytesseract

    image = Image.open(io.BytesIO(data))
    return pytesseract.image_to_string(image)


def _extract_pdf(data: bytes) -> ExtractResult:
    """Extract text from a PDF. Falls back to OCR for pages with no text layer."""
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(data))
    pages_text: List[str] = []
    ocr_used = False

    for page in reader.pages:
        text = (page.extract_text() or "").strip()
        pages_text.append(text)

    # If the PDF has little/no extractable text, it is likely scanned -> OCR.
    total_chars = sum(len(t) for t in pages_text)
    if total_chars < 40:
        try:
            from pdf2image import convert_from_bytes
            import pytesseract

            images = convert_from_bytes(data)
            ocr_pages = [pytesseract.image_to_string(img) for img in images]
            ocr_used = True
            return ExtractResult("\n\n".join(ocr_pages), ocr_used)
        except Exception:
            # pdf2image needs poppler; if unavailable, return whatever text we had.
            pass

    return ExtractResult("\n\n".join(pages_text), ocr_used)


def _extract_docx(data: bytes) -> ExtractResult:
    from docx import Document

    doc = Document(io.BytesIO(data))
    parts = [p.text for p in doc.paragraphs]
    for table in doc.tables:
        for row in table.rows:
            parts.append(" | ".join(cell.text for cell in row.cells))
    return ExtractResult("\n".join(parts), ocr_used=False)


def extract_text(filename: str, content_type: str, data: bytes) -> ExtractResult:
    """Dispatch extraction based on file type."""
    name = filename.lower()

    if name.endswith(".pdf") or content_type == "application/pdf":
        return _extract_pdf(data)

    if name.endswith(".docx") or content_type in (
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ):
        return _extract_docx(data)

    if name.endswith((".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif")) or (
        content_type or ""
    ).startswith("image/"):
        return ExtractResult(_ocr_image_bytes(data), ocr_used=True)

    if name.endswith((".txt", ".md", ".csv")) or (content_type or "").startswith("text/"):
        return ExtractResult(data.decode("utf-8", errors="ignore"), ocr_used=False)

    # Last resort: try to decode as text.
    return ExtractResult(data.decode("utf-8", errors="ignore"), ocr_used=False)


def _normalize(text: str) -> str:
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def chunk_text(text: str, chunk_size: int, overlap: int) -> List[str]:
    """Split text into overlapping word-based chunks.

    Word-based windows keep chunks coherent and avoid cutting mid-word.
    chunk_size / overlap are measured in characters (approximate).
    """
    text = _normalize(text)
    if not text:
        return []

    words = text.split(" ")
    chunks: List[str] = []
    current: List[str] = []
    current_len = 0

    for word in words:
        current.append(word)
        current_len += len(word) + 1
        if current_len >= chunk_size:
            chunks.append(" ".join(current).strip())
            # Build overlap tail by characters.
            tail: List[str] = []
            tail_len = 0
            for w in reversed(current):
                tail_len += len(w) + 1
                tail.insert(0, w)
                if tail_len >= overlap:
                    break
            current = tail
            current_len = sum(len(w) + 1 for w in current)

    if current and " ".join(current).strip():
        chunks.append(" ".join(current).strip())

    return [c for c in chunks if c]


def process_document(
    filename: str,
    content_type: str,
    data: bytes,
    chunk_size: int,
    overlap: int,
) -> Tuple[List[str], bool, int]:
    """Return (chunks, ocr_used, num_chars)."""
    result = extract_text(filename, content_type, data)
    chunks = chunk_text(result.text, chunk_size, overlap)
    return chunks, result.ocr_used, len(result.text)