app.py

import os
import pandas as pd
import chromadb
from google import genai
from sentence_transformers import SentenceTransformer, CrossEncoder
from typing import List, Dict
from flask import Flask, request, jsonify
from flask_cors import CORS
from datetime import datetime
import time

# ======================================================================
# CONFIGURATION
# ======================================================================

DATA_FILE_PATH = "data/QR.csv"

# CORRECTION CRITIQUE: Déplacement de la DB vers /tmp
# Ce répertoire est le seul garanti en écriture sur Hugging Face Spaces.
CHROMA_DB_PATH = "/tmp/bdd_ChromaDB"
COLLECTION_NAME = "qr_data_dual_embeddings"

Q_COLUMN_NAME = "Question"
R_COLUMN_NAME = "Reponse"
SYSTEM_PROMPT_PATH = "data/system_prompt.txt"

# Les chemins des modèles sont conservés (ils se mettront en cache dans /tmp grâce au Dockerfile)
SRC_CROSS_ENCODER = "models/mmarco-mMiniLMv2-L12-H384-v1"
SRC_PARAPHRASE = "models/paraphrase-mpnet-base-v2"

N_RESULTS_RETRIEVAL = 10
N_RESULTS_RERANK = 3

# Clé pour la route RAG (récupérée de l'environnement ou par défaut)
GEMINI_API_KEY_RAG = os.getenv("GEMINI_API_KEY", "AIzaSyDXXY7uSXryTxZ51jQFsSLcPnC_Ivt9V1g")
# NOUVELLE CLÉ demandée, mise en dur pour la route directe
GEMINI_API_KEY_DIRECT = "AIzaSyCpG2G3K0cZmTxWFO-c4OoOrW1fcTYQwgo"

GEMINI_MODEL = "gemini-2.5-flash"

MAX_CONVERSATION_HISTORY = 10

# Configuration pour l'accès externe (host et port)
API_HOST = '0.0.0.0'
API_PORT = 1212  # Le port 1212 est conservé, il doit être configuré dans le README.md

# ======================================================================
# VARIABLES GLOBALES
# ======================================================================

model_cross_encoder: CrossEncoder = None
model_paraphrase: SentenceTransformer = None
collection: chromadb.Collection = None
system_prompt: str = None
gemini_client_rag: genai.Client = None # Client pour la route RAG
gemini_client_direct: genai.Client = None # Client pour la route directe

conversation_histories: Dict[str, List[Dict[str, str]]] = {}
conversation_start_times: Dict[str, str] = {}

# ======================================================================
# CHARGEMENT DES RESSOURCES
# ======================================================================

def load_models():
    """Charge les modèles SentenceTransformer et CrossEncoder."""
    print("⏳ Chargement des modèles...")
    try:
        # Tente de charger localement, sinon télécharge (le cache se fera dans /tmp)
        cross_encoder = CrossEncoder(
            SRC_CROSS_ENCODER if os.path.exists(SRC_CROSS_ENCODER)
            else "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
        )
        paraphrase = SentenceTransformer(
            SRC_PARAPHRASE if os.path.exists(SRC_PARAPHRASE)
            else "sentence-transformers/paraphrase-mpnet-base-v2"
        )
        print("✅ Modèles chargés avec succès.")
        return cross_encoder, paraphrase
    except Exception as e:
        print(f"❌ Erreur chargement modèles: {e}")
        # Note: L'erreur de PermissionError est maintenant gérée par le Dockerfile
        raise

def load_data():
    """Charge le DataFrame depuis le CSV."""
    try:
        if not os.path.exists(DATA_FILE_PATH):
            print(f"⚠️ Fichier {DATA_FILE_PATH} non trouvé. Utilisation d'exemple.")
            df = pd.DataFrame({
                Q_COLUMN_NAME: ["Où est le soleil?", "Qui est l'IA?"],
                R_COLUMN_NAME: ["Le soleil est une étoile.", "L'IA est l'intelligence artificielle."]
            })
        else:
            df = pd.read_csv(DATA_FILE_PATH)
            print(f"✅ {len(df)} lignes chargées depuis {DATA_FILE_PATH}.")
        return df
    except Exception as e:
        print(f"❌ Erreur chargement données: {e}")
        raise

