Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,873 +1,318 @@
 import os
 import streamlit as st
-import requests
-from dotenv import load_dotenv
-import json # Pour l'export JSON
-from datetime import datetime # Pour nommer les fichiers d'export
-import time # Pour les messages d'attente si besoin (non strictement utilisé ici mais utile)
-from PIL import Image  # Pour potentiellement ouvrir/manipuler l'image si nécessaire (basé sur pillow)
-from io import BytesIO # Pour lire les données binaires de l'image
-
-# Import Google Generative AI library
+import json
+from datetime import datetime
 import google.generativeai as genai
-
-# Import DuckDuckGo Search (for web search feature)
-# Assurez-vous d'avoir installé cette bibliothèque : pip install duckduckgo_search
-from duckduckgo_search import DDGS # Importe la classe DDGS
+from duckduckgo_search import DDGS
+from dotenv import load_dotenv
 
 # -----------------------------------------------------------------------------
-# Environment & constants
+# Configuration de l'environnement et des constantes
 # -----------------------------------------------------------------------------
 load_dotenv()
-HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN") # Gardé pour la structure, même si plus de modèles HF actifs
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") # Add Google API Key (assumed from Streamlit Cloud secrets)
 
-# API URL base for Hugging Face Inference API (Gardé pour la structure)
-HF_BASE_API_URL = "https://api-inference.huggingface.co/models/"
+# Configurez vos clés API dans un fichier .env ou dans les secrets de Streamlit Cloud
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+# Configuration de la bibliothèque Google
+if GOOGLE_API_KEY:
+    try:
+        genai.configure(api_key=GOOGLE_API_KEY)
+    except Exception as e:
+        # Cette erreur est critique, donc on l'affiche de manière proéminente.
+        st.error(f"Erreur de configuration de l'API Google : {e}")
+else:
+    # Avertissement si la clé n'est pas trouvée au démarrage.
+    st.warning("La clé API Google (GOOGLE_API_KEY) n'est pas configurée. L'application ne pourra pas fonctionner.")
+
 
 # -----------------------------------------------------------------------------
-# Define Available Models and Default Parameters per Task Type
+# Définition des modèles disponibles
 # -----------------------------------------------------------------------------
-# Structure: {'id': 'model_id', 'name': 'Display Name', 'provider': 'huggingface'/'google', 'type': 'text'/'t2i', 'params': {...} }
+# Remplacé par les modèles Gemini les plus récents et pertinents.
+# 'latest' pointe toujours vers la version la plus récente du modèle.
 AVAILABLE_MODELS = [
-    # --- Google Text Models (Only keeping Flash as requested) ---
     {
         "id": "gemini-1.5-flash-latest",
-        "name": "Gemini 1.5 Flash (Google)",
+        "name": "Gemini 1.5 Flash (Rapide et efficace)",
         "provider": "google",
-        "type": "text",
-        "params": {
-            "max_new_tokens": 200,
-            "temperature": 0.6,
-            "top_p": 0.9,
-        },
     },
-     {
-        "id": "gemini-2.0-flash-lite-001", # Ajouté comme demandé
-        "name": "Gemini 2.0 Flash Lite 001 (Google)",
+    {
+        "id": "gemini-1.5-pro-latest",
+        "name": "Gemini 1.5 Pro (Le plus performant)",
         "provider": "google",
-        "type": "text", # Supposé compatible pour le chat texte
-        "params": {
-             "max_new_tokens": 200,
-             "temperature": 0.6,
-             "top_p": 0.9,
-        },
     },
-     {
-         "id": "gemini-2.0-flash-preview-image-generation", # Ajouté comme demandé - NOTE : Son comportement exact comme modèle de chat texte peut varier.
-         "name": "Gemini 2.0 Flash Preview Image Gen (Google)",
-         "provider": "google",
-         "type": "text", # Supposé compatible pour le chat texte malgré son nom
-         "params": { # Paramètres par défaut pour le texte
-             "max_new_tokens": 200,
-             "temperature": 0.6,
-             "top_p": 0.9,
-         },
-     },
-    # Retiré : gemini-1.5-pro-latest, gemini-pro, gemini-2.5-flash-preview, gemini-2.5-pro-preview
-    # Retiré : Mistral (Hugging Face Text)
-    # Retiré : Stable Diffusion XL Base 1.0 (Hugging Face T2I) et les autres SD.
 ]
 
-# Separate models lists by type (still potentially useful for future expansion)
-MODELS_BY_TYPE = {m_type: [m for m in AVAILABLE_MODELS if m['type'] == m_type] for m_type in set(m['type'] for m in AVAILABLE_MODELS)}
-
-# Default task is now fixed to 'text'
-DEFAULT_TASK = 'text'
-# Find the first text model as default
-first_text_model = next((m for m in AVAILABLE_MODELS if m['type'] == DEFAULT_TASK), None)
-if first_text_model:
-    DEFAULT_MODEL_ID = first_text_model['id']
-elif AVAILABLE_MODELS: # Fallback to the very first model if no text model exists
-    DEFAULT_MODEL_ID = AVAILABLE_MODELS[0]['id']
-else: # No models defined at all
-    DEFAULT_MODEL_ID = None
-
+DEFAULT_MODEL_ID = "gemini-1.5-flash-latest"
 
 # -----------------------------------------------------------------------------
-# Session-state initialisation
-# Initialize all session state variables used BEFORE any UI or logic accesses them.
+# Initialisation du Session State
 # -----------------------------------------------------------------------------
-DEFAULT_SYSTEM_MESSAGE = "Vous êtes KolaChatBot, un assistant IA serviable, créatif et honnête. Répondez en français."
-DEFAULT_STARTER_MESSAGE = "Bonjour ! Je suis KolaChatBot. Choisissez un modèle pour commencer. 🤖"
-
-# Ensure initial task is set (now fixed to 'text')
-st.session_state.selected_task = DEFAULT_TASK
-
-# Ensure initial model ID is set
-if 'selected_model_id' not in st.session_state:
-    st.session_state.selected_model_id = DEFAULT_MODEL_ID
-
-# Initialize all parameter types in session state with hardcoded defaults.
-# These will NOT be automatically reset when model changes, only on explicit button click
-# or if not initialized.
-if 'max_response_length' not in st.session_state: st.session_state.max_response_length = 200
-if 'temperature' not in st.session_state: st.session_state.temperature = 0.6
-if 'top_p' not in st.session_state: st.session_state.top_p = 0.9
-# Image-specific parameters are kept for completeness, but will not be used in text-only mode
-if 'num_inference_steps' not in st.session_state: st.session_state.num_inference_steps = 50
-if 'guidance_scale' not in st.session_state: st.session_state.guidance_scale = 7.5
-if 'image_height' not in st.session_state: st.session_state.image_height = 512
-if 'image_width' not in st.session_state: st.session_state.image_width = 512
-if 'negative_prompt' not in st.session_state: st.session_state.negative_prompt = ""
-
-# Initialize system and starter messages in session state if not exists
-if "system_message" not in st.session_state:
-    st.session_state.system_message = DEFAULT_SYSTEM_MESSAGE
-if "starter_message" not in st.session_state:
-    st.session_state.starter_message = DEFAULT_STARTER_MESSAGE
-
-# Initialize avatars
-if "avatars" not in st.session_state:
-    st.session_state.avatars = {"user": "👤", "assistant": "🤖"}
-
-# Initialize chat history in session state if not exists.
-# Start with the initial assistant message with type 'text'
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message, "type": "text"}]
-
-# State variable to track if a reset was triggered by buttons
-if '_reset_triggered' not in st.session_state:
-    st.session_state._reset_triggered = False
-
-# State variables to detect changes for rerun/reset logic
-# Track task and model ID *before* the radio/selectbox potentially changes them mid-run
-if '_prev_task_before_radio' not in st.session_state: st.session_state._prev_task_before_radio = st.session_state.selected_task
-if '_prev_model_id_before_selectbox' not in st.session_state:
-     st.session_state._prev_model_id_before_selectbox = st.session_state.selected_model_id
-
-# State for web search feature
-if 'enable_web_search' not in st.session_state:
-    st.session_state.enable_web_search = False
+def initialize_session_state():
+    """Initialise toutes les variables nécessaires dans le session state pour éviter les erreurs."""
+    
+    # Messages par défaut
+    DEFAULT_SYSTEM_MESSAGE = "Vous êtes KolaChatBot, un assistant IA serviable, créatif et honnête. Répondez en français."
+    DEFAULT_STARTER_MESSAGE = "Bonjour ! Je suis KolaChatBot. Comment puis-je vous aider aujourd'hui ? 🤖"
+
+    # Initialisation de base
+    if "selected_model_id" not in st.session_state:
+        st.session_state.selected_model_id = DEFAULT_MODEL_ID
+    if "system_message" not in st.session_state:
+        st.session_state.system_message = DEFAULT_SYSTEM_MESSAGE
+    if "starter_message" not in st.session_state:
+        st.session_state.starter_message = DEFAULT_STARTER_MESSAGE
+    if "chat_history" not in st.session_state:
+        st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message, "type": "text"}]
 
