engine/vlm_extractor.py

import base64
import json
import logging
import os
import re

logger = logging.getLogger(__name__)


def _get_hf_token() -> str:
    """Get HF API token from env var or HF Spaces secrets file."""
    token = os.environ.get("HF_API_TOKEN", "")
    if token:
        return token
    # HF Spaces Docker: secrets may be mounted at /run/secrets/<NAME>
    secrets_path = "/run/secrets/HF_API_TOKEN"
    if os.path.exists(secrets_path):
        try:
            with open(secrets_path) as f:
                token = f.read().strip()
            if token:
                logger.info("Found HF_API_TOKEN at %s", secrets_path)
                os.environ["HF_API_TOKEN"] = token
                return token
        except Exception as e:
            logger.warning("Failed to read %s: %s", secrets_path, e)
    return ""


# VLM prompts per document type
_PROMPTS = {
    "sa_id_card": (
        "This is a South African ID smart card (front side). "
        "Extract these fields as JSON: "
        '{"id_number": "13-digit SA ID number (YYMMDDSSSSCCAZ)", '
        '"surname": "...", "names": "...", "date_of_birth": "YYYY-MM-DD", '
        '"sex": "Male or Female", "nationality": "...", '
        '"country_of_birth": "...", "citizenship_status": "SA Citizen or Permanent Resident"}. '
        "The ID number is 13 digits: YYMMDD=DOB, SSSS=gender sequence "
        "(0000-4999=Female, 5000-9999=Male), C=citizenship (0=citizen, 1=resident), "
        "A=usually 8, Z=check digit. Return ONLY valid JSON, no explanation."
    ),
    "sa_id_book": (
        "This is a South African green ID book (paper format, front page). "
        "It has a green security background with printed text. "
        "Extract these fields as JSON: "
        '{"id_number": "13-digit SA ID number", '
        '"surname": "...", "names": "...", "date_of_birth": "YYYY-MM-DD", '
        '"sex": "Male or Female", "nationality": "...", '
        '"country_of_birth": "...", "citizenship_status": "SA Citizen or Permanent Resident"}. '
        "Labels may be in English or Afrikaans (e.g. VAN/SURNAME, VOORNAME/NAMES). "
        "Return ONLY valid JSON, no explanation."
    ),
    "passport": (
        "This is a passport document. Extract these fields as JSON: "
        '{"passport_number": "...", "surname": "...", "given_names": "...", '
        '"date_of_birth": "YYYY-MM-DD", "sex": "Male or Female", '
        '"nationality": "...", "expiry_date": "YYYY-MM-DD", '
        '"issuing_country": "..."}. '
        "If this is a South African passport, also extract: "
        '"id_number": "13-digit SA ID number if visible". '
        "Return ONLY valid JSON, no explanation."
    ),
}

# Default model — must be available on HF Inference API
_DEFAULT_MODEL = "Qwen/Qwen3-VL-8B-Instruct"

# Inference provider — pin explicitly so HF auto-routing can't switch us to a
# provider that has dropped serverless support for the model (e.g. Together,
# which now requires a dedicated endpoint for Qwen3-VL-8B-Instruct).
_VLM_PROVIDER = os.environ.get("OCR_VLM_PROVIDER", "novita")

# Timeout for VLM calls (seconds)
_VLM_TIMEOUT = 90

# Max long edge for images sent to VLM (pixels)
_VLM_MAX_RESOLUTION = 1500