def load_system_prompt():
    """Charge le system prompt."""
    try:
        with open(SYSTEM_PROMPT_PATH, 'r', encoding='utf-8') as f:
            return f.read().strip()
    except FileNotFoundError:
        default = "Tu es un assistant utile et concis. Réponds à la requête de l'utilisateur."
        print(f"⚠️ System prompt non trouvé à {SYSTEM_PROMPT_PATH}. Utilisation du prompt par défaut.")
        return default

def initialize_gemini_client(api_key, client_name):
    """Initialise un client Google Gemini."""
    if api_key == "AIzaSyDXXY7uSXryTxZ51jQFsSLcPnC_Ivt9V1g":
        print(f"⚠️ AVIS pour {client_name}: Clé Gemini par défaut/placeholder détectée.")
    if api_key == "AIzaSyCpG2G3K0cZmTxWFO-c4OoOrW1fcTYQwgo":
        print(f"💡 AVIS pour {client_name}: Clé Gemini directe mise en dur détectée.")
    try:
        print(f"✅ Client Gemini '{client_name}' initialisé.")
        return genai.Client(api_key=api_key)
    except Exception as e:
        print(f"❌ Erreur lors de l'initialisation du client Gemini '{client_name}': {e}")
        raise

# ======================================================================
# CHROMADB SETUP
# ======================================================================

def setup_chromadb_collection(client, df, model_paraphrase):
    """Configure et remplit la collection ChromaDB."""
    total_docs = len(df) * 2

    # S'assurer que le répertoire de la DB existe
    os.makedirs(CHROMA_DB_PATH, exist_ok=True)

    try:
        collection = client.get_or_create_collection(name=COLLECTION_NAME)
    except Exception as e:
        print(f"❌ Erreur lors de l'accès à la collection ChromaDB: {e}")
        raise

    if collection.count() == total_docs and total_docs > 0:
        print(f"✅ Collection déjà remplie ({collection.count()} docs) dans {CHROMA_DB_PATH}.")
        return collection

    if total_docs == 0:
        print("⚠️ DataFrame vide. Collection non remplie.")
        return collection

    print(f"⏳ Remplissage de ChromaDB ({len(df)} lignes) à l'emplacement: {CHROMA_DB_PATH}...")

    docs, metadatas, ids = [], [], []

    for i, row in df.iterrows():
        question = str(row[Q_COLUMN_NAME])
        reponse = str(row[R_COLUMN_NAME])
        meta = {Q_COLUMN_NAME: question, R_COLUMN_NAME: reponse, "source_row": i}

        docs.append(question)
        metadatas.append({**meta, "type": "question"})
        ids.append(f"id_{i}_Q")

        docs.append(reponse)
        metadatas.append({**meta, "type": "reponse"})
        ids.append(f"id_{i}_R")

    embeddings = model_paraphrase.encode(docs, show_progress_bar=False).tolist()

    # Nettoyage et recréation (pour le cas où les données CSV ont changé)
    try:
        client.delete_collection(name=COLLECTION_NAME)
    except:
        pass

    collection = client.get_or_create_collection(name=COLLECTION_NAME)
    collection.add(embeddings=embeddings, documents=docs, metadatas=metadatas, ids=ids)

    print(f"✅ Collection remplie: {collection.count()} documents.")
    return collection

# ======================================================================
# RAG - RETRIEVAL & RERANKING
# ======================================================================