+    # Paramètres de génération
+    if "max_response_length" not in st.session_state:
+        st.session_state.max_response_length = 1024
+    if "temperature" not in st.session_state:
+        st.session_state.temperature = 0.7
+    if "top_p" not in st.session_state:
+        st.session_state.top_p = 0.95
+
+    # Fonctionnalité de recherche web
+    if "enable_web_search" not in st.session_state:
+        st.session_state.enable_web_search = False
+
+    # Stockage temporaire pour les résultats de recherche (pour affichage)
+    if 'last_search_results' not in st.session_state:
+        st.session_state.last_search_results = None
+
+# Appel de la fonction d'initialisation au début du script
+initialize_session_state()
 
 # -----------------------------------------------------------------------------
-# Helper pour formater les exports (Adapté pour images)
+# Fonctions d'export de la conversation
 # -----------------------------------------------------------------------------
 def format_history_to_txt(chat_history: list[dict]) -> str:
-    lines = [f"KolaChatBot Conversation (Export Texte Uniquement) - Exporté le {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"]
+    lines = [f"KolaChatBot Conversation - Exporté le {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"]
     for message in chat_history:
-        if message.get("role") == "system": continue
         role = "Utilisateur" if message["role"] == "user" else "KolaChatBot"
-        content_type = message.get("type", "text")
-        if content_type == "text" or content_type == "t2i_prompt": # Include user prompts for T2I in text export
-            lines.append(f"{role}:\n{message['content']}\n")
-        elif content_type == "t2i" and "prompt" in message: # For AI image response, include the prompt if available
-            lines.append(f"{role} (Image générée):\nPrompt: {message['prompt']}\n(Image non incluse dans l'export TXT)\n")
-        # Ignore other types gracefully
-
+        lines.append(f"--- {role} ---\n{message['content']}\n\n")
     return "".join(lines)
 
 def format_history_to_json(chat_history: list[dict]) -> str:
-    # Export JSON can include image data (base64 or placeholder) or just text.
-    # Storing raw bytes directly isn't JSON serializable. Base64 is possible but makes JSON huge.
-    # Let's export text content and image prompts, but not the image data itself.
-    export_history = []
-    for message in chat_history:
-        if message.get("role") == "system": continue
-        export_msg = {"role": message["role"], "type": message.get("type", "text")}
-
-        if export_msg["type"] == "text":
-            export_msg["content"] = message.get("content", "")
-        elif export_msg["type"] == "t2i_prompt": # User prompt for T2I (kept for handling potential old history)
-             export_msg["content"] = message.get("content", "")
-        elif export_msg["type"] == "t2i" and "prompt" in message: # AI image response (kept for handling potential old history)
-            export_msg["prompt"] = message["prompt"]
-            export_msg["image_placeholder"] = "(Image non incluse dans l'export JSON)" # Indicate image was here
-            if "content" in message: # Include error message if it was an error T2I response
-                 export_msg["error_content"] = message["content"]
-
-        export_history.append(export_msg)
-
-    return json.dumps(export_history, indent=2, ensure_ascii=False)
-
+    export_data = {
+        "export_date": datetime.now().isoformat(),
+        "conversation": chat_history
+    }
+    return json.dumps(export_data, indent=2, ensure_ascii=False)
 
 def format_history_to_md(chat_history: list[dict]) -> str:
     lines = [f"# KolaChatBot Conversation\n*Exporté le {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*\n\n"]
     for message in chat_history:
-        if message.get("role") == "system": continue
-        user_avatar = st.session_state.avatars.get('user', '👤')
-        assistant_avatar = st.session_state.avatars.get('assistant', '🤖')
-        role_label = f"{user_avatar} **Utilisateur**" if message["role"] == "user" else f"{assistant_avatar} **KolaChatBot**"
-
-        content_type = message.get("type", "text")
-
-        lines.append(f"### {role_label}\n\n")
-        if content_type == "text":
-            lines.append(f"{message.get('content', '')}\n\n")
-        elif content_type == "t2i_prompt": # Kept for handling potential old history
-             lines.append(f"*Prompt image:* {message.get('content', '')}\n\n")
-        elif content_type == "t2i": # Kept for handling potential old history
-            prompt_text = message.get('prompt', 'Pas de prompt enregistré')
-            error_text = message.get('content', '') # Potential error message
-            lines.append(f"*Image générée (prompt: {prompt_text})*\n")
-            if error_text:
-                 lines.append(f"*(Erreur de génération : {error_text})*\n")
-            # In actual markdown, you could embed a base64 image or link if saved.
-            # For this simple export, just note it.
-            lines.append("![Image non incluse dans l'export Markdown](placeholder.png)\n\n")
-
-        lines.append("---\n\n")
-
+        avatar = "👤" if message["role"] == "user" else "🤖"
+        role_label = "Utilisateur" if message["role"] == "user" else "KolaChatBot"
+        lines.append(f"### {avatar} {role_label}\n\n{message['content']}\n\n---\n\n")
     return "".join(lines)
 
 # -----------------------------------------------------------------------------