def _prepare_image(image_path: str) -> tuple[str, str]:
    """Read, optionally downscale, and base64-encode an image for VLM.

    Returns (base64_data, mime_type).
    """
    import cv2

    img = cv2.imread(image_path)
    if img is None:
        # Fall back to raw file read
        with open(image_path, "rb") as f:
            data = base64.b64encode(f.read()).decode("utf-8")
        ext = image_path.rsplit(".", 1)[-1].lower()
        mime_map = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png",
                    "bmp": "image/bmp", "webp": "image/webp"}
        return data, mime_map.get(ext, "image/jpeg")

    h, w = img.shape[:2]
    long_edge = max(h, w)

    if long_edge > _VLM_MAX_RESOLUTION:
        scale = _VLM_MAX_RESOLUTION / long_edge
        new_w, new_h = int(w * scale), int(h * scale)
        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        logger.info("Downscaled image from %dx%d to %dx%d for VLM", w, h, new_w, new_h)

    # Encode as JPEG (smaller than PNG)
    _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
    data = base64.b64encode(buf.tobytes()).decode("utf-8")
    return data, "image/jpeg"


def extract_fields_vlm(image_path: str, doc_type: str) -> dict | None:
    """Extract document fields using a VLM via HF Inference API.

    Args:
        image_path: Path to the document image.
        doc_type: One of 'sa_id_card', 'sa_id_book', 'passport'.

    Returns:
        Dict with extracted fields, or None if VLM fails.
    """
    api_token = _get_hf_token()
    if not api_token:
        logger.info("HF_API_TOKEN not set, skipping VLM extraction")
        return None

    model = os.environ.get("OCR_VLM_MODEL", _DEFAULT_MODEL)
    prompt = _PROMPTS.get(doc_type)
    if not prompt:
        logger.warning("No VLM prompt for doc_type: %s", doc_type)
        return None

    try:
        from huggingface_hub import InferenceClient

        # Downscale large images to reduce base64 payload and speed up VLM
        image_data, mime_type = _prepare_image(image_path)

        logger.info("Calling VLM (%s via %s) for doc_type=%s", model, _VLM_PROVIDER, doc_type)
        client = InferenceClient(provider=_VLM_PROVIDER, token=api_token, timeout=_VLM_TIMEOUT)
        response = client.chat.completions.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{image_data}"
                            },
                        },
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                }
            ],
            max_tokens=512,
        )

        content = response.choices[0].message.content

        # Parse JSON from response (may be wrapped in ```json ... ```)
        result = _parse_vlm_response(content)
        if result:
            result["source"] = "vlm"
            logger.info("VLM extraction successful: %d fields", len(result) - 1)
        return result

    except Exception as e:
        logger.warning("VLM extraction failed: %s", e)
        return None


def _parse_vlm_response(content: str) -> dict | None:
    """Extract JSON dict from VLM response text.

    Handles responses wrapped in ```json ... ```, plain JSON,
    and Qwen3 thinking mode (<think>...</think> prefix).
    """
    # Strip Qwen3 thinking tags — JSON is always after </think>
    think_end = content.find("</think>")
    if think_end != -1:
        content = content[think_end + len("</think>"):].strip()

    # Try to extract JSON from markdown code block
    json_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", content, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group(1))
        except json.JSONDecodeError:
            pass

    # Try to find raw JSON object
    json_match = re.search(r"\{[^{}]*\}", content, re.DOTALL)
    if json_match:
        try:
            return json.loads(json_match.group(0))
        except json.JSONDecodeError:
            pass

    # Try the entire content as JSON
    try:
        parsed = json.loads(content.strip())
        if isinstance(parsed, dict):
            return parsed
    except json.JSONDecodeError:
        pass

    logger.warning("Could not parse JSON from VLM response: %.200s", content)
    return None


def warm_vlm() -> bool:
    """Check if VLM is available by making a simple text request.

    Returns True if the model is responsive.
    """
    api_token = _get_hf_token()
    if not api_token:
        return False

    model = os.environ.get("OCR_VLM_MODEL", _DEFAULT_MODEL)

    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(provider=_VLM_PROVIDER, token=api_token, timeout=10)
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": "Reply with OK"}],
            max_tokens=5,
        )
        ok = bool(response.choices[0].message.content)
        logger.info("VLM warm-up: model=%s provider=%s loaded=%s", model, _VLM_PROVIDER, ok)
        return ok
    except Exception as e:
        logger.warning("VLM warm-up failed: %s", e)
        return False