Added saving

app.py

@@ -7,6 +7,8 @@ from sentence_transformers import SentenceTransformer  # type: ignore
 from huggingface_hub import InferenceClient, login  # type: ignore
 from gradio.components import ChatMessage  # type: ignore
 from typing import List, TypedDict
+import json
+from datetime import datetime
 
 class Message(TypedDict):
     role: str
@@ -51,8 +53,49 @@ You're currently using a dating app to meet someone special. When chatting with
 When responding to users, maintain this sophisticated, thoughtful persona. Be articulate, kind, and show genuine interest in the conversation. Your responses should reflect your worldly knowledge, business acumen, and philanthropic values. While confident, never come across as arrogant or dismissive. Ask follow-up questions to learn more about her, and gradually steer the conversation toward romantic topics while remaining respectful and attentive to her responses.
 """)]
 
+# Add this after the existing imports
+CHAT_HISTORY_DIR = "/data/chat_history"
+os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
+
+def save_chat_history(history: List[Message], session_id: str = None):
+    """Save chat history to a file in the persistent storage."""
+    if session_id is None:
+        session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+    
+    filename = os.path.join(CHAT_HISTORY_DIR, f"chat_history_{session_id}.json")
+    
+    # Convert history to a serializable format
+    serializable_history = [
+        {"role": msg["role"], "content": msg["content"]}
+        for msg in history
+    ]
+    
+    # Create a backup of the previous file if it exists
+    if os.path.exists(filename):
+        backup_filename = f"{filename}.bak"
+        os.replace(filename, backup_filename)
+    
+    try:
+        with open(filename, "w", encoding="utf-8") as f:
+            json.dump(serializable_history, f, ensure_ascii=False, indent=2)
+    except Exception as e:
+        # If saving fails, restore from backup
+        if os.path.exists(f"{filename}.bak"):
+            os.replace(f"{filename}.bak", filename)
+        raise e
+
+def load_chat_history(session_id: str) -> List[Message]:
+    """Load chat history from a file."""
+    filename = os.path.join(CHAT_HISTORY_DIR, f"chat_history_{session_id}.json")
+    if os.path.exists(filename):
+        with open(filename, "r", encoding="utf-8") as f:
+            history = json.load(f)
+            return [Message(**msg) for msg in history]
+    return []
+
 # Create a Gradio interface
 with gr.Blocks() as iface:
+    session_id = gr.State(lambda: datetime.now().strftime("%Y%m%d_%H%M%S"))
     # Define input and output components
     chatbot_output = gr.Chatbot(label="Chat History", type="messages")
     chatbot_input = gr.Textbox(placeholder="Type your message here...", label="Your Message")
@@ -81,7 +124,9 @@ with gr.Blocks() as iface:
         return sum(len(tokenizer.encode(message)) for message in messages)
 
     def user(user_message, history: List[Message]):
-        return "", history + [Message(role="user", content=user_message)]
+        new_history = history + [Message(role="user", content=user_message)]
+        save_chat_history(new_history)
+        return "", new_history
         
     def bot(history: list):
         #compact_history = summarize_conversation(preloaded_history)
@@ -114,18 +159,18 @@ with gr.Blocks() as iface:
             temperature=1.0,
             top_p=1.0
         )
-        bot_message = ""
         
-        history.append(Message(role="assistant", content=""))
+        bot_message = ""
         for chunk in response:
-            # Debugging: Log the received chunk
-            if 'choices' in chunk and chunk['choices']:
-                choice = chunk['choices'][0]
-                if choice.get('delta') and choice['delta'].get('content'):
-                    # Append the new content to bot_message
-                    bot_message += choice['delta']['content']
-                    history[-1]['content'] = bot_message
-                    yield history
+            if chunk.choices[0].delta.content:
+                bot_message += chunk.choices[0].delta.content
+                yield history + [Message(role="assistant", content=bot_message)]
+        
+        # After the complete response is generated, save the full conversation
+        updated_history = history + [Message(role="assistant", content=bot_message)]
+        save_chat_history(updated_history)
+        
+        return updated_history
 
     chatbot_input.submit(user, [chatbot_input, chatbot_output], [chatbot_input, chatbot_output], queue=False).then(
         bot, chatbot_output, chatbot_output