diff --git a/modules/ui_components.py b/modules/ui_components.py
index 883d2d5..aa9776a 100644
--- a/modules/ui_components.py
+++ b/modules/ui_components.py
@@ -17,7 +17,8 @@ def get_openrouter_stats():
 def render_user_msg(text):
     return Div(P(text, style="margin:0; font-weight:bold; white-space: pre-wrap;"), cls="user-bubble")
 
-def render_ai_msg(text, intention, nom_fournisseur, nom_modele, couleur, execution_time=None, trace=None):
+# --- LA FONCTION CORRIGÉE AVEC SES 10 ARGUMENTS ---
+def render_ai_msg(text, intention, nom_fournisseur, nom_modele, couleur, execution_time=None, trace=None, prompt_associe="", thematique="general", msg_id=""):
     time_badge = Span(f"⏱️ {execution_time}s", style="margin-right: 10px; color:#aaa; font-size:0.8em;") if execution_time else ""
     
     trace_html = ""
@@ -32,7 +33,25 @@ def render_ai_msg(text, intention, nom_fournisseur, nom_modele, couleur, executi
         time_badge, style="border-bottom: 1px solid #333; padding-bottom: 10px; margin-bottom: 10px; display: flex; flex-wrap: wrap; gap: 5px;"
     )
     
+    # --- BOUTON FORCER RAG (HTMX) ---
+    bouton_forcer = ""
+    if msg_id and prompt_associe and thematique != "aucun":
+        safe_prompt = prompt_associe.replace('"', '\\"').replace('\n', ' ')
+        safe_reponse = text.replace('"', '\\"').replace('\n', ' ')
+        valeurs_htmx = f'{{"prompt": "{safe_prompt}", "reponse": "{safe_reponse}", "theme": "{thematique}"}}'
+        
+        bouton_forcer = Button("💾 Mémoriser", 
+            hx_post="/api/forcer_rag", 
+            hx_vals=valeurs_htmx, 
+            hx_target=f"#badge-rag-{msg_id}", 
+            hx_swap="outerHTML", 
+            id=f"badge-rag-{msg_id}", 
+            cls="msg-action-btn", 
+            style="color: #4caf50; border-color: #4caf50;"
+        )
+
     actions = Div(
+        bouton_forcer,
         Button("📄 PDF", onclick="downloadResponsePDF(this)", cls="msg-action-btn", title="Télécharger en PDF"),
         Button("📘 DOC", onclick="downloadResponseDOC(this)", cls="msg-action-btn", title="Télécharger en Word"),
         Button("📝 TXT", onclick="downloadResponseTXT(this)", cls="msg-action-btn", title="Télécharger en Texte"),
@@ -42,9 +61,6 @@ def render_ai_msg(text, intention, nom_fournisseur, nom_modele, couleur, executi
         cls="msg-actions"
     )
     
-    # Si le texte contient une image générée (détection basique Markdown ou HTML)
-    # L'affichage natif s'occupera du rendu via marked.js
-    
     return Div(
         trace_html,
         encart_info,
@@ -79,7 +95,10 @@ def render_main_page(chat_id, utilisateur_nom="xavier", role="Super Admin"):
                     msg.get('modele', ''), 
                     msg.get('couleur', '#333'), 
                     msg.get('execution_time', ''),
-                    msg.get('trace', [])
+                    msg.get('trace', []),
+                    msg.get('prompt_associe', ''),
+                    msg.get('thematique', 'general'),
+                    msg.get('msg_id', '')
                 ))
     if messages_ui: messages_ui.append(Script("setTimeout(renderMarkdown, 150);"))
 
@@ -150,7 +169,6 @@ def render_main_page(chat_id, utilisateur_nom="xavier", role="Super Admin"):
                 Form(
                     Input(type="hidden", name="chat_id", value=chat_id),
                     Div(
-                        # --- NOUVEAU SÉLECTEUR DE MODE MÉDIA ---
                         Div(
                             Label(Input(type="radio", name="media_mode", value="texte", checked=True, onchange="checkMediaPaymentWarning()"), "📝 Texte", cls="media-btn"),
                             Label(Input(type="radio", name="media_mode", value="image", onchange="checkMediaPaymentWarning()"), "🎨 Image", cls="media-btn"),