Create llm_engine.py

llm_engine.py

@@ -0,0 +1,355 @@
+"""
+llm_engine.py — Moteur d'inférence LLM
+Gère le chargement des modèles, la génération de texte et le fallback.
+"""
+
+import time
+import logging
+from typing import Optional, Tuple, Generator
+
+logging.basicConfig(level=logging.INFO, format="[%(name)s] %(levelname)s: %(message)s")
+logger = logging.getLogger("LLMEngine")
+
+from config import (
+    LLM_MODEL, QA_MODEL, MAX_NEW_TOKENS, TEMPERATURE,
+    TOP_P, REPETITION_PENALTY, DO_SAMPLE
+)
+from utils import clean_response, is_valid_response, format_error_message
+from prompts import build_chat_prompt, build_qa_context
+
+
+class LLMEngine:
+    """
+    Moteur d'inférence principal avec système de fallback en cascade.
+    
+    Cascade de fallback :
+    1. Modèle LLM principal (génération chat)
+    2. Modèle QA (question-answering sur contexte)
+    3. Réponse de fallback statique
+    
+    Cette architecture garantit qu'une réponse est toujours retournée,
+    même si les modèles principaux ne sont pas disponibles.
+    """
+
+    def __init__(self):
+        self.text_pipeline = None      # Pipeline génération de texte (LLM)
+        self.qa_pipeline = None        # Pipeline question-answering (fallback)
+        self.models_loaded = False
+        self._load_models()
+
+    def _load_models(self) -> None:
+        """
+        Charge les modèles IA de manière sécurisée.
+        Utilise lazy loading — ne bloque pas le démarrage si un modèle échoue.
+        """
+        logger.info(f"Chargement du modèle LLM : {LLM_MODEL}")
+        
+        # Import différé pour éviter les erreurs si transformers n'est pas installé
+        try:
+            from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+            import torch
+            
+            # Détection du device (GPU si disponible, sinon CPU)
+            device = 0 if torch.cuda.is_available() else -1
+            device_name = "GPU (CUDA)" if device == 0 else "CPU"
+            logger.info(f"Device sélectionné : {device_name}")
+            
+            # ── Pipeline LLM principal ──────────────────────────────────────
+            try:
+                self.text_pipeline = pipeline(
+                    "text-generation",
+                    model=LLM_MODEL,
+                    device=device,
+                    # Paramètres d'optimisation mémoire
+                    torch_dtype="auto",           # Sélectionne float16 sur GPU
+                    trust_remote_code=False,
+                )
+                logger.info(f"✅ LLM chargé : {LLM_MODEL}")
+            except Exception as e:
+                logger.warning(f"⚠️ Échec LLM principal : {e}")
+                self.text_pipeline = None
+            
+            # ── Pipeline QA de fallback ─────────────────────────────────────
+            try:
+                self.qa_pipeline = pipeline(
+                    "question-answering",
+                    model=QA_MODEL,
+                    device=device,
+                )
+                logger.info(f"✅ QA pipeline chargé : {QA_MODEL}")
+            except Exception as e:
+                logger.warning(f"⚠️ Échec QA pipeline : {e}")
+                self.qa_pipeline = None
+            
+            self.models_loaded = self.text_pipeline is not None or self.qa_pipeline is not None
+            
+        except ImportError:
+            logger.error("❌ transformers non installé. Installez avec : pip install transformers torch")
+            self.models_loaded = False
+
+    def generate(
+        self,
+        user_message: str,
+        conversation_history: list,
+        domain: Optional[str] = None,
+        max_tokens: int = MAX_NEW_TOKENS,
+    ) -> Tuple[str, str]:
+        """
+        Génère une réponse pour le message utilisateur.
+        
+        Cascade de fallback :
+        1. LLM principal → génération de texte contextualisée
+        2. QA pipeline → extraction de réponse depuis contexte
+        3. Message d'erreur informatif
+        
+        Args:
+            user_message: Le message de l'utilisateur
+            conversation_history: Historique récent de la conversation
+            domain: Domaine détecté pour les prompts spécialisés
+            max_tokens: Nombre maximum de tokens à générer
+        
+        Returns:
+            Tuple (réponse, source) où source ∈ {'llm', 'qa', 'fallback'}
+        """
+        start_time = time.time()
+        
+        # ── Tentative 1 : LLM principal ────────────────────────────────────
+        if self.text_pipeline is not None:
+            response, source = self._generate_with_llm(
+                user_message, conversation_history, domain, max_tokens
+            )
+            if response:
+                elapsed = time.time() - start_time
+                logger.info(f"[LLM] Réponse générée en {elapsed:.2f}s ({source})")
+                return response, source
+        
+        # ── Tentative 2 : QA pipeline ──────────────────────────────────────
+        if self.qa_pipeline is not None:
+            response, source = self._generate_with_qa(user_message, domain)
+            if response:
+                elapsed = time.time() - start_time
+                logger.info(f"[QA] Réponse extraite en {elapsed:.2f}s ({source})")
+                return response, source
+        
+        # ── Fallback final ─────────────────────────────────────────────────
+        logger.warning("Tous les modèles ont échoué. Retour message de fallback.")
+        fallback = self._get_fallback_response(user_message, domain)
+        return fallback, "fallback"
+
+    def _generate_with_llm(
+        self,
+        user_message: str,
+        conversation_history: list,
+        domain: Optional[str],
+        max_tokens: int,
+    ) -> Tuple[Optional[str], str]:
+        """
+        Génère avec le LLM principal (pipeline text-generation).
+        
+        Utilise le format ChatML pour structurer le prompt.
+        Extrait uniquement la partie 'assistant' de la sortie.
+        
+        Returns:
+            Tuple (réponse nettoyée, 'llm') ou (None, 'llm_failed')
+        """
+        try:
+            # Construction du prompt formaté
+            prompt = build_chat_prompt(conversation_history, user_message, domain)
+            
+            # Génération
+            outputs = self.text_pipeline(
+                prompt,
+                max_new_tokens=max_tokens,
+                temperature=TEMPERATURE,
+                top_p=TOP_P,
+                repetition_penalty=REPETITION_PENALTY,
+                do_sample=DO_SAMPLE,
+                return_full_text=False,   # Retourne uniquement la partie générée
+                pad_token_id=self.text_pipeline.tokenizer.eos_token_id,
+            )
+            
+            if not outputs or not outputs[0]:
+                return None, "llm_empty"
+            
+            generated_text = outputs[0].get("generated_text", "")
+            
+            # Extraction de la réponse : prend tout avant le prochain <|user|>
+            if "<|user|>" in generated_text:
+                generated_text = generated_text.split("<|user|>")[0]
+            if "<|system|>" in generated_text:
+                generated_text = generated_text.split("<|system|>")[0]
+            
+            # Nettoyage
+            response = clean_response(generated_text)
+            
+            if is_valid_response(response):
+                return response, "llm"
+            
+            return None, "llm_short"
+            
+        except Exception as e:
+            logger.error(f"Erreur génération LLM : {e}")
+            return None, "llm_error"
+
+    def _generate_with_qa(
+        self,
+        user_message: str,
+        domain: Optional[str],
+    ) -> Tuple[Optional[str], str]:
+        """
+        Extrait une réponse via le pipeline question-answering.
+        
+        Utilise un contexte enrichi selon le domaine détecté.
+        Plus fiable que le LLM pour les questions factuelles courtes.
+        
+        Returns:
+            Tuple (réponse, 'qa') ou (None, 'qa_failed')
+        """
+        try:
+            context = build_qa_context(domain)
+            
+            result = self.qa_pipeline(
+                question=user_message,
+                context=context,
+                max_answer_len=256,
+            )
+            
+            answer = result.get("answer", "").strip()
+            score = result.get("score", 0)
+            
+            logger.info(f"[QA] Score de confiance : {score:.3f}")
+            
+            # Accepter la réponse si la confiance est suffisante
+            if score > 0.1 and is_valid_response(answer):
+                return answer, "qa"
+            
+            return None, "qa_low_confidence"
+            
+        except Exception as e:
+            logger.error(f"Erreur QA pipeline : {e}")
+            return None, "qa_error"
+
+    def _get_fallback_response(self, user_message: str, domain: Optional[str]) -> str:
+        """
+        Génère une réponse de secours informative basée sur des règles simples.
+        
+        Analyse les mots-clés de la question pour retourner une réponse
+        pertinente depuis un mini-dictionnaire intégré.
+        
+        Args:
+            user_message: La question de l'utilisateur
+            domain: Domaine détecté
+        
+        Returns:
+            Réponse textuelle de fallback
+        """
+        msg_lower = user_message.lower()
+        
+        # Mini-base de réponses intégrées (fallback ultime)
+        fallback_rules = {
+            # Réseaux
+            ("switch", "commutateur"): (
+                "Un **switch** (commutateur) est un équipement réseau de couche 2 (OSI) "
+                "qui interconnecte des appareils dans un réseau local (LAN). "
+                "Il utilise les adresses MAC pour acheminer les trames vers le bon port. "
+                "Commandes Cisco de base :\n"
+                "```\nSwitch> enable\nSwitch# show mac address-table\nSwitch# show interfaces\n```"
+            ),
+            ("routeur", "router"): (
+                "Un **routeur** est un équipement réseau de couche 3 (OSI) qui interconnecte "
+                "plusieurs réseaux différents. Il utilise les adresses IP et une table de routage "
+                "pour acheminer les paquets.\n"
+                "Commandes Cisco de base :\n"
+                "```\nRouter> enable\nRouter# show ip route\nRouter# show ip interface brief\n```"
+            ),
+            ("vlan",): (
+                "Un **VLAN** (Virtual LAN) permet de segmenter logiquement un réseau physique "
+                "en plusieurs réseaux virtuels isolés. Configuration Cisco :\n"
+                "```\nSwitch(config)# vlan 10\nSwitch(config-vlan)# name SERVEURS\n"
+                "Switch(config)# interface fa0/1\nSwitch(config-if)# switchport mode access\n"
+                "Switch(config-if)# switchport access vlan 10\n```"
+            ),
+            ("ospf",): (
+                "**OSPF** (Open Shortest Path First) est un protocole de routage dynamique "
+                "à état de lien (Link-State). Il utilise l'algorithme de Dijkstra pour calculer "
+                "les meilleurs chemins. Configuration Cisco :\n"
+                "```\nRouter(config)# router ospf 1\n"
+                "Router(config-router)# network 192.168.1.0 0.0.0.255 area 0\n```"
+            ),
+            # Cybersécurité
+            ("vpn",): (
+                "Un **VPN** (Virtual Private Network) crée un tunnel chiffré entre deux points "
+                "sur Internet, assurant confidentialité et intégrité des données. "
+                "Types principaux : Site-to-Site (deux réseaux), Remote Access (nomade). "
+                "Protocoles : IPSec, OpenVPN, WireGuard, SSL/TLS."
+            ),
+            ("firewall", "pare-feu"): (
+                "Un **pare-feu** (firewall) filtre le trafic réseau selon des règles de sécurité. "
+                "Types : stateless (filtre par paquet), stateful (suit les connexions), "
+                "applicatif (inspecte le contenu - NGFW). "
+                "Il constitue la première ligne de défense du réseau."
+            ),
+            # IA/ML
+            ("machine learning", "apprentissage automatique"): (
+                "Le **Machine Learning** est un sous-domaine de l'IA où les algorithmes "
+                "apprennent automatiquement à partir de données. "
+                "3 types principaux :\n"
+                "- **Supervisé** : données étiquetées (classification, régression)\n"
+                "- **Non supervisé** : données non étiquetées (clustering)\n"
+                "- **Par renforcement** : apprentissage par récompenses"
+            ),
+            ("llm", "grand modèle de langage"): (
+                "Un **LLM** (Large Language Model) est un modèle de langage entraîné sur "
+                "d'immenses corpus textuels. Il utilise l'architecture Transformer. "
+                "Exemples : GPT-4, Claude, LLaMA, Mistral. "
+                "Ils excellent dans la génération de texte, la traduction, le code et le Q&A."
+            ),
+        }
+        
+        # Cherche la première règle correspondante
+        for keywords, response in fallback_rules.items():
+            if any(kw in msg_lower for kw in keywords):
+                return response
+        
+        # Réponse générique selon le domaine
+        domain_generic = {
+            "réseaux": (
+                "Je n'ai pas trouvé de réponse précise à votre question sur les réseaux. "
+                "Pour approfondir ce sujet, je recommande :\n"
+                "- La documentation Cisco (cisco.com/c/en/us/support)\n"
+                "- Les cours CCNA sur NetAcad (netacad.com)\n"
+                "- Packet Tracer pour la pratique en simulation"
+            ),
+            "cybersécurité": (
+                "Je n'ai pas trouvé de réponse précise à votre question en cybersécurité. "
+                "Ressources recommandées :\n"
+                "- OWASP (owasp.org) pour la sécurité applicative\n"
+                "- SANS Institute (sans.org) pour les formations\n"
+                "- TryHackMe / HackTheBox pour la pratique"
+            ),
+            "ia": (
+                "Je n'ai pas trouvé de réponse précise à votre question sur l'IA/ML. "
+                "Ressources recommandées :\n"
+                "- Coursera / DeepLearning.AI (Andrew Ng)\n"
+                "- Hugging Face (huggingface.co) pour les modèles\n"
+                "- Papers With Code (paperswithcode.com)"
+            ),
+        }
+        
+        if domain and domain in domain_generic:
+            return domain_generic[domain]
+        
+        return (
+            "Je n'ai pas pu générer une réponse complète à votre question. "
+            "Pourriez-vous la reformuler ou la préciser ? "
+            "WENDAA AI couvre les domaines : réseaux, cybersécurité, IA/ML et data."
+        )
+
+    def get_status(self) -> dict:
+        """Retourne le statut des modèles chargés."""
+        return {
+            "llm_loaded": self.text_pipeline is not None,
+            "qa_loaded": self.qa_pipeline is not None,
+            "model_name": LLM_MODEL,
+            "qa_model_name": QA_MODEL,
+        }