def retrieve_and_rerank(query_text, collection, model_paraphrase, model_cross_encoder):
    """Récupère et rerank les résultats."""
    print(f"🔍 Récupération pour: '{query_text[:40]}...'")

    query_emb = model_paraphrase.encode([query_text]).tolist()
    results = collection.query(
        query_embeddings=query_emb,
        n_results=N_RESULTS_RETRIEVAL,
        include=['documents', 'metadatas', 'distances']
    )

    if not results['ids'][0]:
        print("⚠️ Aucun résultat trouvé.")
        return pd.DataFrame()

    candidates = []
    cross_input = []

    for i, doc in enumerate(results['documents'][0]):
        meta = results['metadatas'][0][i]
        candidates.append({
            'question': meta[Q_COLUMN_NAME],
            'reponse': meta[R_COLUMN_NAME],
            'doc_type': meta.get('type'),
            'text_reranked': doc,
            'initial_distance': results['distances'][0][i]
        })
        cross_input.append([query_text, doc])

    scores = model_cross_encoder.predict(cross_input)
    for i, score in enumerate(scores):
        candidates[i]['rerank_score'] = score

    df = pd.DataFrame(candidates).sort_values('rerank_score', ascending=False)
    df = df.drop_duplicates(subset=['question', 'reponse'], keep='first')

    return df.head(N_RESULTS_RERANK)

def generate_rag_prompt(query_text, df_results, conversation_history):
    """Génère le prompt RAG final."""
    context = []
    if not df_results.empty:
        for _, row in df_results.iterrows():
            context.append(f"Q: {row['question']}\nR: {row['reponse']}")

    context_str = "\n---\n".join(context)

    history_str = ""
    if conversation_history:
        history_str = "HISTORIQUE:\n"
        # Ajout du contexte pour le LLM, mais on ne veut pas l'historique complet
        # On va limiter l'historique à l'affichage si on dépasse MAX_CONVERSATION_HISTORY
        display_history = conversation_history[-(MAX_CONVERSATION_HISTORY * 2):]
        for msg in display_history:
            role = "USER" if msg["role"] == "user" else "ASSISTANT"
            # On utilise 'content' pour le texte du message
            history_str += f"{role}: {msg['content']}\n"
        history_str += "\n"

    return f"""{history_str}UTILISATEUR: {query_text}

CONTEXTE (si utile):
[{context_str}]

INSTRUCTIONS:
- Réponds brièvement et clairement, en un minimum de 2 paragraphe et un maximum de 3 paragraphe, entre chaque paragraphe tu écrira un #
- Reste professionnel, met des emojis dans ta réponse tout en restant profesionnel
- Réponds uniquement à la question posée
- Ne mentionne NOVA que si l'utilisateur l'évoque clairement
- Mentionne obligatoirement Mathieu dans ta réponse"""

# ======================================================================
# GESTION HISTORIQUE
# ======================================================================

def get_conversation_history(session_id):
    """Récupère l'historique d'une session."""
    return conversation_histories.get(session_id, [])

def add_to_history(session_id, role, content):
    """Ajoute un message à l'historique."""
    if session_id not in conversation_histories:
        conversation_histories[session_id] = []

    conversation_histories[session_id].append({"role": role, "content": content})

    # Limiter la taille de l'historique conservé en mémoire
    if len(conversation_histories[session_id]) > MAX_CONVERSATION_HISTORY * 2:
        conversation_histories[session_id] = conversation_histories[session_id][-(MAX_CONVERSATION_HISTORY * 2):]

def clear_history(session_id):
    """Efface l'historique d'une session."""
    conversation_histories[session_id] = []

# ======================================================================
# CALL GEMINI
# ======================================================================

def call_gemini(final_prompt, system_prompt, gemini_client):
    """
    Appelle Google Gemini avec une logique de réessai en cas d'échec de l'API.
    Maximum de 10 tentatives.
    """
    MAX_RETRIES = 10
    
    # S'assurer que le client est bien initialisé
    if gemini_client is None:
        return "Erreur: Client Gemini non initialisé."

    for attempt in range(MAX_RETRIES):
        try:
            print(f"    📞 Tentative d'appel Gemini #{attempt + 1}...")
            # L'API Python de Google lève des exceptions `APIError` pour les échecs, 
            # y compris ceux qui correspondent aux 5xx.
            response = gemini_client.models.generate_content(
                model=GEMINI_MODEL,
                contents=f"{system_prompt}\n\n{final_prompt}"
            )
            # Si la réponse réussit, on sort de la boucle
            return response.text.replace("*", "")

        except Exception as e:
            # Ici, on capture toute erreur d'API ou de connexion.
            # On considère cela comme une erreur de service transitoire pour les réessais.
            error_message = str(e)
            print(f"    ❌ Erreur Gemini (Tentative {attempt + 1}/{MAX_RETRIES}): {error_message}")

            if attempt < MAX_RETRIES - 1:
                # Si ce n'est pas la dernière tentative, on attend avant de réessayer
                sleep_time = 2  # Attente de 2 secondes
                print(f"    😴 Attente de {sleep_time} secondes avant de réessayer...")
                time.sleep(sleep_time)
            else:
                # Dernière tentative échouée
                print("    🛑 Toutes les tentatives de réessai ont échoué.")
                return f"Erreur fatale après {MAX_RETRIES} tentatives: {error_message}"

    # Ne devrait jamais être atteint, mais par sécurité
    return "Erreur inconnue dans la boucle de réessai de Gemini."