-# LLM API helper (Unified call logic)
+# Fonctions principales (Recherche Web et Appel API)
 # -----------------------------------------------------------------------------
-# Removed build_mistral_prompt as no Mistral models
-
-def call_hf_inference(model_id: str, payload_inputs: any, params: dict, model_type: str) -> any:
-    """
-    Calls the Hugging Face Inference API for either text generation or text-to-image.
-    This function is now only relevant for T2I calls if you were to add back T2I models.
-    """
-    if not HUGGINGFACEHUB_API_TOKEN:
-        return "Erreur d'API Hugging Face: Le token HUGGINGFACEHUB_API_TOKEN est introuvable."
-
-    headers = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}
-    url = f"{HF_BASE_API_URL}{model_id}"
-
-    # Since we removed HF text models, this branch is not expected to be hit in this config.
-    # If you add HF text models back, ensure the prompt building is done by the caller.
-    if model_type == 'text':
-         # This case should not be hit with the current AVAILABLE_MODELS list
-         st.error("Erreur interne: Modèle texte Hugging Face non pris en charge dans cette configuration.")
-         return {"role": "assistant", "content": "Erreur interne: Modèle texte Hugging Face non pris en charge dans cette configuration.", "type": "text"}
-
-
-    elif model_type == 't2i':
-        # This case should not be hit with the current AVAILABLE_MODELS list
-        st.error("Erreur interne: Modèle Text-to-Image Hugging Face non pris en charge dans cette configuration.")
-        return {"role": "assistant", "content": "Erreur interne: Modèle Text-to-Image Hugging Face non pris en charge dans cette configuration.", "type": "text", "prompt": payload_inputs}
-
-    else:
-        return f"Erreur interne: Type de modèle Hugging Face '{model_type}' inconnu."
-
-
-    # The code below this point would only be reached if you added back HF models.
-    # The structure is kept but not actively used in this config.
-    payload = {}
-    response_parser = None
-    response_is_json = False # Default to False for safety
-
-    # ... (rest of the hf inference logic is now effectively unreachable) ...
-    # To avoid linting errors or confusion, let's just make it return an error for any other type.
-    return f"Erreur interne: Type de modèle Hugging Face '{model_type}' non géré dans l'appel API."
-
-
-
-def call_google_api(model_id: str, system_message: str, chat_history_for_api: list[dict], params: dict) -> str:
-    """
-    Calls the Google Generative AI API (Text models only).
-    Note: The system_message is included for consistency but IGNORED in the API call.
-    """
-    if not GOOGLE_API_KEY:
-        return "Erreur d'API Google: La clé API Google est introuvable."
-
+def perform_web_search(query: str, num_results: int = 5) -> tuple[str, list]:
+    """Effectue une recherche web avec DuckDuckGo et retourne le contexte et les sources."""
+    st.session_state.last_search_results = None  # Réinitialiser à chaque recherche
     try:
-        genai.configure(api_key=GOOGLE_API_KEY)
+        with DDGS() as ddgs:
+            results = list(ddgs.text(keywords=query, region='fr-fr', max_results=num_results))
+            if not results:
+                return "Aucun résultat de recherche trouvé.", []
+
+            formatted_context = ""
+            source_details = []
+            for i, res in enumerate(results):
+                formatted_context += f"[Source {i+1}]\nTitre: {res.get('title', 'N/A')}\nExtrait: {res.get('body', 'N/A')}\nURL: {res.get('href', 'N/A')}\n\n"
+                source_details.append(res)
+            
+            st.session_state.last_search_results = source_details  # Sauvegarder pour affichage ultérieur
+            return formatted_context, source_details
     except Exception as e:
-        return f"Erreur de configuration Google API: {e}. Vérifiez votre GOOGLE_API_KEY."
+        return f"Erreur lors de la recherche web: {e}", []
 
-    try:
-        model = genai.GenerativeModel(model_id)
-
-        # Prepare history for the Gemini API
-        # The standard format is a list of dicts: [{'role': 'user', 'parts': [...]}, {'role': 'model', 'parts': ...]}]
-        # `chat_history_for_api` here is the list of messages relevant to the model turns (filtered by text type).
-        gemini_history_parts = []
+def get_gemini_response_stream(model_id: str, system_prompt: str, chat_history_for_api: list[dict], params: dict):
+    """Appelle l'API Google Gemini et retourne un générateur (stream) pour la réponse."""
+    if not GOOGLE_API_KEY:
+        yield "Erreur: La clé API Google n'est pas configurée. Veuillez l'ajouter pour continuer."
+        return
 
-        # System message is ignored for API call (as per previous fix)
+    try:
+        model = genai.GenerativeModel(
+            model_id,
+            system_instruction=system_prompt
+        )
 
+        gemini_history = []
         for msg in chat_history_for_api:
-            if msg.get("type", "text") == "text":
-                # Map roles: Streamlit 'user' -> Google 'user', Streamlit 'assistant' -> Google 'model'
-                role = 'user' if msg['role'] == 'user' else 'model'
-                gemini_history_parts.append({"role": role, "parts": [msg['content']]})
+            role = 'user' if msg['role'] == 'user' else 'model'
+            gemini_history.append({"role": role, "parts": [msg['content']]})
+        
+        current_prompt = gemini_history.pop()['parts'][0]
 
         generation_config = genai.types.GenerationConfig(
-            max_output_tokens=params.get("max_new_tokens", 200),
-            temperature=params.get("temperature", 0.6),
-            top_p=params.get("top_p", 0.9),
+            max_output_tokens=params.get("max_new_tokens"),
+            temperature=params.get("temperature"),
+            top_p=params.get("top_p"),
         )
