src/rag_handler.py

import os
import logging
from typing import Optional
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter

logger = logging.getLogger(__name__)

# Variables globales pour l'initialisation différée
_embeddings_model = None
_rag_handler_instance = None

# Utiliser /tmp qui est toujours writable dans les conteneurs
VECTOR_STORE_PATH = "/tmp/vector_store" 

def get_embeddings_model():
    """Obtient le modèle d'embeddings avec initialisation différée."""
    global _embeddings_model
    if _embeddings_model is None:
        try:
            from langchain_huggingface import HuggingFaceEmbeddings
            logger.info("Initialisation du modèle d'embeddings...")
            _embeddings_model = HuggingFaceEmbeddings(
                model_name='sentence-transformers/all-MiniLM-L6-v2',
                model_kwargs={'device': 'cpu'},
                encode_kwargs={'normalize_embeddings': True}
            )
            logger.info("✅ Modèle d'embeddings initialisé avec succès")
        except Exception as e:
            logger.error(f"❌ Erreur lors de l'initialisation du modèle d'embeddings: {e}")
            _embeddings_model = None
    return _embeddings_model

class RAGHandler:
    def __init__(self, knowledge_base_path: str = "/app/knowledge_base", lazy_init: bool = True):
        """
        Initialise le RAG Handler.
        
        Args:
            knowledge_base_path (str): Le chemin vers le dossier contenant les documents de connaissances (.md).
            lazy_init (bool): Si True, initialise le vector store seulement lors de la première utilisation.
        """
        self.knowledge_base_path = knowledge_base_path
        self.embeddings = None
        self.vector_store = None
        self._initialized = False
        
        # S'assurer que le répertoire /tmp/vector_store existe
        os.makedirs(VECTOR_STORE_PATH, exist_ok=True)
        
        if not lazy_init:
            self._initialize()

    def _initialize(self):
        """Initialise le RAG Handler de manière différée."""
        if self._initialized:
            return
        
        try:
            logger.info("Initialisation du RAG Handler...")
            self.embeddings = get_embeddings_model()
            
            if self.embeddings is None:
                logger.error("Impossible d'initialiser les embeddings")
                return
            
            self.vector_store = self._load_or_create_vector_store(self.knowledge_base_path)
            self._initialized = True
            logger.info("✅ RAG Handler initialisé avec succès")
            
        except Exception as e:
            logger.error(f"❌ Erreur lors de l'initialisation du RAG Handler: {e}")
            self._initialized = False

    def _load_documents(self, path: str) -> list:
        """Charge les documents depuis un chemin de répertoire spécifié."""
        try:
            if not os.path.exists(path):
                logger.warning(f"Répertoire {path} non trouvé")
                return []
                
            loader = DirectoryLoader(
                path,
                glob="**/*.md", 
                loader_cls=TextLoader,
                loader_kwargs={"encoding": "utf-8"}
            )
            logger.info(f"Chargement des documents depuis : {path}")
            documents = loader.load()
            logger.info(f"✅ {len(documents)} documents chargés")
            return documents
        except Exception as e:
            logger.error(f"❌ Erreur lors du chargement des documents: {e}")
            return []

    def _create_vector_store(self, knowledge_base_path: str) -> Optional[FAISS]:
        """Crée et sauvegarde la base de données vectorielle à partir des documents."""
        try:
            documents = self._load_documents(knowledge_base_path)
            if not documents:
                logger.warning("Aucun document trouvé - création d'un vector store vide")
                # Créer un document fictif pour initialiser le vector store
                from langchain.schema import Document
                dummy_doc = Document(
                    page_content="Document de test pour initialiser le vector store",
                    metadata={"source": "dummy"}
                )
                documents = [dummy_doc]
            
            logger.info(f"{len(documents)} documents chargés. Création des vecteurs...")
            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
            texts = text_splitter.split_documents(documents)
            
            vector_store = FAISS.from_documents(texts, self.embeddings)
            
            # Sauvegarder dans /tmp
            try:
                vector_store.save_local(VECTOR_STORE_PATH)
                logger.info(f"✅ Vector store créé et sauvegardé dans : {VECTOR_STORE_PATH}")
            except Exception as save_error:
                logger.warning(f"⚠️ Impossible de sauvegarder le vector store: {save_error}")
                # Continuer sans sauvegarde - le vector store reste en mémoire
            
            return vector_store
            
        except Exception as e:
            logger.error(f"❌ Erreur lors de la création du vector store: {e}")
            return None

    def _load_or_create_vector_store(self, knowledge_base_path: str) -> Optional[FAISS]:
        """Charge le vector store s'il existe, sinon le crée."""
        try:
            index_path = os.path.join(VECTOR_STORE_PATH, "index.faiss")
            if os.path.exists(index_path):
                logger.info(f"Chargement du vector store existant depuis : {VECTOR_STORE_PATH}")
                return FAISS.load_local(
                    VECTOR_STORE_PATH, 
                    embeddings=self.embeddings, 
                    allow_dangerous_deserialization=True 
                )
            else:
                logger.info("Aucun vector store trouvé. Création d'un nouveau...")
                return self._create_vector_store(knowledge_base_path)
        except Exception as e:
            logger.error(f"❌ Erreur lors du chargement/création du vector store: {e}")
            # En cas d'échec total, retourner None plutôt que planter
            return None

    def get_relevant_feedback(self, query: str, k: int = 1) -> list[str]:
        """Recherche les k conseils les plus pertinents pour une requête."""
        # Initialisation différée si nécessaire
        if not self._initialized:
            self._initialize()
        
        if not self.vector_store:
            logger.warning("Vector store non disponible - retour de conseils génériques")
            return [
                "Préparez vos réponses aux questions comportementales",
                "Montrez votre motivation pour le poste",
                "Donnez des exemples concrets de vos réalisations"
            ]
        
        try:
            results = self.vector_store.similarity_search(query, k=k)
            feedback = [doc.page_content for doc in results if doc.page_content.strip()]
            
            # Fallback si pas de résultats pertinents
            if not feedback:
                return ["Conseil général: Préparez-vous bien pour les entretiens futurs."]
                
            return feedback
        except Exception as e:
            logger.error(f"❌ Erreur lors de la recherche: {e}")
            return ["Conseil général: Travaillez sur vos compétences de communication."]

# Fonction pour obtenir une instance partagée
def get_rag_handler() -> Optional[RAGHandler]:
    """Obtient une instance partagée du RAG Handler."""
    global _rag_handler_instance
    if _rag_handler_instance is None:
        try:
            _rag_handler_instance = RAGHandler(lazy_init=True)
        except Exception as e:
            logger.error(f"❌ Erreur lors de la création du RAG Handler: {e}")
            _rag_handler_instance = None
    return _rag_handler_instance

if __name__ == '__main__':
    print("Test du RAG Handler avec /tmp vector store...")
    handler = RAGHandler(knowledge_base_path="/app/knowledge_base", lazy_init=False)
    
    test_query = "gestion du stress"
    feedback = handler.get_relevant_feedback(test_query, k=2)
    
    print(f"\nTest de recherche pour : '{test_query}'")
    if feedback:
        print("Feedback trouvé :")
        for f in feedback:
            print(f"- {f[:150]}...")
    else:
        print("Aucun feedback trouvé.")