better history format

app.py

@@ -15,15 +15,22 @@ firebase_admin.initialize_app(cred)
 db = firestore.client()
 
 
-def save_message(user_id, message):
+def save_message(user_id, message, role="user"):
     # Enregistrer le message dans Firestore
     doc_ref = db.collection("conversations").document(user_id)
     doc_ref.set(
         {
             "messages": firestore.ArrayUnion(
-                [{"message": message, "timestamp": datetime.now()}]
+                [
+                    {
+                        "role": role,
+                        "message": f"[{role.upper()}] {message} [/{role.upper()}]",
+                        "timestamp": datetime.now(),
+                    }
+                ]
             )
-        }
+        },
+        merge=True,
     )
 
 
@@ -31,8 +38,8 @@ def load_history(user_id):
     # Charger l'historique de conversation de l'utilisateur
     doc_ref = db.collection("conversations").document(user_id)
     doc = doc_ref.get()
-    if doc.exists:
-        messages = doc.to_dict()["messages"]
+    if doc:
+        messages = doc.to_dict().get("messages", [])
         return messages
     else:
         return []
@@ -93,17 +100,13 @@ def format_prompt(message, history, user_id=""):
 
     prompt = "<s>"
 
-    for user_prompt, bot_response in history:
-        prompt += f"[USER] {user_prompt} [/USER]"
-        if (
-            history
-            and isinstance(history[-1], dict)
-            and user_prompt == history[-1].get("0")
-        ):
-            # N'ajouter les balises [BOT] qu'autour de la dernière réponse du bot
-            prompt += f" [BOT]{filter_sensitive_info(bot_response)}[/BOT] "
-        else:
-            prompt += f" {filter_sensitive_info(bot_response)} "
+    for item in history:
+        role = item.get("role")
+        content = item.get("message")
+        if role == "user":
+            prompt += f"[USER] {content} [/USER]"
+        elif role == "bot":
+            prompt += f" [BOT]{filter_sensitive_info(content)}[/BOT] "
     prompt += f"</s>[INST][USER] {message} [/USER][/INST]"
 
     # Ajouter le contexte statique et dynamique à chaque fois que l'utilisateur pose une question
@@ -163,19 +166,17 @@ def generate(
         return_full_text=False,
     )
     output = ""
-    bot_response = ""
 
     for response in stream:
         response_text = filter_sensitive_info(response.token.text)
         response_text = remove_length_info(response_text)
         output += response_text
-        bot_response += response_text
         yield output
 
     if user_id != "":
         # Sauvegarder le message de l'utilisateur et la réponse complète du bot dans Firestore
-        save_message(user_id, prompt)
-        save_message(user_id, bot_response)
+        save_message(user_id, prompt, role="user")
+        save_message(user_id, output, role="bot")
 
     return output