-
-        # Call generate_content *without* the 'system_instruction' argument
-        response = model.generate_content(
-            contents=gemini_history_parts,
+        
+        response_stream = model.generate_content(
+            current_prompt,
             generation_config=generation_config,
-            request_options={'timeout': 180}
+            history=gemini_history,
+            stream=True
         )
 
-        # Process the response
-        if response.candidates:
-            if response.candidates[0].content and response.candidates[0].content.parts:
-                generated_text = "".join(part.text for part in response.candidates[0].content.parts)
-                return generated_text.strip()
-            elif hasattr(response.candidates[0], 'finish_reason') and response.candidates[0].finish_reason:
-                # Model finished generation early
-                return f"API Google: La génération s'est terminée prématurément. Raison: {response.candidates[0].finish_reason}. Prompt feedback: {response.prompt_feedback}"
-            else:
-                # Candidate exists but no content parts, might be an unhandled scenario
-                return f"Erreur API Google: Réponse vide ou inattendue après appel réussi. Debug info: {response.candidates[0]}"
-        elif hasattr(response, 'prompt_feedback') and response.prompt_feedback and response.prompt_feedback.block_reason:
-            # Prompt was blocked for safety reasons
-            return f"API Google: Votre message a été bloqué ({response.prompt_feedback.block_reason}). Raison détaillée: {response.prompt_feedback.safety_ratings}"
-        else:
-            # No candidates and no block feedback - unknown error
-            return f"API Google: Aucune réponse générée pour une raison inconnue. Debug info: {response}"
+        for chunk in response_stream:
+            if chunk.parts:
+                yield chunk.text
 
     except Exception as e:
-        # Catch any other exceptions during the API call
-        return f"Erreur lors de l'appel API Google: {e}"
-
-
-def _format_time(seconds):
-    """Formats seconds into a human-readable string."""
-    if not isinstance(seconds, (int, float)) or seconds < 0: return "N/A"
-    minutes = int(seconds // 60)
-    remaining_seconds = int(seconds % 60)
-    if minutes > 0: return f"{minutes} min {remaining_seconds} sec"
-    return f"{remaining_seconds} sec"
+        yield f"Erreur lors de l'appel à l'API Google: {e}"
 
 # -----------------------------------------------------------------------------
-# Web Search Function (for RAG)
+# Configuration de la page Streamlit
 # -----------------------------------------------------------------------------
-def perform_web_search(query: str, num_results: int = 3) -> str:
-    """Performs a web search using DuckDuckGo and returns formatted results."""
-    try:
-        # Use DDGS.text for simple text search (DDGS class)
-        # It's best practice to use DDGS as a context manager to ensure proper session closure
-        with DDGS() as ddgs:
-            # Set region for potentially better results
-            results = ddgs.text(keywords=query, max_results=num_results, region='fr-fr') # Added region
-
-            # Convert generator to list for iteration
-            results_list = list(results)
+st.set_page_config(page_title="KolaChatBot IA", page_icon="🤖", layout="wide")
 
-        if not results_list:
-            return "Aucun résultat de recherche trouvé."
-
-        formatted_results = []
-        for i, res in enumerate(results_list):
-            formatted_results.append(f"Titre: {res.get('title', 'N/A')}")
-            formatted_results.append(f"URL: {res.get('href', 'N/A')}")
-            formatted_results.append(f"Extrait: {res.get('body', 'N/A')}") # Removed extra newline
-        return "\n".join(formatted_results)
-    except Exception as e:
-        return f"Erreur lors de la recherche web: {e}"
+st.title("🤖 KolaChatBot IA")
+selected_model_info = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.selected_model_id), None)
+st.markdown(f"*Modèle actuel : **{selected_model_info['name']}***")
 
 # -----------------------------------------------------------------------------