# ======================================================================
# PROCESSUS DE RÉPONSE - RAG
# ======================================================================

def get_answer_rag_process(query_text, collection, model_paraphrase, model_cross_encoder, conversation_history):
    """Exécute le processus RAG complet."""
    print(f"\n{'='*50}")
    print(f"🚀 Traitement RAG: '{query_text}'")
    print(f"{'='*50}")

    df_results = retrieve_and_rerank(query_text, collection, model_paraphrase, model_cross_encoder)
    final_prompt = generate_rag_prompt(query_text, df_results, conversation_history)

    return final_prompt

# ======================================================================
# PROCESSUS DE RÉPONSE - DIRECT
# ======================================================================

def get_answer_direct_process(query_text):
    """Génère le prompt direct sans RAG."""
    return f"UTILISATEUR: {query_text}"

# ======================================================================
# INITIALISATION GLOBALE
# ======================================================================

def initialize_global_resources():
    """Initialise tous les modèles et ressources."""
    global model_cross_encoder, model_paraphrase, collection, system_prompt
    global gemini_client_rag, gemini_client_direct

    print("\n" + "="*50)
    print("⚙️  INITIALISATION RAG & Clients Gemini")
    print("="*50)

    try:
        model_cross_encoder, model_paraphrase = load_models()
        df = load_data()
        system_prompt = load_system_prompt()
        # Initialisation des deux clients
        gemini_client_rag = initialize_gemini_client(GEMINI_API_KEY_RAG, "RAG (Env/Default)")
        gemini_client_direct = initialize_gemini_client(GEMINI_API_KEY_DIRECT, "Direct (Hardcoded)")
    except Exception:
        # L'erreur est déjà print dans les fonctions de chargement
        return False

    try:
        print(f"⏳ Initialisation de ChromaDB à l'emplacement: {CHROMA_DB_PATH}")
        # Le PersistentClient créera les fichiers dans le chemin spécifié (maintenant dans /tmp)
        chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
        collection = setup_chromadb_collection(chroma_client, df, model_paraphrase)
        print("✅ INITIALISATION COMPLÈTE\n")
        return True
    except Exception as e:
        print(f"❌ Erreur lors de l'initialisation de ChromaDB ou du remplissage: {e}")
        return False

# ======================================================================
# FLASK API
# ======================================================================

app = Flask(__name__)
# CORS activé, permet les requêtes depuis n'importe quelle origine
CORS(app)

@app.route('/status', methods=['GET'])
def api_status():
    """Route de ping pour vérifier l'état de l'API."""
    return jsonify({"status": "everything is good"}), 200

@app.route('/api/get_answer', methods=['POST'])
def api_get_answer():
    """Endpoint principal pour obtenir une réponse avec RAG."""
    # Le client RAG utilise la clé d'environnement/par défaut
    if any(x is None for x in [model_cross_encoder, model_paraphrase, collection, system_prompt, gemini_client_rag]):
        return jsonify({"error": "Ressources RAG non chargées. Veuillez vérifier les logs d'initialisation."}), 500

    try:
        data = request.get_json()
        query_text = data.get('query_text')
        session_id = data.get('session_id', 'archive')

        if not query_text:
            generic_message = "Problème avec l'API, veuillez réessayer plus tard."
            return jsonify({"error": generic_message}), 500

        # Récupère historique
        history = get_conversation_history(session_id)

        # Génère prompt RAG
        rag_prompt = get_answer_rag_process(query_text, collection, model_paraphrase, model_cross_encoder, history)

        # Appelle Gemini avec le client RAG
        response = call_gemini(rag_prompt, system_prompt, gemini_client_rag)

        # Sauvegarde réponse
        add_to_history(session_id, "user", query_text)
        add_to_history(session_id, "assistant", response)

        return jsonify({"generated_response": response})

    except Exception as e:
        print(f"❌ Erreur générale de l'API RAG: {e}")
        generic_message = "Problème avec l'API RAG, veuillez réessayer plus tard."
        return jsonify({"error": generic_message}), 500

