Orchestrator/privat/NoteGraph/ingestion/extraction/pdf.py

"""PDF text extraction using pdfplumber.

Detects scanned pages (no extractable text) and delegates to:
1. Mistral Vision API (preferred — handles handwriting well)
2. Tesseract OCR (fallback if no LLM configured)
"""

import base64
import io
import logging
from pathlib import Path

import pdfplumber

from .base import ExtractionResult

logger = logging.getLogger(__name__)


class PDFExtractor:
    """Extracts text from PDF files, with Vision LLM or OCR fallback for scanned pages."""

    def __init__(
        self,
        ocr_enabled: bool = True,
        ocr_language: str = "deu+eng",
        llm_provider: str = "mistral",
        llm_model: str = "mistral-small-latest",
        mistral_api_key: str = "",
    ):
        self.ocr_enabled = ocr_enabled
        self.ocr_language = ocr_language
        self.llm_provider = llm_provider
        self.llm_model = llm_model
        self.mistral_api_key = mistral_api_key
        self._use_vision = bool(mistral_api_key)

    def can_handle(self, file_path: Path) -> bool:
        """Return True for .pdf files."""
        return file_path.suffix.lower() == ".pdf"

    def extract(self, file_path: Path) -> ExtractionResult:
        """Extract text from PDF, using Vision LLM or OCR for scanned pages."""
        pages_text: list[str] = []
        has_direct_text = False
        has_vision_text = False

        with pdfplumber.open(file_path) as pdf:
            for page_num, page in enumerate(pdf.pages, start=1):
                text = page.extract_text() or ""
                text = text.strip()

                if text:
                    pages_text.append(text)
                    has_direct_text = True
                elif self.ocr_enabled:
                    # Scanned page — try Vision LLM first, then Tesseract
                    logger.debug(
                        "Page %d of %s has no text, attempting vision OCR",
                        page_num, file_path,
                    )
                    ocr_text = self._vision_ocr_page(page, page_num, file_path)
                    if ocr_text:
                        pages_text.append(ocr_text)
                        has_vision_text = True

        # Determine extraction method
        if has_direct_text and has_vision_text:
            method = "mixed"
        elif has_vision_text:
            method = "vision-ocr"
        else:
            method = "direct"

        return ExtractionResult(
            text="\n\n".join(pages_text),
            source_file=file_path,
            extraction_method=method,
            metadata={"page_count": len(pages_text)},
        )

    def _vision_ocr_page(self, page, page_num: int, file_path: Path) -> str:
        """OCR a PDF page using Mistral Vision API, with Tesseract fallback."""
        if self._use_vision:
            try:
                return self._mistral_vision_ocr(page, page_num)
            except Exception as e:
                logger.warning(
                    "Vision OCR failed for page %d of %s: %s, trying Tesseract",
                    page_num, file_path, e,
                )

        # Fallback to Tesseract
        return self._tesseract_ocr(page, page_num, file_path)

    def _mistral_vision_ocr(self, page, page_num: int) -> str:
        """Send page image to Mistral Vision for text extraction."""
        import httpx

        # Convert page to PNG image
        image = page.to_image(resolution=200)
        img_buffer = io.BytesIO()
        image.original.save(img_buffer, format="PNG")
        img_bytes = img_buffer.getvalue()
        img_b64 = base64.b64encode(img_bytes).decode("utf-8")

        # Call Mistral Vision API via LiteLLM-compatible format
        response = httpx.post(
            "https://api.mistral.ai/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {self.mistral_api_key}",
                "Content-Type": "application/json",
            },
            json={
                "model": self.llm_model,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "Extract ALL text from this image. This is a handwritten note or document. "
                                        "Transcribe everything you can read, preserving structure (paragraphs, lists, headings). "
                                        "If text is in German, keep it in German. Return ONLY the extracted text, no commentary.",
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{img_b64}",
                                },
                            },
                        ],
                    }
                ],
                "max_tokens": 4000,
            },
            timeout=60.0,
        )
        response.raise_for_status()
        data = response.json()
        text = data["choices"][0]["message"]["content"]
        logger.info("Vision OCR page %d: extracted %d chars", page_num, len(text))
        return text.strip()

    def _tesseract_ocr(self, page, page_num: int, file_path: Path) -> str:
        """Fallback: OCR a PDF page using Tesseract."""
        try:
            import pytesseract
            image = page.to_image(resolution=300)
            text = pytesseract.image_to_string(image.original, lang=self.ocr_language)
            return text.strip()
        except Exception as e:
            logger.warning(
                "Tesseract OCR failed for page %d of %s: %s", page_num, file_path, e
            )
            return ""