-# Generation Functions (Simplified as task is fixed to Text)
-# -----------------------------------------------------------------------------
-
-def get_text_response(selected_model_id: str, system_prompt: str, full_chat_history: list[dict]):
-    """
-    Handles text generation request using the selected text model.
-    This is the primary generation function in this text-only configuration.
-    """
-    # Model type compatibility is checked by the caller before calling this function.
-    selected_model_info = next((m for m in AVAILABLE_MODELS if m['id'] == selected_model_id), None)
-    if not selected_model_info: # Should not happen due to caller check, but safety
-        return {"role": "assistant", "content": f"Erreur interne: Modèle texte '{selected_model_id}' introuvable.", "type": "text"}
-
-    model_provider = selected_model_info['provider']
-    model_id = selected_model_info['id']
-    # model_type is 'text' here as checked by the caller
-    params = { # Use parameters from session state for text generation
-        "max_new_tokens": st.session_state.max_response_length,
-        "temperature": st.session_state.temperature,
-        "top_p": st.session_state.top_p,
-    }
-
-    # Create a list of actual conversation turns relevant for text models.
-    # This excludes the initial assistant starter message if it's just the default one
-    # and there's more history. It should only include text messages.
-    actual_conversation_history_for_prompt = [
-        msg for msg in full_chat_history
-        if msg.get("type", "text") == "text" # Only include messages explicitly typed as 'text'
-        and not (msg['role'] == 'assistant' and msg['content'] == st.session_state.starter_message and len(full_chat_history) > 1)
-    ]
-
-    # For text models, the history for prompt building should end with the user's text message.
-    # The last message in the *full_chat_history* is guaranteed to be the user's latest input
-    # (which triggered this generation). We need to ensure that user input was *text type*.
-    if full_chat_history and full_chat_history[-1].get("type", "text") != "text":
-        # This indicates an issue with the calling logic or user input handling
-        st.error("Logique d'historique interne invalide: La dernière message utilisateur n'est pas du texte pour la génération de texte.")
-        return {"role": "assistant", "content": "Erreur interne: Le dernier message utilisateur n'est pas du texte pour la génération de texte.", "type": "text"}
-
-    response_content = "" # Initialize outside condition
-
-    # --- Web Search (RAG) Integration for Google models ---
-    if st.session_state.enable_web_search and model_provider == 'google':
-        last_user_query = full_chat_history[-1]['content'] # Get the user's last actual query
-        
-        # Display "Recherche en cours..." message
-        st.session_state.chat_history.append({"role": "assistant", "content": f"*(KolaChatBot recherche sur le web pour '{last_user_query}'...)*", "type": "text"})
-        st.rerun() # Force a rerun to display this message immediately
-
-        # This part runs on the *next* rerun
-        try:
-             # Perform the actual web search
-             search_results = perform_web_search(last_user_query)
-             
-             # Create a copy of the history for the API call and inject the search results
-             temp_api_history = []
-             for msg in actual_conversation_history_for_prompt:
-                  temp_api_history.append(msg.copy()) # Make sure to copy the dict
-
-             # Prepend search results to the *last user message* in the temp history
-             context_message_content = f"Voici des informations pertinentes du web sur le sujet de la dernière question :\n```\n{search_results}\n```\n\nEn te basant sur ces informations (si pertinentes) et notre conversation, réponds à la question : "
-
-             if temp_api_history and temp_api_history[-1]['role'] == 'user' and temp_api_history[-1]['type'] == 'text':
-                  temp_api_history[-1]['content'] = context_message_content + temp_api_history[-1]['content']
-             else:
-                  # Fallback: add context as a separate user message if last message isn't a simple user text
-                  temp_api_history.append({"role": "user", "content": context_message_content, "type": "text"})
-
-             # Call the Google API with the modified history
-             response_content = call_google_api(model_id, system_prompt, temp_api_history, params)
-             
-             # Update the search message in chat history to indicate results were used
-             # Find the message we added earlier and update its content
-             if st.session_state.chat_history and st.session_state.chat_history[-1]["content"].startswith("*(KolaChatBot recherche web"):
-                  st.session_state.chat_history[-1]["content"] = f"*(KolaChatBot a terminé la recherche web. Résultats intégrés à la réponse.)*"
-
-
-        except Exception as e:
-             response_content = f"Erreur lors de la recherche web ou de l'intégration des résultats : {e}"
-             # If search failed, update the temporary message to show the error
-             if st.session_state.chat_history and st.session_state.chat_history[-1]["content"].startswith("*(KolaChatBot recherche web"):
-                 st.session_state.chat_history[-1]["content"] = f"*(KolaChatBot a rencontré une erreur lors de la recherche web : {e})*"
-
-
-    else: # No web search or not a Google model
-        if model_provider == 'google':
-            # This path is for Google models WITHOUT web search enabled
-            response_content = call_google_api(model_id, system_prompt, actual_conversation_history_for_prompt, params)
-        else:
-             # This branch is effectively unreachable with the current AVAILABLE_MODELS list (only Google Text models)
-             response_content = f"Erreur interne: Fournisseur API '{model_provider}' inconnu pour le modèle texte '{model_id}'."
-
-
-    return {"role": "assistant", "content": response_content, "type": "text"}
-
-
-# This function is no longer used in this text-only configuration.
-def get_image_response(selected_model_id: str, user_prompt: str):
-     return {"role": "assistant", "content": "La génération d'images n'est pas disponible dans cette version de l'application.", "type": "text", "prompt": user_prompt}
-
-
-# -----------------------------------------------------------------------------
-# Streamlit page configuration
-# -----------------------------------------------------------------------------
-st.set_page_config(page_title="KolaChatBot - Multi-Modèles IA", page_icon="🤖", layout="wide")
-st.title("🤖 KolaChatBot - Multi-Modèles IA")
-
-# Dynamically display the currently selected model and task in the description
-# Get the info for the *currently selected* model to display its name
-selected_model_info_display = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.get('selected_model_id')), None)
-selected_model_name_display = selected_model_info_display['name'] if selected_model_info_display else 'Aucun modèle sélectionné' # Fallback text
-
-
-# Current task is now fixed to 'text'
-current_task_label = "Génération de Texte (Conversation)"
-st.markdown(f"*Tâche actuelle : **{current_task_label}***\n*Modèle : **`{selected_model_name_display}`** *")
-
-
-# -----------------------------------------------------------------------------
-# Manuel d'utilisation (Update for text-only and RAG)
-# -----------------------------------------------------------------------------
-with st.expander("📖 Manuel d'utilisation de KolaChatBot", expanded=False):
-    st.markdown("""
-    Bienvenue sur KolaChatBot - Une application de chat IA ! Voici comment tirer le meilleur parti de notre assistant IA :
-
-    **1. Comment interagir ?**
-    - **La tâche est fixée sur Texte (conversation).**
-    - **Entrer votre prompt :** Tapez votre message dans la zone de texte en bas et appuyez sur Entrée.
-
-    **2. Paramètres dans la barre latérale (Sidebar) :**
-    Ajustez les paramètres pour la génération de texte.
-
-    *   **Sélection du Modèle :** Choisissez un modèle Google Gemini Flash pour converser.
-        - **Important :** Assurez-vous que la clé API Google (`GOOGLE_API_KEY`) est configurée dans vos secrets.
-        - Changer de modèle **ne réinitialise pas** automatiquement les autres paramètres. Les paramètres actuels seront utilisés lors de la prochaine génération.
-    *   **Option "Activer la recherche web" :** Activez cette option pour que l'IA (via les modèles Google Gemini) recherche sur le web et utilise les résultats pour enrichir ses réponses.
-    *   **Paramètres de Génération :** Ajustez la longueur maximale de la réponse, la température (créativité), et le Top-P (sampling).
-    *   **Message Système / Personnalité :** Définissez le rôle ou le style que l'IA de texte doit adopter.
-    *   **Message de Bienvenue de l'IA :** Personnalisez le message initial.
-    *   **Sélection des Avatars.**
-    *   **Gestion de la Conversation :** Utilisez les boutons pour appliquer les paramètres actuels (et démarrer une nouvelle conversation) ou simplement effacer l'historique actuel.
-    *   **Exporter la Conversation :** Téléchargez l'historique.
-
-    **3. Limitations :**
-    - **L'application est configurée uniquement pour la tâche de génération de texte.** Seuls les modèles Google Gemini Flash sont disponibles. Tenter d'utiliser un autre type de modèle entraînerait une erreur.
-    - La fonction "Recherche Web" est actuellement implémentée uniquement pour les modèles Google Gemini et peut entraîner des délais supplémentaires. Les résultats de recherche web s'affichent temporairement dans le chat avant la réponse de l'IA.
-    - Notez que l'application de votre **Message Système / Personnalité** aux modèles Google dépend de la prise en charge de la fonctionnalité `system_instruction` par la version de la bibliothèque Google Generative AI et le modèle spécifique utilisé. Si vous obtenez des erreurs (`unexpected keyword argument 'system_instruction'`), le message système sera ignoré pour ces modèles.
-    - Cette application ne stocke pas les données de manière permanente. L'historique est présent dans la session tant que l'application tourne.
-
-    Amusez-vous bien avec KolaChatBot !
-    """)
-# -----------------------------------------------------------------------------
-# Sidebar settings
+# Barre latérale (Sidebar)
 # -----------------------------------------------------------------------------
 with st.sidebar:
-    st.header("🛠��� Configuration de KolaChatBot")
-
-    # Task is fixed to 'text'
-    st.subheader("🎯 Tâche IA : Texte (fixée)")
+    st.header("🛠️ Configuration")
 
     st.subheader("🧠 Sélection du Modèle")
-    all_available_models = AVAILABLE_MODELS
-
-    if not all_available_models:
-        st.warning("Aucun modèle disponible défini dans la liste AVAILABLE_MODELS.")
-        st.session_state.selected_model_id = None
-    else:
-        model_options = {model['id']: model['name'] for model in all_available_models}
-
-        current_model_index = 0
-        if st.session_state.selected_model_id in model_options:
-            current_model_index = list(model_options.keys()).index(st.session_state.selected_model_id)
-        elif all_available_models:
-             default_model_from_list = next((m for m in AVAILABLE_MODELS), None)
-             if default_model_from_list:
-                 st.session_state.selected_model_id = default_model_from_list['id']
-                 current_model_index = 0
-             else:
-                 st.session_state.selected_model_id = None
-
-
-        st.session_state._prev_model_id_before_selectbox = st.session_state.selected_model_id
-        selected_model_id_from_selectbox_value = st.selectbox(
-            "Choisir le modèle :",
-            options=list(model_options.keys()),
-            format_func=lambda x: model_options[x],
-            index=current_model_index,
-            key="selected_model_id",
-            help="Sélectionnez un modèle disponible (uniquement modèles texte pris en charge dans cette version)."
-        )
-
-
-    # --- Model Change Detection (without auto-resetting params) ---
-    model_id_changed_this_run = st.session_state.selected_model_id != st.session_state.get('_prev_model_id_before_selectbox')
-
-    # If model changed (and not triggered by an explicit button reset), reset history
-    if model_id_changed_this_run and not st.session_state.get('_reset_triggered', False):
+    model_options = {model['id']: model['name'] for model in AVAILABLE_MODELS}
+    
+    def on_model_change():
+        """Callback pour réinitialiser la conversation lors du changement de modèle."""
         st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message, "type": "text"}]
