anonymisation/vlm_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VLM Manager — Analyse visuelle des PDF via Ollama (qwen3-vl)
-------------------------------------------------------------
Couche complémentaire aux regex/NER : envoie chaque page PDF comme image
à un VLM local (Ollama) pour détecter visuellement les PII.

Dégradation gracieuse : si Ollama est indisponible, le pipeline continue sans VLM.

Dépendances : aucune (utilise uniquement urllib de la stdlib).
"""
from __future__ import annotations

import base64
import io
import json
import logging
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

try:
    from PIL import Image
except ImportError:
    Image = None  # type: ignore

log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass
class VlmConfig:
    """Configuration pour le VLM Ollama."""
    base_url: str = "http://localhost:11434"
    model: str = "qwen2.5vl:7b"
    timeout: int = 180
    max_image_size: int = 2048   # pixels (côté le plus long)
    temperature: float = 0.1
    num_predict: int = 8192
    min_confidence: float = 0.5


# ---------------------------------------------------------------------------
# Mapping catégories VLM → (PiiHit.kind, clé PLACEHOLDERS)
# ---------------------------------------------------------------------------

VLM_CATEGORY_MAP: Dict[str, Tuple[str, str]] = {
    "NOM":              ("VLM_NOM",         "NOM"),
    "PRENOM":           ("VLM_NOM",         "NOM"),
    "ADRESSE":          ("VLM_ADRESSE",     "ADRESSE"),
    "TELEPHONE":        ("VLM_TEL",         "TEL"),
    "EMAIL":            ("VLM_EMAIL",       "EMAIL"),
    "DATE_NAISSANCE":   ("VLM_DATE_NAISS",  "DATE_NAISSANCE"),
    "NIR":              ("VLM_NIR",         "NIR"),
    "IPP":              ("VLM_IPP",         "IPP"),
    "CODE_POSTAL":      ("VLM_CP",          "CODE_POSTAL"),
    "VILLE":            ("VLM_VILLE",       "VILLE"),
    "RPPS":             ("VLM_RPPS",        "RPPS"),
    # Identifiants médicaux / traçabilité
    "NUMERO_PATIENT":   ("VLM_NUM_PATIENT", "DOSSIER"),
    "NUMERO_LOT":       ("VLM_NUM_LOT",     "MASK"),
    "NUMERO_ORDONNANCE":("VLM_NUM_ORD",     "DOSSIER"),
    "NUMERO_SEJOUR":    ("VLM_NDA",         "NDA"),
    "NDA":              ("VLM_NDA",         "NDA"),
    "SERVICE":          ("VLM_SERVICE",     "MASK"),
    "ETABLISSEMENT":    ("VLM_ETAB",        "ETAB"),
    "DATE":             ("VLM_DATE",        "DATE"),
    "AGE":              ("VLM_AGE",         "AGE"),
}


# ---------------------------------------------------------------------------
# Prompt système
# ---------------------------------------------------------------------------

_SYSTEM_PROMPT = (
    "Tu identifies les données personnelles et identifiants traçables dans les documents "
    "médicaux français. Réponds uniquement en JSON."
)

_USER_PROMPT_TEMPLATE = """\
Identifie TOUTES les informations permettant d'identifier un patient dans cette page de document médical.
Le document peut être pivoté ou manuscrit — lis dans toutes les orientations.

