Refactor message handling, add undo/retry features

app.py

@@ -3,6 +3,7 @@ import requests
 import json
 import os
 import uuid
+import traceback
 
 ENDPOINT = os.environ.get("ENDPOINT", None)
 TOKEN = os.environ.get("TOKEN", None)
@@ -33,10 +34,7 @@ def convert_to_messages(message, history):
 
     return messages
 
-def chat_cosmosllama(message: str, 
-                     history: list,
-                     conversation_id: str
-) -> str:
+def chat_cosmosllama(message: str, history: list, conversation_id: str, parent_message_id: str) -> (str, str):
     messages = convert_to_messages(message, history)
 
     headers = {
@@ -44,52 +42,98 @@ def chat_cosmosllama(message: str,
         "Content-Type": "application/json"
     }
 
-    response = requests.post(f"https://model.cosmosytu.online{ENDPOINT}", json={"messages": messages, "conversation_id": conversation_id}, headers=headers, stream=True)
+    body = {"messages": messages}
+    body["conversation_id"] = conversation_id
+    
+    if parent_message_id:
+        body["parent_message_id"] = parent_message_id
+
+    response = requests.post(f"https://model.cosmosytu.online{ENDPOINT}", json=body, headers=headers, stream=True)
+
+    full_response = ""
+    new_parent_message_id = parent_message_id
 
     if response.status_code == 200:
-        full_response = ""
         for line in response.iter_lines():
             if line:
                 decoded_line = line.decode('utf-8')
                 try:
                     json_line = json.loads(decoded_line[6:])
 
+                    if "assistantMessageId" in json_line:
+                        new_parent_message_id = json_line["assistantMessageId"]
+
                     if "CosmosLlama" in json_line:
                         full_response += json_line['CosmosLlama']
-                        yield full_response
-                except Exception as e:
-                    print(e)
-                    pass
+                        yield full_response, new_parent_message_id
+                except Exception:
+                    traceback.print_exc()
     else:
-        yield "Bir hata oluştu."
+        yield "Bir hata oluştu.", new_parent_message_id
 
 def get_conversation_id():
     return str(uuid.uuid4())
 
-def reset_conversation():
-    new_conversation_id = get_conversation_id()
-    return new_conversation_id
+def clear_conversation():
+    return get_conversation_id(), []
 
-chatbot=gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Cosmos LLaMa')
+chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Cosmos LLaMa')
 
 with gr.Blocks(fill_height=True, css=css) as demo:
+    conversation_id_state = gr.State(value=get_conversation_id())
+    parent_message_history = gr.State(value=[])
     
-    conversation_id_state = gr.State(value=get_conversation_id)
-    
+    def handle_user_message(message, history, conversation_id_state, parent_message_history):
+        is_parent_message_set = False
+        response_generator = chat_cosmosllama(message, history, conversation_id_state, parent_message_history[-1] if parent_message_history else None)
+        for response, updated_parent_message_id in response_generator:
+            if not is_parent_message_set:
+                parent_message_history.append(updated_parent_message_id)
+                is_parent_message_set = True
+            yield response
+
+    def undo_message(history, parent_message_history):
+        if parent_message_history:
+            parent_message_history.pop()
+            history.pop()
+            return history, parent_message_history
+        return history, parent_message_history
+
+    def retry_message(history, conversation_id_state, parent_message_history):
+        if parent_message_history and history:
+            parent_message_history.pop()
+            history.pop()
+            return history, parent_message_history
+        return history, parent_message_history
+
     interface = gr.ChatInterface(
-        fn=chat_cosmosllama,
+        fn=handle_user_message,
         chatbot=chatbot,
         fill_height=True,
         cache_examples=False,
-        additional_inputs=[conversation_id_state],
-        undo_btn=None,
+        additional_inputs=[conversation_id_state, parent_message_history],
         retry_btn="🔄 Tekrar Cevapla",
+        undo_btn="↩️ Geri Al",
         clear_btn="🗑️ Temizle",
         submit_btn="Gönder",
         stop_btn="Durdur"
     )
 
-    interface.retry_btn.click(fn=reset_conversation, inputs=[], outputs=[conversation_id_state])
+    demo.load(fn=clear_conversation, inputs=[], outputs=[conversation_id_state, parent_message_history])
+
+    interface.undo_btn.click(
+        fn=undo_message,
+        inputs=[chatbot, parent_message_history],
+        outputs=[chatbot, parent_message_history]
+    )
+
+    interface.retry_btn.click(
+        fn=retry_message,
+        inputs=[chatbot, conversation_id_state, parent_message_history],
+        outputs=[chatbot, parent_message_history]
+    )
+
+    interface.clear_btn.click(fn=clear_conversation, inputs=[], outputs=[conversation_id_state, parent_message_history])
     interface.textbox.placeholder="Mesajınızı buraya giriniz..."
 
 if __name__ == "__main__":