@app.route('/api/gemini_only', methods=['POST'])
def api_gemini_only():
    """NOUVELLE ROUTE : Endpoint pour les requêtes directes à Gemini sans RAG. Utilise la clé mise en dur."""
    # Le client direct utilise la clé mise en dur
    if gemini_client_direct is None:
        return jsonify({"error": "Client Gemini direct non initialisé. Vérifiez les logs."}), 500

    try:
        data = request.get_json()
        query_text = data.get('query_text')
        # On peut optionally récupérer un 'system_prompt_direct' pour customiser, sinon on utilise le prompt par défaut
        custom_system_prompt = data.get('system_prompt', system_prompt) 

        if not query_text:
            return jsonify({"error": "Paramètre 'query_text' manquant."}), 400

        print(f"\n{'='*50}")
        print(f"⚡ Traitement Direct: '{query_text}'")
        print(f"{'='*50}")

        # Génère le prompt final (juste la question)
        final_prompt = get_answer_direct_process(query_text)

        # Appelle Gemini avec le client direct
        # On utilise le 'system_prompt' par défaut ou un custom s'il est fourni
        response = call_gemini(final_prompt, custom_system_prompt, gemini_client_direct)

        # Pas d'ajout à l'historique de conversation ici car c'est une route directe sans session RAG/Historique

        return jsonify({"generated_response": response})

    except Exception as e:
        print(f"❌ Erreur générale de l'API Direct: {e}")
        generic_message = "Problème avec l'API directe, veuillez réessayer plus tard."
        return jsonify({"error": generic_message}), 500


@app.route('/api/clear_history', methods=['POST'])
def api_clear_history():
    """Efface l'historique d'une session."""
    try:
        data = request.get_json()
        session_id = data.get('session_id', 'archive')
        clear_history(session_id)

        return jsonify({"message": f"Historique effacé: {session_id}"})
    except Exception as e:
        generic_message = "Problème avec l'API, veuillez réessayer plus tard."
        return jsonify({"error": generic_message}), 500

# ======================================================================
# MAIN
# ======================================================================

if __name__ == '__main__':
    print("start app.py")
    if initialize_global_resources():

        # Récupération de l'adresse IP si possible (pour l'affichage)
        try:
            import socket
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))  # Connecte à un serveur externe pour trouver l'IP locale utilisée
            local_ip = s.getsockname()[0]
            s.close()
        except Exception:
            local_ip = "127.0.0.1"  # Fallback si échec

        print("\n" + "="*50)
        print("🌐 SERVEUR DÉMARRÉ")
        print(f"✅ API accessible à l'URL (via l'interface réseau locale): http://{local_ip}:{API_PORT}")
        print(f"✅ Route Status: http://{local_ip}:{API_PORT}/status")
        print(f"✅ Route RAG (avec Historique): http://{local_ip}:{API_PORT}/api/get_answer")
        print(f"✅ Route DIRECTE (Clé spéciale): http://{local_ip}:{API_PORT}/api/gemini_only")
        print(f"💡 N'oubliez pas de configurer 'app_port: 1212' et 'sdk: docker' dans votre README.md !")
        print("="*50 + "\n")

        # L'utilisation de host='0.0.0.0' dans app.run() permet l'accès depuis l'extérieur
        app.run(host=API_HOST, port=API_PORT, debug=False)
    else:
        print("❌ Impossible de démarrer le serveur. Veuillez vérifier les logs pour les erreurs d'initialisation.")