-        st.info(f"⚠️ Modèle changé. La conversation a été réinitialisée. Les paramètres restent les mêmes jusqu'à application explicite.")
-        # Do not set _reset_triggered here
-        st.rerun()
-
-
-    # Check required API Key for the selected model
-    if st.session_state.selected_model_id:
-        current_model_info_check = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.selected_model_id), None)
-        if current_model_info_check:
-            current_model_provider = current_model_info_check['provider']
-            if current_model_provider == 'huggingface': # Should not happen with current models
-                if not HUGGINGFACEHUB_API_TOKEN:
-                    st.warning("❌ Le token Hugging Face est manquant (`HUGGINGFACEHUB_API_TOKEN`). Les modèles Hugging Face ne fonctionneront pas.")
-            elif current_model_provider == 'google':
-                if not GOOGLE_API_KEY:
-                    st.warning("❌ La clé API Google est manquante (`GOOGLE_API_KEY`). Les modèles Google ne fonctionneront pas.")
-
-
-    # --- Dynamic Parameter Settings based on the SELECTED MODEL's TYPE ---
-    st.subheader("⚙️ Paramètres")
-    current_selected_model_info = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.selected_model_id), None)
-
-    if st.session_state.selected_model_id and current_selected_model_info:
-        expander_title = f"Ajuster Paramètres : {current_selected_model_info['name']}" # Simplified title
-        with st.expander(expander_title, expanded=True):
-            # Since only text models are available, only display text parameters
-            if current_selected_model_info['type'] == 'text':
-                 st.session_state.max_response_length = st.number_input(
-                     "Max New Tokens (longueur max réponse) :",
-                     min_value=20, max_value=8192, # Increased max_value
-                     value=st.session_state.max_response_length,
-                     step=10,
-                     key="max_new_tokens_input",
-                     help="Longueur maximale de la réponse de l'IA (en jetons ou tokens).",
-                 )
-                 st.session_state.temperature = st.slider(
-                     "Temperature (créativité) :",
-                     min_value=0.0, max_value=2.0,
-                     value=st.session_state.temperature,
-                     step=0.01,
-                     key="temperature_input",
-                     help="Contrôle le caractère aléatoire des réponses. Plus élevé = plus créatif/imprévisible.",
-                 )
-                 st.session_state.top_p = st.slider(
-                     "Top-P (sampling) :",
-                     min_value=0.01, max_value=1.0,
-                     value=st.session_state.top_p,
-                     step=0.01,
-                     key="top_p_input",
-                     help="Contrôle la diversité en limitant les options de tokens. Plus bas = moins diversifié. 1.0 = désactivé.",
-                 )
-
-                 st.session_state.system_message = st.text_area(
-                     "Message Système / Personnalité :",
-                     value=st.session_state.system_message,
-                     height=100,
-                     key="system_message_input",
-                     help="Décrivez le rôle ou le style que l'IA de texte doit adopter. Notez que cette consigne peut être ignorée par certains modèles/versions API (voir Limitations).",
-                 )
-                 st.session_state.starter_message = st.text_area(
-                     "Message de Bienvenue de l'IA :",
-                     value=st.session_state.starter_message,
-                     height=100,
-                     key="starter_message_input",
-                     help="Le premier message que l'IA affichera au début d'une nouvelle conversation textuelle ou après un reset.",
-                 )
-                 # Web Search Checkbox - only displayed for Google models
-                 if current_selected_model_info['provider'] == 'google':
-                     st.session_state.enable_web_search = st.checkbox(
-                         "Activer la recherche web (pour les modèles Google)",
-                         value=st.session_state.enable_web_search,
-                         key="web_search_checkbox",
-                         help="Si coché, les modèles Google Gemini effectueront une recherche DuckDuckGo et utiliseront les résultats pour répondre à votre prompt."
-                     )
-                 else: # If a non-Google model were added back
-                      # Ensure checkbox state is not changed and it's disabled
-                      # Use a unique key based on model_id to prevent state interference if model changes
-                      st.session_state.enable_web_search = st.checkbox(
-                         "Activer la recherche web (pour les modèles Google)",
-                         value=False, # Always False if not Google
-                         key=f"web_search_checkbox_{current_selected_model_info.get('id', 'nongoogle')}", # Unique key
-                         disabled=True,
-                         help="La recherche web n'est disponible que pour les modèles Google Gemini.",
-                     )
-
-            elif current_selected_model_info['type'] == 't2i': # This branch should not be hit
-                 st.info("Note: Ce modèle est de type 'Image', mais l'application ne supporte que la tâche 'Texte'.")
-
-        if current_selected_model_info and current_selected_model_info['type'] not in ['text', 't2i']: # Should not be hit
-             st.info(f"Type de modèle sélectionné ('{current_selected_model_info['type']}') inconnu ou non pris en charge pour l'affichage des paramètres.")
-
-    else:
-         st.info("Aucun modèle sélectionné, les paramètres ne sont pas disponibles.")
+        st.toast(f"Modèle changé. Conversation réinitialisée.")
+
+    st.selectbox(
+        "Choisir le modèle :",
+        options=list(model_options.keys()),
+        format_func=lambda x: model_options[x],
+        key="selected_model_id",
+        on_change=on_model_change,
+        help="Changer de modèle démarre une nouvelle conversation."
+    )
+    
+    if not GOOGLE_API_KEY:
+        st.error("❌ Clé API Google manquante.")
 
+    st.subheader("⚙️ Paramètres de Génération")
+    with st.expander("Ajuster les paramètres", expanded=False):
+        st.slider("Max Tokens", 128, 8192, key="max_response_length", step=128)
+        st.slider("Température", 0.0, 2.0, key="temperature", step=0.05)
+        st.slider("Top-P", 0.0, 1.0, key="top_p", step=0.05)
 
-    # Flag used to detect if 'Appliquer Paramètres' or 'Effacer' was clicked
-    # and prevent regeneration on the immediate rerun
-    if st.session_state.get('_reset_triggered', False):
-         st.session_state._reset_triggered = False
+    st.subheader("👤 Personnalisation")
+    st.text_area("Message Système / Personnalité", height=100, key="system_message")
+    st.text_area("Message de bienvenue", height=100, key="starter_message")
 
+    st.subheader("🌐 Recherche Web (RAG)")
+    st.checkbox("Activer la recherche web", key="enable_web_search")
 
-    st.subheader("🔄 Gestion de la Conversation")
-    # Reset button that applies *all* parameters (system/starter/model) and starts new conversation
-    if st.button("♻️ Appliquer Paramètres & Nouvelle Conversation", type="primary", help="Applique les paramètres actuels et démarre une nouvelle conversation texte en effaçant l'historique."):
-        # Parameters are already updated in session state via key linking
-        # Reset history.
+    st.subheader("🔄 Gestion")
+    col1, col2 = st.columns(2)
+    if col1.button("♻️ Nouvelle Conv.", use_container_width=True, help="Démarre une nouvelle conversation."):
         st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message, "type": "text"}]
