rpa_vision_v3/core/detection/ollama_client.py

"""
OllamaClient - Client pour Vision-Language Models via Ollama

Interface pour communiquer avec des VLM (Qwen, LLaVA, etc.) via Ollama.
"""

import logging
from typing import Dict, List, Optional, Any
import requests
import json
import base64
from pathlib import Path
from PIL import Image
import io

logger = logging.getLogger(__name__)


class OllamaClient:
    """
    Client Ollama pour VLM
    
    Permet d'envoyer des images et prompts à un VLM via l'API Ollama.
    """
    
    def __init__(self, 
                 endpoint: str = "http://localhost:11434",
                 model: str = "qwen3-vl:8b",
                 timeout: int = 180):
        """
        Initialiser le client Ollama
        
        Args:
            endpoint: URL de l'API Ollama
            model: Nom du modèle VLM à utiliser
            timeout: Timeout en secondes
        """
        self.endpoint = endpoint.rstrip('/')
        self.model = model
        self.timeout = timeout
        self._check_connection()
    
    def _check_connection(self) -> bool:
        """Vérifier la connexion à Ollama"""
        try:
            response = requests.get(f"{self.endpoint}/api/tags", timeout=5)
            if response.status_code == 200:
                models = response.json().get('models', [])
                model_names = [m['name'] for m in models]
                if self.model not in model_names:
                    logger.warning(f" Model '{self.model}' not found in Ollama")
                    logger.info(f"Available models: {model_names}")
                return True
        except Exception as e:
            logger.warning(f" Cannot connect to Ollama at {self.endpoint}: {e}")
            return False
        return False
    
    def generate(self,
                prompt: str,
                image_path: Optional[str] = None,
                image: Optional[Image.Image] = None,
                system_prompt: Optional[str] = None,
                temperature: float = 0.1,
                max_tokens: int = 500,
                force_json: bool = False,
                assistant_prefill: Optional[str] = None,
                num_ctx: Optional[int] = None,
                extra_images_b64: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        Générer une réponse du VLM via l'API chat d'Ollama.

        Pour les modèles thinking (qwen3-vl), on utilise la technique du
        "assistant prefill" : un message assistant pré-rempli est ajouté
        après le message user, forçant le modèle à continuer directement
        sans phase de thinking. Cela résout le bug Ollama 0.18.x où
        think=false est ignoré par le renderer qwen3-vl-thinking.

        Sans prefill : le modèle pense 500+ tokens puis répond (~180s)
        Avec prefill : le modèle répond directement (~1-5s)

        Args:
            prompt: Prompt textuel
            image_path: Chemin vers une image (optionnel)
            image: Image PIL (optionnel)
            system_prompt: Prompt système (optionnel)
            temperature: Température de génération
            max_tokens: Nombre max de tokens
            force_json: Forcer la sortie JSON (non recommandé pour qwen3-vl)
            assistant_prefill: Début de réponse pré-rempli (auto-détecté si None)
            num_ctx: Context window (défaut 2048, augmenter pour batch)
            extra_images_b64: Images supplémentaires en base64 à envoyer avec le prompt.
                Ajoutées après l'image principale. Utile pour le VLM multi-image
                (ex: screenshot + crop de référence).

        Returns:
            Dict avec 'response', 'success', 'error'
        """
        try:
            # Préparer l'image si fournie
            image_data = None
            if image_path:
                image_data = self._encode_image_from_path(image_path)
            elif image:
                image_data = self._encode_image_from_pil(image)

            # Nettoyer le prompt — retirer /no_think et /nothink du texte
            effective_prompt = prompt.replace("/no_think\n", "").replace("/no_think", "")
            effective_prompt = effective_prompt.replace("/nothink ", "").replace("/nothink", "")
            effective_prompt = effective_prompt.strip()

            # Construire le message utilisateur
            user_message = {"role": "user", "content": effective_prompt}
            all_images = []
            if image_data:
                all_images.append(image_data)
            if extra_images_b64:
                all_images.extend(extra_images_b64)
            if all_images:
                user_message["images"] = all_images

            # Construire les messages
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            messages.append(user_message)

            # Déterminer si le modèle est un modèle thinking (qwen3)
            is_thinking_model = "qwen3" in self.model.lower()

            # WORKAROUND Ollama 0.18.x : think=false est ignoré par le
            # renderer qwen3-vl-thinking. On utilise un assistant prefill
            # pour forcer le modèle à skip le thinking et répondre directement.
            # Le prefill est choisi en fonction du format attendu.
            # IMPORTANT : avec image, sans prefill le thinking dépasse 180s.
            prefill_used = None
            if is_thinking_model:
                if assistant_prefill is not None:
                    prefill_used = assistant_prefill
                elif force_json:
                    prefill_used = "{"
                elif all_images:
                    # Avec image(s), le thinking est catastrophique (>180s).
                    # Prefill générique pour forcer une réponse directe.
                    prefill_used = "Based on the image,"

                if prefill_used is not None:
                    messages.append({
                        "role": "assistant",
                        "content": prefill_used
                    })

            # num_ctx par défaut à 2048 (correspondant au default du modèle
            # chargé en mémoire). Changer num_ctx force un rechargement du
            # KV cache (~30s de pénalité), donc ne l'augmenter que pour les
            # requêtes batch qui dépassent la limite (image + prompt long).
            effective_num_ctx = num_ctx or 2048

            payload = {
                "model": self.model,
                "messages": messages,
                "stream": False,
                "options": {
                    "temperature": temperature,
                    "num_predict": max_tokens,
                    "num_ctx": effective_num_ctx,
                    "top_k": 1
                }
            }

            # Garder think=false au cas où une future version d'Ollama le
            # corrige — le prefill reste le mécanisme principal
            if is_thinking_model:
                payload["think"] = False

            if force_json:
                payload["format"] = "json"

            # Envoyer la requête via l'API chat
            response = requests.post(
                f"{self.endpoint}/api/chat",
                json=payload,
                timeout=self.timeout
            )

            if response.status_code == 200:
                result = response.json()
                content = result.get("message", {}).get("content", "")

                # Reconstituer la réponse complète en ajoutant le prefill
                if prefill_used and content:
                    content = prefill_used + content

                return {
                    "response": content,
                    "success": True,
                    "error": None
                }
            else:
                return {
                    "response": "",
                    "success": False,
                    "error": f"HTTP {response.status_code}: {response.text}"
                }

        except Exception as e:
            return {
                "response": "",
                "success": False,
                "error": str(e)
            }
    
    def detect_ui_elements(self, image_path: str) -> Dict[str, Any]:
        """
        Détecter les éléments UI dans une image
        
        Args:
            image_path: Chemin vers le screenshot
        
        Returns:
            Dict avec liste d'éléments détectés
        """
        prompt = """Analyze this screenshot and list all interactive UI elements you can see.
For each element, provide:
- Type (button, text_input, checkbox, radio, dropdown, tab, link, icon, table_row, menu_item)
- Position (approximate x, y coordinates)
- Label or text content
- Semantic role (primary_action, cancel, submit, form_input, search_field, navigation, settings, close)

Format your response as JSON."""

        result = self.generate(
            prompt, image_path=image_path, temperature=0.1,
            assistant_prefill="[",
        )
        
        if result["success"]:
            try:
                # Parser la réponse JSON
                elements = json.loads(result["response"])
                return {"elements": elements, "success": True}
            except json.JSONDecodeError:
                # Si pas JSON valide, retourner texte brut
                return {"elements": [], "success": False, "raw_response": result["response"]}
        
        return {"elements": [], "success": False, "error": result["error"]}
    
    def classify_element_type(self, 
                             element_image: Image.Image,
                             context: Optional[str] = None) -> Dict[str, Any]:
        """
        Classifier le type d'un élément UI
        
        Args:
            element_image: Image de l'élément
            context: Contexte additionnel
        
        Returns:
            Dict avec 'type' et 'confidence'
        """
        types_list = "button, text_input, checkbox, radio, dropdown, tab, link, icon, table_row, menu_item"
        
        prompt = f"""What type of UI element is this?
Choose ONLY ONE from: {types_list}

Respond with just the type name, nothing else."""

        if context:
            prompt += f"\n\nContext: {context}"

        result = self.generate(
            prompt, image=element_image, temperature=0.1,
            assistant_prefill="The type is:",
        )
        
        if result["success"]:
            # Retirer le prefill du début pour extraire le type
            raw = result["response"]
            if raw.startswith("The type is:"):
                raw = raw[len("The type is:"):]
            element_type = raw.strip().lower()
            # Valider que c'est un type connu
            valid_types = types_list.split(", ")
            if element_type in valid_types:
                return {"type": element_type, "confidence": 0.9, "success": True}
            else:
                # Essayer de trouver le type le plus proche
                for vtype in valid_types:
                    if vtype in element_type:
                        return {"type": vtype, "confidence": 0.7, "success": True}
        
        return {"type": "unknown", "confidence": 0.0, "success": False}
    
    def classify_element_role(self,
                             element_image: Image.Image,
                             element_type: str,
                             context: Optional[str] = None) -> Dict[str, Any]:
        """
        Classifier le rôle sémantique d'un élément
        
        Args:
            element_image: Image de l'élément
            element_type: Type de l'élément
            context: Contexte additionnel
        
        Returns:
            Dict avec 'role' et 'confidence'
        """
        roles_list = "primary_action, cancel, submit, form_input, search_field, navigation, settings, close, delete, edit, save"
        
        prompt = f"""This is a {element_type}. What is its semantic role or purpose?
Choose ONLY ONE from: {roles_list}

Respond with just the role name, nothing else."""

        if context:
            prompt += f"\n\nContext: {context}"

        result = self.generate(
            prompt, image=element_image, temperature=0.1,
            assistant_prefill="The role is:",
        )
        
        if result["success"]:
            # Retirer le prefill du début pour extraire le rôle
            raw = result["response"]
            if raw.startswith("The role is:"):
                raw = raw[len("The role is:"):]
            role = raw.strip().lower()
            # Valider que c'est un rôle connu
            valid_roles = roles_list.split(", ")
            if role in valid_roles:
                return {"role": role, "confidence": 0.9, "success": True}
            else:
                # Essayer de trouver le rôle le plus proche
                for vrole in valid_roles:
                    if vrole in role:
                        return {"role": vrole, "confidence": 0.7, "success": True}
        
        return {"role": "unknown", "confidence": 0.0, "success": False}
    
    def extract_text(self, image: Image.Image) -> Dict[str, Any]:
        """
        Extraire le texte d'une image
        
        Args:
            image: Image PIL
        
        Returns:
            Dict avec 'text' extrait
        """
        prompt = "Extract all visible text from this image. Return only the text, nothing else."

        result = self.generate(
            prompt, image=image, temperature=0.1,
            assistant_prefill="Text:",
        )
        
        if result["success"]:
            # Retirer le prefill du début pour extraire le texte
            raw = result["response"]
            if raw.startswith("Text:"):
                raw = raw[len("Text:"):]
            return {"text": raw.strip(), "success": True}

        return {"text": "", "success": False, "error": result["error"]}

    # Taille minimum pour une classification fiable par le VLM
    _MIN_CLASSIFY_SIZE = 40

    def classify_element_complete(self, element_image: Image.Image) -> Dict[str, Any]:
        """
        Classifier complètement un élément UI en UN SEUL appel VLM.

        Optimisations :
        - Skip les crops < 40px (le VLM ne peut rien en tirer)
        - Retry 1 fois si réponse vide
        - Extraction JSON robuste (cherche {…} même dans du texte)
        """
        # Skip les images trop petites — deviner par la forme
        if (element_image.width < self._MIN_CLASSIFY_SIZE
                or element_image.height < self._MIN_CLASSIFY_SIZE):
            ratio = element_image.width / max(element_image.height, 1)
            if ratio > 3:
                guessed = "link"
            elif element_image.width < 24 and element_image.height < 24:
                guessed = "icon"
            else:
                guessed = "button"
            return {
                "type": guessed, "role": "unknown", "text": "",
                "confidence": 0.3, "success": True,
            }

        # Le system prompt contraint le thinking de qwen3-vl et réduit
        # drastiquement le nombre de tokens gaspillés en réflexion interne.
        # Sans system prompt, le modèle pense 500-800 tokens et épuise le budget.
        # Avec, il ne pense que 100-400 tokens et produit du JSON fiable.
        system_prompt = "You are a JSON-only UI classifier. No thinking. No explanation. Output raw JSON only."

        prompt = """Classify this UI element. Reply with ONLY a JSON object, nothing else.

Types: button, text_input, checkbox, radio, dropdown, tab, link, icon, table_row, menu_item
Roles: primary_action, cancel, submit, form_input, search_field, navigation, settings, close, delete, edit, save

Example 1: {"type": "button", "role": "submit", "text": "OK"}
Example 2: {"type": "text_input", "role": "form_input", "text": ""}
Example 3: {"type": "icon", "role": "close", "text": "X"}

Your answer:"""

        # Retry une fois si réponse vide
        for attempt in range(2):
            result = self.generate(
                prompt,
                image=element_image,
                system_prompt=system_prompt,
                temperature=0.1,
                max_tokens=300,
                force_json=False,
                assistant_prefill="{"
            )

            if not result["success"]:
                continue

            response_text = result["response"].strip()
            if not response_text:
                if attempt == 0:
                    continue
                break

            # Extraction JSON robuste
            parsed = self._extract_json_from_response(response_text)
            if parsed is not None:
                return self._validate_classification(parsed)

        return {
            "type": "unknown", "role": "unknown", "text": "",
            "confidence": 0.0, "success": False,
            "error": "VLM returned empty or unparseable response"
        }

    def _extract_json_from_response(self, text: str) -> Optional[Dict]:
        """Extrait un objet JSON d'une réponse VLM, même si entouré de texte."""
        import re as _re

        # Nettoyer le markdown
        if "```" in text:
            lines = text.split("\n")
            text = "\n".join([l for l in lines if not l.startswith("```")])
            text = text.strip()

        # Essai 1 : parse direct
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            pass

        # Essai 2 : trouver {…} dans le texte
        match = _re.search(r'\{[^{}]+\}', text)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass

        # Essai 3 : fixer les single quotes
        fixed = text.replace("'", '"')
        match = _re.search(r'\{[^{}]+\}', fixed)
        if match:
            try:
                return json.loads(match.group())
            except json.JSONDecodeError:
                pass

        logger.debug(f"Cannot extract JSON from VLM response: {text[:100]}")
        return None

    def _validate_classification(self, data: Dict) -> Dict[str, Any]:
        """Valide et normalise un résultat de classification."""
        valid_types = ["button", "text_input", "checkbox", "radio", "dropdown",
                      "tab", "link", "icon", "table_row", "menu_item"]
        valid_roles = ["primary_action", "cancel", "submit", "form_input",
                      "search_field", "navigation", "settings", "close",
                      "delete", "edit", "save"]

        elem_type = str(data.get("type", "unknown")).lower().strip()
        elem_role = str(data.get("role", "unknown")).lower().strip()
        elem_text = str(data.get("text", ""))

        if elem_type not in valid_types:
            elem_type = "unknown"
        if elem_role not in valid_roles:
            elem_role = "unknown"

        return {
            "type": elem_type, "role": elem_role, "text": elem_text,
            "confidence": 0.85, "success": True
        }

    def _encode_image_from_path(self, image_path: str) -> str:
        """Encoder une image depuis un fichier en base64"""
        with open(image_path, 'rb') as f:
            return base64.b64encode(f.read()).decode('utf-8')
    
    def _encode_image_from_pil(self, image: Image.Image) -> str:
        """Encoder une image PIL en base64 avec prétraitement optimisé"""
        # 1. Convertir en RGB si nécessaire (évite erreurs PNG transparent)
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # 1b. Minimum 32x32 (requis par qwen3-vl, sinon Ollama panic)
        min_size = 32
        if image.width < min_size or image.height < min_size:
            new_w = max(image.width, min_size)
            new_h = max(image.height, min_size)
            image = image.resize((new_w, new_h), Image.NEAREST)

        # 2. Redimensionnement intelligent : max 1280px sur le côté long
        max_size = 1280
        if max(image.size) > max_size:
            ratio = max_size / max(image.size)
            new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
            image = image.resize(new_size, Image.Resampling.LANCZOS)

        # 3. Sauvegarder en JPEG qualité 90 (plus léger, meilleur pour VLM)
        buffer = io.BytesIO()
        image.save(buffer, format='JPEG', quality=90)
        return base64.b64encode(buffer.getvalue()).decode('utf-8')
    
    def list_models(self) -> List[str]:
        """Lister les modèles disponibles dans Ollama"""
        try:
            response = requests.get(f"{self.endpoint}/api/tags", timeout=5)
            if response.status_code == 200:
                models = response.json().get('models', [])
                return [m['name'] for m in models]
        except Exception as e:
            logger.error(f"Error listing models: {e}")
        return []
    
    def pull_model(self, model_name: str) -> bool:
        """
        Télécharger un modèle dans Ollama
        
        Args:
            model_name: Nom du modèle à télécharger
        
        Returns:
            True si succès
        """
        try:
            logger.info(f"Pulling model {model_name}...")
            response = requests.post(
                f"{self.endpoint}/api/pull",
                json={"name": model_name},
                stream=True,
                timeout=600
            )
            
            if response.status_code == 200:
                for line in response.iter_lines():
                    if line:
                        data = json.loads(line)
                        if 'status' in data:
                            logger.info(f"  {data['status']}")
                return True
        except Exception as e:
            logger.error(f"Error pulling model: {e}")
        return False


# ============================================================================
# Fonctions utilitaires
# ============================================================================

def create_ollama_client(model: str = "qwen3-vl:8b",
                        endpoint: str = "http://localhost:11434") -> OllamaClient:
    """
    Créer un client Ollama
    
    Args:
        model: Nom du modèle VLM
        endpoint: URL de l'API Ollama
    
    Returns:
        OllamaClient configuré
    """
    return OllamaClient(endpoint=endpoint, model=model)


def check_ollama_available(endpoint: str = "http://localhost:11434") -> bool:
    """
    Vérifier si Ollama est disponible
    
    Args:
        endpoint: URL de l'API Ollama
    
    Returns:
        True si disponible
    """
    try:
        response = requests.get(f"{endpoint}/api/tags", timeout=5)
        return response.status_code == 200
    except (requests.RequestException, ConnectionError, TimeoutError):
        return False