shira-hermes/api/services/ocr.py

"""OCR fallback — extract text from images and scanned PDFs via Claude Vision."""

from __future__ import annotations

import base64
import io
import logging
import os

from openai import AsyncOpenAI

logger = logging.getLogger("shira.ocr")

OCR_PROMPT = (
    "Extract ALL visible text from this image. Return ONLY the raw text content — no commentary, "
    "no formatting notes, no headers like 'Here is the text:'. Preserve the original language "
    "(Hebrew, English, Arabic, etc.) and line breaks. If the image contains a form or table, "
    "preserve the structure with tabs or spaces."
)

MAX_PDF_PAGES = 20
PDF_RENDER_DPI = 180


def _build_vision_messages(images_b64: list[tuple[str, str]], prompt: str = OCR_PROMPT) -> list[dict]:
    """Build OpenAI-format messages with vision content. images_b64 is list of (mime, base64)."""
    content = [{"type": "text", "text": prompt}]
    for mime, b64 in images_b64:
        content.append({
            "type": "image_url",
            "image_url": {"url": f"data:{mime};base64,{b64}"},
        })
    return [{"role": "user", "content": content}]


async def _call_vision(images_b64: list[tuple[str, str]], prompt: str = OCR_PROMPT) -> str:
    """Send images to Claude Vision via ai-gateway and return extracted text."""
    client = AsyncOpenAI(
        base_url=os.environ.get("AI_GATEWAY_URL", "http://localhost:3000").rstrip("/") + "/v1",
        api_key=os.environ.get("AI_GATEWAY_API_KEY", ""),
    )
    response = await client.chat.completions.create(
        model=os.environ.get("CLAUDE_MODEL", "sonnet"),
        messages=_build_vision_messages(images_b64, prompt),
        max_tokens=4096,
    )
    return (response.choices[0].message.content or "").strip()


async def ocr_image(base64_data: str, mime_type: str) -> str:
    """OCR a single image via Claude Vision."""
    logger.info("[ocr] image mime=%s size=%d", mime_type, len(base64_data))
    return await _call_vision([(mime_type, base64_data)])


async def ocr_pdf(base64_data: str) -> str:
    """Render PDF pages as images and OCR each via Claude Vision. Returns combined text."""
    try:
        import fitz  # PyMuPDF
    except ImportError:
        return "❌ PyMuPDF לא מותקן — לא ניתן לסרוק PDF כתמונה."

    pdf_bytes = base64.b64decode(base64_data)
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    page_count = min(len(doc), MAX_PDF_PAGES)
    logger.info("[ocr] pdf pages=%d (limit=%d)", len(doc), MAX_PDF_PAGES)

    pages_text = []
    for i in range(page_count):
        page = doc[i]
        pix = page.get_pixmap(dpi=PDF_RENDER_DPI)
        img_bytes = pix.tobytes("png")
        img_b64 = base64.b64encode(img_bytes).decode("ascii")
        try:
            text = await _call_vision([("image/png", img_b64)])
            pages_text.append(f"=== עמוד {i + 1} ===\n{text}")
        except Exception as e:
            logger.error("[ocr] page %d failed: %s", i + 1, e)
            pages_text.append(f"=== עמוד {i + 1} ===\n[שגיאה: {e}]")

    doc.close()
    if len(doc) > MAX_PDF_PAGES:
        pages_text.append(f"\n[הוגבל ל-{MAX_PDF_PAGES} עמודים מתוך {len(doc)}]")
    return "\n\n".join(pages_text)


def _extract_docx_xml_text(docx_zip) -> str:
    """Parse word/document.xml from a DOCX and return plain text.

    The DOCX text payload is a flat sequence of <w:t>...</w:t> runs grouped
    inside <w:p> paragraphs. We turn each paragraph into one text line.
    """
    import re
    try:
        xml = docx_zip.read("word/document.xml").decode("utf-8", errors="replace")
    except KeyError:
        return ""

    # Anchor the run-open to either `<w:t>` (no attrs) or `<w:t ...>` (with attrs).
    # The `(?:\s[^>]*)?` keeps us from also matching `<w:tab>` / `<w:tbl>` etc.
    run_re = re.compile(r"<w:t(?:\s[^>]*)?>(.*?)</w:t>", flags=re.DOTALL)
    para_re = re.compile(r"<w:p(?:\s[^>]*)?>.*?</w:p>", flags=re.DOTALL)

    paragraphs = []
    for p_match in para_re.finditer(xml):
        block = p_match.group(0)
        line = "".join(run_re.findall(block)).strip()
        if line:
            paragraphs.append(line)
    return "\n".join(paragraphs).strip()


async def ocr_docx_images(base64_data: str) -> str:
    """Extract text + images from a DOCX archive.

    Order of operations:
    1. Read word/document.xml for the actual text content (the LLM needs this).
    2. OCR any embedded images via Claude Vision (signatures, scanned inserts).

    Previously this only OCR'd images, so a text-heavy DOCX with one signature
    image returned only `[signature]` — and downstream LLMs hallucinated the rest.
    """
    import zipfile

    docx_bytes = base64.b64decode(base64_data)
    try:
        z = zipfile.ZipFile(io.BytesIO(docx_bytes))
    except zipfile.BadZipFile:
        return "❌ הקובץ אינו DOCX תקין."

    sections = []

    # 1. Text from document.xml
    xml_text = _extract_docx_xml_text(z)
    logger.info("[ocr] docx xml text chars=%d", len(xml_text))
    if xml_text:
        sections.append(f"=== טקסט מהמסמך ({len(xml_text)} תווים) ===\n{xml_text}")

    # 2. OCR each embedded image
    media = [n for n in z.namelist() if n.startswith("word/media/")]
    logger.info("[ocr] docx images found=%d", len(media))
    for i, name in enumerate(media[:MAX_PDF_PAGES]):
        ext = name.rsplit(".", 1)[-1].lower()
        mime = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "gif": "image/gif"}.get(ext, "image/png")
        img_bytes = z.read(name)
        img_b64 = base64.b64encode(img_bytes).decode("ascii")
        try:
            text = await _call_vision([(mime, img_b64)])
            sections.append(f"=== תמונה {i + 1} ({name}) ===\n{text}")
        except Exception as e:
            logger.error("[ocr] docx image %s failed: %s", name, e)

    if not sections:
        return "❌ לא נמצא טקסט וגם לא תמונות ב-DOCX."
    return "\n\n".join(sections)


async def extract_text_from_bytes(base64_data: str, mime_type: str, file_name: str = "") -> str:
    """Route to the right OCR method based on mime type. Returns extracted text."""
    mt = (mime_type or "").lower()

    if mt.startswith("image/"):
        return await ocr_image(base64_data, mt)

    if mt == "application/pdf":
        return await ocr_pdf(base64_data)

    if mt in (
        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "application/msword",
    ) or file_name.lower().endswith((".docx", ".doc")):
        return await ocr_docx_images(base64_data)

    return f"❌ סוג קובץ לא נתמך ל-OCR: {mime_type}"