-        st.info("Paramètres appliqués et conversation réinitialisée.")
-        st.session_state._reset_triggered = True # Set flag
-        st.rerun() # Rerun immediately to refresh UI
-
-    # Clear history button - simpler, just history
-    if st.button("🗑️ Effacer la Conversation Actuelle", help="Efface l'historique de conversation mais conserve les paramètres actuels."):
-        # Reset history.
+        st.toast("Nouvelle conversation démarrée.")
+        st.rerun()
+    if col2.button("🗑️ Effacer", type="primary", use_container_width=True, help="Efface l'historique de la conversation."):
         st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message, "type": "text"}]
-        st.info("Conversation actuelle effacée.")
-        st.session_state._reset_triggered = True # Set flag
-        st.rerun() # Rerun immediately to show empty chat with starter message
-
-
-    st.subheader("📥 Exporter la Conversation")
-    # Check if there's actual history beyond the initial assistant message
-    export_possible = len([msg for msg in st.session_state.chat_history if msg.get("role") != "system"]) > 1
+        st.toast("Conversation effacée.")
+        st.rerun()
 
-    if "chat_history" in st.session_state and export_possible :
+    st.subheader("📥 Exporter")
+    if len(st.session_state.chat_history) > 1:
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-
-        col_export1, col_export2, col_export3 = st.columns(3)
-        with col_export1:
-            st.download_button(
-                label="TXT",
-                data=format_history_to_txt(st.session_state.chat_history),
-                file_name=f"kolachat_conversation_{timestamp}.txt",
-                mime="text/plain",
-                use_container_width=True
-            )
-        with col_export2:
-            st.download_button(
-                label="JSON",
-                data=format_history_to_json(st.session_state.chat_history),
-                file_name=f"kolachat_conversation_{timestamp}.json",
-                mime="application/json",
-                use_container_width=True
-            )
-        with col_export3:
-            st.download_button(
-                label="Markdown",
-                data=format_history_to_md(st.session_state.chat_history),
-                file_name=f"kolachat_conversation_{timestamp}.md",
-                mime="text/markdown",
-                use_container_width=True
-            )
+        st.download_button("TXT", format_history_to_txt(st.session_state.chat_history), f"kolachat_{timestamp}.txt")
+        st.download_button("JSON", format_history_to_json(st.session_state.chat_history), f"kolachat_{timestamp}.json")
+        st.download_button("Markdown", format_history_to_md(st.session_state.chat_history), f"kolachat_{timestamp}.md")
     else:
-        st.caption("Aucune conversation substantielle ou exportable pour le moment.")
-
+        st.caption("La conversation est vide.")
+        
     st.divider()
-    st.subheader("Crédits")
     st.markdown("""
-    **Auteur :** Sidoine K. YEBADOKPO
-    *Expert en Analyse de Données*
-
-    📧 [syebadokpo@gmail.com](mailto:syebadokpo@gmail.com)
+    **Auteur :** Sidoine K. YEBADOKPO  
+    *Expert en Analyse de Données*  
+    📧 syebadokpo@gmail.com  
     📞 +229 96 91 13 46
     """)
 
-
 # -----------------------------------------------------------------------------
-# Main Chat / Input / Generation Logic
+# Interface de Chat Principale
 # -----------------------------------------------------------------------------
-# Use a container for the chat history display area with a fixed height and border
-chat_interface_container = st.container(height=600, border=True)
-
-with chat_interface_container:
-    # Embed a div with a specific ID for the chat history messages
-    st.markdown("<div id='chat_history_container_scroll_area'>", unsafe_allow_html=True)
-
-    # Display existing messages from chat_history.
-    for message in st.session_state.chat_history:
-        if message["role"] == "system": continue
-
-        avatar_type = st.session_state.avatars.get(message["role"], "❓")
-
-        with st.chat_message(message["role"], avatar=avatar_type):
-            message_type = message.get("type", "text")
-            if message_type == "text":
-                st.markdown(message.get("content", ""))
-            elif message_type == "t2i_prompt": # Should not happen with current config
-                st.markdown(f"Prompt image : *{message.get('content', 'Prompt vide')}*")
-            elif message_type == "t2i" and ("image_data" in message or "content" in message): # Should not happen with current config
-                 st.markdown(f"Prompt: *{message.get('prompt', 'Pas de prompt enregistré')}*")
-                 if "image_data" in message:
-                      try: st.image(Image.open(BytesIO(message["image_data"])), caption="Image générée")
-                      except Exception as e: st.error(f"Erreur affichage image : {e}")
-                 elif message.get('content'): st.error(f"Échec génération image: {message.get('content')}")
-                 else: st.warning("Réponse image format inattendu.")
-
-
-    # Close the div for the scroll area
-    st.markdown("</div>", unsafe_allow_html=True)
-
-    # JavaScript to scroll to the bottom
-    # Needs to be inside the container's 'with' block or just below it.
-    # Running it after all messages are rendered ensures it scrolls to the *new* last message.
-    st.markdown(
-        """
-        <script>
-        var chatHistory = document.getElementById('chat_history_container_scroll_area');
-        if (chatHistory) {
-            chatHistory.scrollTop = chatHistory.scrollHeight;
-        }
-        </script>
-        """,
-        unsafe_allow_html=True,
-    )
-
-
-# Area for the chat input box
-input_container = st.container()
-with input_container:
-    if st.session_state.selected_model_id is not None:
-        current_selected_model_info = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.selected_model_id), None)
-        placeholder_text = f"Entrez votre message pour la conversation texte..." # Fixed placeholder text as task is text
-        user_input = st.chat_input(placeholder=placeholder_text, disabled=False)
-    else:
-         user_input = st.chat_input(placeholder=f"Veuillez configurer un modèle disponible pour démarrer...", disabled=True)
-         st.info(f"Veuillez sélectionner un modèle disponible dans la barre latérale pour commencer. Il n'y a actuellement aucun modèle configuré ou sélectionné.")
-
-
-# --- Logic to handle user input and trigger generation ---
-# This part runs only when user_input is NOT None (i.e., after the user hits Enter)
-# AND a model is selected AND the reset flag is not set (from buttons)
-if user_input and st.session_state.selected_model_id and not st.session_state.get('_reset_triggered', False):
-    # 1. Append the user's message to the chat history
-    user_message_entry = {"role": "user", "content": user_input, "type": "text"} # Always text input now
-    st.session_state.chat_history.append(user_message_entry)
-
-    # --- Trigger a rerun to immediately show the user message ---
-    # This restarts the script, redrawing the chat (now including the user's message).
-    # The code below the chat_input block will then execute on this rerun.
-    st.rerun()
-
-
-# --- This block executes on reruns, check if processing is needed ---
-# Processing is needed if the last message was from the user AND we haven't added an assistant response yet,
-# AND a model is selected, AND a reset was NOT just triggered by a button.
-if st.session_state.chat_history and st.session_state.chat_history[-1]["role"] == "user" and st.session_state.selected_model_id and not st.session_state.get('_reset_triggered', False):
-
-    # Find info about the currently selected model
-    selected_model_info = next((m for m in AVAILABLE_MODELS if m['id'] == st.session_state.selected_model_id), None)
+# Affichage de l'historique des messages
+for message in st.session_state.chat_history:
+    avatar = "👤" if message["role"] == "user" else "🤖"
+    with st.chat_message(message["role"], avatar=avatar):
+        st.markdown(message["content"])
+        # Afficher les sources si elles sont attachées au message de l'assistant
+        if message.get("sources"):
+            with st.expander("Sources web consultées", expanded=False):
+                for i, source in enumerate(message["sources"]):
+                    st.markdown(f"**{i+1}. {source.get('title', 'Titre inconnu')}**\n"
+                                f"[*Source*]({source.get('href', '#')})\n"
+                                f"> {source.get('body', 'Aucun extrait disponible.')}\n---")
+
+# Logique de traitement de l'entrée utilisateur
+if prompt := st.chat_input(f"Envoyer un message...", disabled=not GOOGLE_API_KEY):
+    # Ajouter et afficher le message de l'utilisateur
+    st.session_state.chat_history.append({"role": "user", "content": prompt, "type": "text"})
+    with st.chat_message("user", avatar="👤"):
+        st.markdown(prompt)
+
+    # Préparer et afficher la réponse de l'assistant
+    with st.chat_message("assistant", avatar="🤖"):
+        history_for_api = st.session_state.chat_history.copy()
+        
+        # Logique de Recherche Web (RAG)
+        if st.session_state.enable_web_search:
+            with st.spinner("KolaChatBot recherche sur le web..."):
+                search_context, sources = perform_web_search(prompt)
+            
+            if sources:
+                rag_prompt = (
+                    "En te basant STRICTEMENT sur les informations suivantes issues d'une recherche web, réponds à la question de l'utilisateur. "
+                    "Cite tes sources en utilisant le format [Source X] après chaque phrase concernée.\n\n"
+                    f"--- CONTEXTE DE LA RECHERCHE ---\n{search_context}\n--- FIN DU CONTEXTE ---\n\n"
+                    f"Question de l'utilisateur : {prompt}"
+                )
+                history_for_api[-1]['content'] = rag_prompt
+            else:
+                 st.toast("La recherche web n'a pas fourni de résultats.")
 
