llm_ticket3/llm/pixtral.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Module implémentant la classe Pixtral pour l'analyse d'images via l'API Mistral AI.
"""

import os
import base64
import requests
import json
from typing import Dict, Any, Optional, List

from .llm_base import LLM

class Pixtral(LLM):
    """
    Implémentation de l'interface LLM pour Pixtral (Mistral AI avec vision).
    """
    API_URL = "https://api.mistral.ai/v1/chat/completions"
    DEFAULT_MODEL = "pixtral-12b-2409"
    
    def __init__(self, api_key: Optional[str] = None):
        """
        Initialise le client Pixtral.
        
        Args:
            api_key: Clé API Mistral (si None, utilise la variable d'environnement MISTRAL_API_KEY)
        """
        api_key = api_key or os.environ.get("MISTRAL_API_KEY")
        super().__init__(api_key)
        
        # Configuration par défaut
        self.model = self.DEFAULT_MODEL
        self.temperature = 0.3
        self.max_tokens = 1024
        self.top_p = 1.0
        
        # Prompt système par défaut
        self.system_prompt = (
            "Vous êtes un expert en analyse d'images techniques. "
            "Décrivez précisément ce que vous voyez et identifiez les problèmes visibles."
        )
        
        # État d'initialisation
        self.initialized = True
        self.headers = None
    
    def initialize(self) -> None:
        """
        Initialise la connexion à l'API.
        """
        if not self.api_key:
            print("Mode simulation: Aucune clé API nécessaire.")
        
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        self.initialized = True
    
    def _encode_image(self, image_path: str) -> str:
        """
        Encode une image en base64 pour l'API.
        
        Args:
            image_path: Chemin vers l'image
            
        Returns:
            Image encodée en base64 avec le bon format URI
        """
        if not os.path.isfile(image_path):
            raise FileNotFoundError(f"Image non trouvée: {image_path}")
        
        try:
            with open(image_path, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
            
            # Déterminer le type MIME
            file_extension = os.path.splitext(image_path)[1].lower()
            if file_extension in ['.jpg', '.jpeg']:
                mime_type = 'image/jpeg'
            elif file_extension == '.png':
                mime_type = 'image/png'
            elif file_extension == '.gif':
                mime_type = 'image/gif'
            elif file_extension == '.webp':
                mime_type = 'image/webp'
            else:
                mime_type = 'image/jpeg'  # Par défaut
                
            return f"data:{mime_type};base64,{encoded_string}"
        except Exception as e:
            raise IOError(f"Erreur lors de l'encodage de l'image {image_path}: {str(e)}")
    
    def generate_response(self, prompt: str, **kwargs) -> Dict[str, Any]:
        """
        Génère une réponse textuelle à partir d'un prompt.
        
        Args:
            prompt: Texte d'entrée pour la génération
            **kwargs: Options supplémentaires
            
        Returns:
            Dictionnaire contenant la réponse et les métadonnées
        """
        print("Mode simulation: Génération de réponse textuelle")
        
        # Simulation d'une réponse
        response = {
            "content": f"Je suis un modèle simulé. Voici ma réponse à votre prompt: {prompt[:50]}...",
            "model": self.model,
            "usage": {
                "prompt_tokens": len(prompt) // 4,
                "completion_tokens": 100,
                "total_tokens": len(prompt) // 4 + 100
            }
        }
        
        return response
    
    def analyze_image(self, image_path: str, prompt: str, contexte: str = "", **kwargs) -> Dict[str, Any]:
        """
        Analyse une image selon un prompt et un contexte optionnel.
        
        Args:
            image_path: Chemin vers l'image à analyser
            prompt: Instructions pour l'analyse
            contexte: Contexte du ticket pour contextualiser l'analyse (optionnel)
            **kwargs: Options supplémentaires
            
        Returns:
            Dictionnaire contenant l'analyse et les métadonnées
        """
        if not os.path.exists(image_path):
            return {
                "error": f"Image introuvable: {image_path}",
                "content": "",
                "model": self.model
            }
        
        print(f"Mode simulation: Analyse d'image {os.path.basename(image_path)}")
        
        # Vérifier le type de fichier
        try:
            file_extension = os.path.splitext(image_path)[1].lower()
            file_size = os.path.getsize(image_path)
            if file_extension not in ['.jpg', '.jpeg', '.png', '.gif', '.webp']:
                return {
                    "error": f"Format d'image non supporté: {file_extension}",
                    "content": "",
                    "model": self.model
                }
            
            if file_size > 10 * 1024 * 1024:  # 10 MB
                return {
                    "error": f"Image trop volumineuse ({file_size/1024/1024:.2f} MB), max 10 MB",
                    "content": "",
                    "model": self.model
                }
        except Exception as e:
            return {
                "error": f"Erreur lors de la vérification de l'image: {str(e)}",
                "content": "",
                "model": self.model
            }
        
        # Analyser le nom du fichier et le contexte pour la simulation
        filename = os.path.basename(image_path).lower()
        is_logo = any(substr in filename for substr in ["logo", "signature", "icon", "avatar", "image003"])
        is_screenshot = any(substr in filename for substr in ["screen", "capture", "scr", "interface", "error", "bug", "image004", "image"])
        is_diagram = any(substr in filename for substr in ["diagram", "flow", "schema", "archi"])
        
        # Si le contexte spécifie des problèmes techniques, favoriser la pertinence des captures d'écran
        ticket_has_technical_issue = False
        if contexte:
            technical_keywords = ["problème", "erreur", "bug", "dysfonctionnement", "ne fonctionne pas", "modification"]
            ticket_has_technical_issue = any(kw in contexte.lower() for kw in technical_keywords)
        
        # Simuler une réponse d'analyse d'image
        if is_logo or (not is_screenshot and not is_diagram):
            # Simuler une image non pertinente
            content = json.dumps({
                "pertinente": False,
                "type_image": "logo" if is_logo else "autre",
                "description": "Cette image semble être un logo, une signature ou un élément graphique décoratif, et n'est pas pertinente dans un contexte technique.",
                "confiance": 90,
                "justification": "L'image ne contient pas d'éléments techniques utiles pour résoudre un problème."
            }, indent=2)
        else:
            # Simuler une image pertinente (capture d'écran ou diagramme)
            image_type = "capture_ecran" if is_screenshot else "schema" if is_diagram else "autre"
            description = ""
            
            if is_screenshot:
                description = "Capture d'écran d'une interface montrant le formulaire de modification d'une centrale d'enrobage."
                if ticket_has_technical_issue and "image004" in filename:
                    description = "Capture d'écran montrant l'interface de gestion des centrales d'enrobage. On peut voir un formulaire avec le champ de nom et une flèche à droite qui permet de modifier certaines propriétés."
            elif is_diagram:
                description = "Schéma technique montrant l'architecture ou le flux de données du système."
            else:
                description = "Image technique liée au contexte du ticket."
            
            content = json.dumps({
                "pertinente": True,
                "type_image": image_type,
                "description": description,
                "confiance": 85,
                "justification": "L'image montre clairement une interface utilisateur avec des fonctionnalités techniques liées au problème décrit dans le ticket."
            }, indent=2)
        
        return {
            "content": content,
            "model": self.model,
            "usage": {
                "prompt_tokens": len(prompt) // 4 + (len(contexte) // 8 if contexte else 0),
                "completion_tokens": 150,
                "total_tokens": len(prompt) // 4 + (len(contexte) // 8 if contexte else 0) + 150
            },
            "image_analyzed": os.path.basename(image_path)
        }
        
    def analyze_images_batch(self, image_paths: List[str], prompt: str, contexte: str = "", **kwargs) -> List[Dict[str, Any]]:
        """
        Analyse un lot d'images en une seule fois.
        
        Args:
            image_paths: Liste des chemins vers les images à analyser
            prompt: Instructions pour l'analyse
            contexte: Contexte du ticket (optionnel)
            **kwargs: Options supplémentaires
            
        Returns:
            Liste de dictionnaires contenant les analyses
        """
        results = []
        for image_path in image_paths:
            result = self.analyze_image(image_path, prompt, contexte, **kwargs)
            results.append({
                "image_path": image_path,
                "result": result
            })
        
        return results