Catégories :
- NOM, PRENOM : noms et prénoms de patients, médecins, infirmiers, soignants
- ADRESSE : adresses postales
- TELEPHONE : numéros de téléphone
- DATE_NAISSANCE : dates de naissance
- DATE : toutes les autres dates (consultation, séjour, transfusion, intervention…)
- NIR : numéro de sécurité sociale
- IPP : identifiant permanent du patient (ex: BA172948)
- NDA : numéro de dossier administratif / numéro de séjour
- NUMERO_PATIENT : tout numéro identifiant un patient (numéro EFS, numéro d'ordonnance…)
- NUMERO_LOT : numéros de lots de produits sanguins (PSL), codes numériques de traçabilité
- CODE_POSTAL, VILLE : codes postaux et villes
- ETABLISSEMENT : noms d'hôpitaux, cliniques (ex: CH COTE BASQUE, CHU BORDEAUX)
- SERVICE : noms de services hospitaliers (ex: CANCEROLOGIE HDJ, REANIMATION)
- AGE : âge du patient (ex: 85A, 62 ans)
- RPPS : numéro RPPS du médecin

Règles :
- Texte EXACT visible sur l'image (copie fidèle, y compris manuscrit)
- Inclure TOUS les identifiants numériques (manuscrits ou imprimés)
- Inclure les noms d'établissements et de services hospitaliers
- Ne PAS inclure : médicaments, diagnostics, termes médicaux purs, résultats de labo

Réponds en JSON : {{"entites": [{{"categorie": "NOM", "texte": "DUPONT", "confiance": 0.95}}]}}
Si aucune PII : {{"entites": []}}"""


# ---------------------------------------------------------------------------
# VlmManager
# ---------------------------------------------------------------------------

class VlmManager:
    """Gestionnaire VLM via Ollama. Même pattern que NerModelManager."""

    def __init__(self, config: Optional[VlmConfig] = None):
        self._config = config or VlmConfig()
        self._loaded = False
        self._model_name: Optional[str] = None

    # ---- public API ----

    def is_loaded(self) -> bool:
        return self._loaded

    def load(self, model: Optional[str] = None) -> None:
        """Vérifie la connexion Ollama et la disponibilité du modèle."""
        cfg = self._config
        if model:
            cfg.model = model
        self._model_name = cfg.model

        # 1) Vérifier qu'Ollama répond
        try:
            req = urllib.request.Request(
                f"{cfg.base_url}/api/tags",
                method="GET",
            )
            with urllib.request.urlopen(req, timeout=10) as resp:
                data = json.loads(resp.read().decode("utf-8"))
        except Exception as e:
            raise RuntimeError(f"Ollama indisponible ({cfg.base_url}) : {e}") from e

        # 2) Vérifier que le modèle est disponible
        available = [m.get("name", "") for m in data.get("models", [])]
        # Normaliser : "qwen3-vl:8b" matche "qwen3-vl:8b" ou "qwen3-vl:8b-..."
        model_found = any(
            a == cfg.model or a.startswith(cfg.model.split(":")[0] + ":")
            for a in available
        )
        if not model_found:
            raise RuntimeError(
                f"Modèle '{cfg.model}' non trouvé dans Ollama. "
                f"Disponibles : {', '.join(available) or '(aucun)'}. "
                f"Lancez : ollama pull {cfg.model}"
            )

        self._loaded = True
        log.info("VLM prêt : %s via %s", cfg.model, cfg.base_url)

    def unload(self) -> None:
        self._loaded = False
        self._model_name = None

    @staticmethod
    def models_catalog() -> Dict[str, str]:
        return {
            "Qwen2.5-VL 7B (Ollama)": "qwen2.5vl:7b",
            "Qwen3-VL 8B (Ollama)": "qwen3-vl:8b",
        }

    # ---- analyse d'une page ----

    def analyze_page_image(
        self,
        image: "Image.Image",
        page_number: int = 0,
        existing_pii: Optional[List[str]] = None,
    ) -> List[Dict[str, Any]]:
        """Envoie une image de page à Ollama et retourne les entités détectées.

        Returns:
            Liste de dicts avec clés : categorie, texte, confiance
        """
        if not self._loaded:
            return []
        if Image is None:
            log.warning("Pillow non disponible, VLM ignoré")
            return []

        cfg = self._config

        # Redimensionner l'image
        img = _resize_image(image, cfg.max_image_size)

        # Encoder en base64
        img_b64 = _image_to_base64(img)

        # Construire le prompt utilisateur
        user_prompt = _USER_PROMPT_TEMPLATE
        if existing_pii:
            user_prompt += (
                "\n\nPII déjà détectés (vérifie et cherche ceux qui manquent) : "
                + ", ".join(existing_pii[:20])
            )

        # Appel API Ollama — format: json force une sortie JSON valide
        payload = {
            "model": cfg.model,
            "messages": [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": user_prompt,
                    "images": [img_b64],
                },
            ],
            "format": "json",
            "stream": False,
            "options": {
                "temperature": cfg.temperature,
                "num_predict": cfg.num_predict,
            },
        }

        try:
            body = json.dumps(payload).encode("utf-8")
            req = urllib.request.Request(
                f"{cfg.base_url}/api/chat",
                data=body,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=cfg.timeout) as resp:
                result = json.loads(resp.read().decode("utf-8"))
        except urllib.error.URLError as e:
            log.warning("VLM appel échoué (page %d) : %s", page_number, e)
            return []
        except Exception as e:
            log.warning("VLM erreur inattendue (page %d) : %s", page_number, e)
            return []

        # Extraire le contenu de la réponse
        # Qwen3 peut mettre la réponse dans "content" ou "thinking"
        content = ""
        msg = result.get("message", {})
        if isinstance(msg, dict):
            content = msg.get("content", "")
            # Fallback : si content vide, chercher dans thinking (mode Qwen3)
            if not content.strip():
                content = msg.get("thinking", "")
        elif isinstance(msg, str):
            content = msg

        # Parser le JSON de réponse (défensif)
        entities = _parse_vlm_response(content)

        log.info("VLM page %d : %d entités détectées", page_number, len(entities))
        return entities


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _resize_image(img: "Image.Image", max_size: int) -> "Image.Image":
    """Redimensionne l'image si un côté dépasse max_size, en conservant le ratio."""
    w, h = img.size
    if max(w, h) <= max_size:
        return img
    ratio = max_size / max(w, h)
    new_w = int(w * ratio)
    new_h = int(h * ratio)
    return img.resize((new_w, new_h), Image.LANCZOS)


def _image_to_base64(img: "Image.Image") -> str:
    """Encode une image PIL en base64 (PNG)."""
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("ascii")


def _parse_vlm_response(content: str) -> List[Dict[str, Any]]:
    """Parse la réponse du VLM en liste d'entités. Gère JSON brut, markdown code blocks,
    et JSON noyé dans du texte de raisonnement (thinking)."""
    if not content or not content.strip():
        return []

    import re
    text = content.strip()

    # Tentative 1 : JSON direct
    try:
        data = json.loads(text)
        return _extract_entities(data)
    except json.JSONDecodeError:
        pass

    # Tentative 2 : extraire un bloc ```json ... ``` ou ``` ... ```
    m = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
    if m:
        try:
            data = json.loads(m.group(1).strip())
            return _extract_entities(data)
        except json.JSONDecodeError:
            pass

    # Tentative 3 : chercher un bloc JSON contenant "entites" ou "entities"
    # Gérer les accolades imbriquées en trouvant le bon bloc
    for keyword in ['"entites"', '"entities"']:
        idx = text.find(keyword)
        if idx < 0:
            continue
        # Remonter jusqu'au { ouvrant
        brace_start = text.rfind("{", 0, idx)
        if brace_start < 0:
            continue
        # Trouver le } fermant correspondant (gestion profondeur)
        depth = 0
        for i in range(brace_start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(text[brace_start:i + 1])
                        return _extract_entities(data)
                    except json.JSONDecodeError:
                        break
        break

    # Tentative 4 : chercher le premier { ... } (fallback)
    brace_start = text.find("{")
    brace_end = text.rfind("}")
    if brace_start >= 0 and brace_end > brace_start:
        try:
            data = json.loads(text[brace_start:brace_end + 1])
            return _extract_entities(data)
        except json.JSONDecodeError:
            pass

    # Tentative 5 : réparation JSON tronqué (num_predict dépassé)
    # Le VLM a pu couper la réponse au milieu d'un objet entité
    brace_start = text.find("{")
    if brace_start >= 0:
        fragment = text[brace_start:]
        # Trouver la dernière entité complète (se terminant par })
        last_complete = fragment.rfind("}")
        if last_complete > 0:
            truncated = fragment[:last_complete + 1]
            # Fermer le tableau et l'objet si nécessaire
            for suffix in ["", "]}", "]}}"]:
                try:
                    data = json.loads(truncated + suffix)
                    entities = _extract_entities(data)
                    if entities:
                        log.info("VLM : JSON tronqué réparé (%d entités récupérées)", len(entities))
                        return entities
                except json.JSONDecodeError:
                    continue

    log.warning("VLM : impossible de parser la réponse JSON : %s", text[:200])
    return []


def _extract_entities(data: Any) -> List[Dict[str, Any]]:
    """Extrait la liste d'entités depuis la structure JSON parsée."""
    raw_list = []
    if isinstance(data, dict):
        # Structure attendue : {"entites": [...]}
        raw_list = data.get("entites") or data.get("entities") or []
        if not isinstance(raw_list, list):
            raw_list = []
    elif isinstance(data, list):
        raw_list = data

    result = []
    for e in raw_list:
        if not isinstance(e, dict):
            continue
        texte = e.get("texte") or e.get("text") or ""
        if not texte:
            continue
        # Accepter les entités sans catégorie (default NOM)
        categorie = e.get("categorie") or e.get("category") or "NOM"
        result.append({
            "categorie": categorie.upper(),
            "texte": texte,
            "confiance": float(e.get("confiance", e.get("confidence", 0.8))),
        })
    return result