-    if not selected_model_info:
-        error_message = {"role": "assistant", "content": f"Erreur: Modèle sélectionné '{st.session_state.selected_model_id}' introuvable.", "type": "text"}
-        st.session_state.chat_history.append(error_message)
-        st.rerun()
-    else:
-        current_model_display_name = selected_model_info['name']
-        current_model_type = selected_model_info['type'] # Get type from the *selected model* info
-        last_user_message = st.session_state.chat_history[-1]
-        last_user_message_type = last_user_message.get("type", "text")
-        last_user_message_content = last_user_message.get("content", "")
-
-
-        # === Task/Model Compatibility Check (Simplified as task is fixed to 'text') ===
-        # Check if the selected model is of type 'text'
-        if current_model_type != 'text':
-            # If the model type is not text, display compatibility error
-            error_message_content = f"Erreur de compatibilité : L'application ne supporte que les modèles de texte, mais le modèle sélectionné ('{current_model_display_name}' - Type: '{current_model_type}') n'est pas un modèle de texte.\nVeuillez sélectionner un modèle de texte pour continuer la conversation."
-            error_message_entry = {"role": "assistant", "content": error_message_content, "type": "text"}
-            st.session_state.chat_history.append(error_message_entry)
-            st.rerun()
-        else:
-            # === Task and Model are Compatible (Text task, Text model) - Proceed with Generation ===
-            assistant_response_entry = None
-
-            with chat_interface_container: # Place spinner within the chat container where response will appear
-                assistant_avatar = st.session_state.avatars.get("assistant", "❓")
-                with st.chat_message("assistant", avatar=assistant_avatar):
-                    # The spinner message is shown here, wrapping the API call
-                    with st.spinner(f"KolaChatBot utilise {current_model_display_name} pour générer... 🤔"):
-                        # Call get_text_response (which now handles RAG internally for Google)
-                        assistant_response_entry = get_text_response(
-                            selected_model_id=st.session_state.selected_model_id,
-                            system_prompt=st.session_state.system_message,
-                            full_chat_history=st.session_state.chat_history # Pass full history
-                        )
-                    # Display the content from the response entry object returned by get_text_response
-                    # This happens *after* the spinner block
-                    if assistant_response_entry and assistant_response_entry.get("type") == "text":
-                        st.markdown(assistant_response_entry.get("content", "Erreur: Réponse texte vide."))
-                    elif assistant_response_entry: # Handle errors returned as text from the generation function
-                        st.error(f"Erreur lors de la génération: {assistant_response_entry.get('content', 'Raison inconnue.')}")
-                    else: # Handle None response from get_text_response
-                        st.error("La fonction de génération de texte n'a pas renvoyé de réponse valide.")
-
-
-            # --- Append the generated response entry to chat history ---
-            if assistant_response_entry is not None:
-                # Check if the last message added was the search message placeholder.
-                # If so, replace it with the actual response. Otherwise, just append.
-                if st.session_state.chat_history[-1]["role"] == "assistant" and \
-                   st.session_state.chat_history[-1]["content"].startswith("*(KolaChatBot recherche web"):
-                     # Remove the temporary search message
-                     st.session_state.chat_history.pop()
-                     # Append the actual response
-                     st.session_state.chat_history.append(assistant_response_entry)
-                else:
-                     # Just append the response if no search message was pending
-                     st.session_state.chat_history.append(assistant_response_entry)
-
-                # No need for another st.rerun() here, the UI is already updated by the first rerun.
-
-
-# --- Clear reset flag at the very end of the script if it was set ---
-# This ensures the flag is only True for one full rerun cycle after a button click
-if st.session_state.get('_reset_triggered', False):
-     st.session_state._reset_triggered = False
+        # Appel API et Affichage en Streaming
+        params = {
+            "max_new_tokens": st.session_state.max_response_length,
+            "temperature": st.session_state.temperature,
+            "top_p": st.session_state.top_p,
+        }
+        
+        # Utiliser st.write_stream pour afficher la réponse mot par mot
+        response_content = st.write_stream(get_gemini_response_stream(
+            st.session_state.selected_model_id,
+            st.session_state.system_message,
+            history_for_api,
+            params
+        ))
+
+    # Sauvegarder la réponse complète de l'assistant dans l'historique
+    assistant_message = {"role": "assistant", "content": response_content, "type": "text"}
+    if st.session_state.get('last_search_results'):
+        assistant_message["sources"] = st.session_state.last_search_results
+        st.session_state.last_search_results = None  # Nettoyer
+
+    st.session_state.chat_history.append(assistant_message)
+    # Un seul `rerun` est nécessaire si les sources ont été utilisées pour les réafficher correctement.
+    if "sources" in assistant_message:
+        st.rerun()