Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import requests
 import json
 import os
+import datetime
 
 # Constants
 SPACE_URL = "https://abanm-dubs.hf.space/api/generate"
@@ -12,101 +13,56 @@ CHAT_HISTORY_DIR = "chat_histories"  # Directory to save chat histories
 # Ensure the directory exists
 os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
 
-# Streamlit Configurations (must come first)
+# Streamlit Configurations
 st.set_page_config(page_title="🐶 DUBSChat", layout="wide")
 
-# Sidebar Configuration
-with st.sidebar:
-    dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
-    st.markdown("### Dubs Recall")
-    
-    # Save chat history
-    if st.button("💾 Save Chat"):
-        if "messages" in st.session_state:
-            file_name = st.text_input("Enter filename to save chat:", "chat_history.json")
-            if file_name:
-                file_path = os.path.join(CHAT_HISTORY_DIR, file_name)
-                with open(file_path, "w") as f:
-                    json.dump(st.session_state["messages"], f)
-                st.success(f"Chat saved as {file_name}!")
-
-    # Load chat history
-    file_to_load = st.selectbox("Select a chat to load:", os.listdir(CHAT_HISTORY_DIR))
-    if st.button("📂 Load Chat"):
-        file_path = os.path.join(CHAT_HISTORY_DIR, file_to_load)
-        if os.path.exists(file_path):
-            with open(file_path, "r") as f:
-                st.session_state["messages"] = json.load(f)
-            st.success(f"Loaded chat from {file_to_load}!")
-        else:
-            st.error(f"File {file_to_load} does not exist.")
-
-    # Reset chat
-    if st.button("🔄 Reset Chat"):
-        st.session_state["messages"] = [
-            {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
-            {"role": "assistant", "content": "Hello! How can I assist you today?"}
-        ]
-
-# Template for serializing chat history
-CHAT_TEMPLATE = """
-{% for message in messages %}
-{% if message['role'] == 'system' %}{{'<|system|>\n' + message['content'] + '<|end|>\n'}}
-{% elif message['role'] == 'user' %}{{'<|user|>\n' + message['content'] + '<|end|>\n'}}
-{% elif message['role'] == 'assistant' %}{{'<|assistant|>\n' + message['content'] + '<|end|>\n'}}
-{% endif %}
-{% endfor %}
-{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% else %}{{ eos_token }}{% endif %}
-"""
-
-def format_chat_history(messages, add_generation_prompt=True):
+# Utility Functions
+def save_chat_history(session_name, messages):
     """
-    Format the chat history using the provided template.
-    
-    Args:
-        messages (list): List of message dictionaries with 'role' and 'content'.
-        add_generation_prompt (bool): Whether to include an assistant generation prompt.
-    
-    Returns:
-        str: Formatted chat history string.
+    Save the chat history to a JSON file.
     """
-    from jinja2 import Template
+    file_path = os.path.join(CHAT_HISTORY_DIR, f"{session_name}.json")
+    with open(file_path, "w") as f:
+        json.dump(messages, f)
 
-    template = Template(CHAT_TEMPLATE)
-    return template.render(messages=messages, add_generation_prompt=add_generation_prompt, eos_token=EOS_TOKEN)
+def load_chat_history(file_name):
+    """
+    Load the chat history from a JSON file.
+    """
+    file_path = os.path.join(CHAT_HISTORY_DIR, file_name)
+    with open(file_path, "r") as f:
+        return json.load(f)
 
-def call_api(chat_history):
+def get_saved_sessions():
+    """
+    Get the list of saved chat sessions.
     """
-    Calls the backend API with the formatted chat history.
+    return [f.replace(".json", "") for f in os.listdir(CHAT_HISTORY_DIR) if f.endswith(".json")]
+
+# Sidebar Configuration
+with st.sidebar:
+    st.title("📂 Sessions")
+    saved_sessions = get_saved_sessions()
     
-    Args:
-        chat_history (str): Formatted chat history string.
+    # Display saved sessions
+    if saved_sessions:
+        selected_session = st.radio("Load a past session:", saved_sessions)
+        if st.button("Load Session"):
+            st.session_state["messages"] = load_chat_history(f"{selected_session}.json")
+            st.success(f"Loaded session: {selected_session}")
+    else:
+        st.write("No saved sessions yet.")
     
-    Yields:
-        str: Streamed response chunks from the backend API.
-    """
-    headers = {
-        "Authorization": f"Bearer {dubs_key}",
-        "Content-Type": "application/json",
-    }
-    payload = {"prompt": chat_history}
-
-    try:
-        response = requests.post(SPACE_URL, json=payload, headers=headers, stream=True)
-        response.raise_for_status()
-
-        for chunk in response.iter_lines(decode_unicode=True):
-            if chunk.strip():
-                try:
-                    data = json.loads(chunk)
-                    if "response" in data:
-                        yield data["response"]
-                except json.JSONDecodeError:
-                    yield "Error decoding response chunk."
-    except requests.exceptions.Timeout:
-        yield "Error: The API call timed out."
-    except requests.exceptions.RequestException as e:
-        yield f"Error: {str(e)}"
+    # Create new session
+    if st.button("Start New Session"):
+        session_name = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        st.session_state["messages"] = [
+            {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
+            {"role": "assistant", "content": "Hello! How can I assist you today?"}
+        ]
+        st.session_state["session_name"] = session_name
+        save_chat_history(session_name, st.session_state["messages"])
+        st.success(f"New session started: {session_name}")
 
 # Chat History Initialization
 if "messages" not in st.session_state:
@@ -114,6 +70,8 @@ if "messages" not in st.session_state:
         {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
         {"role": "assistant", "content": "Hello! How can I assist you today?"}
     ]
+    st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    save_chat_history(st.session_state["session_name"], st.session_state["messages"])
 
 # Main Chat UI
 st.title("🐶 DUBSChat")
@@ -133,15 +91,30 @@ if prompt := st.chat_input():
     st.chat_message("user").write(prompt)
 
     # Serialize chat history for API
-    chat_history = format_chat_history(st.session_state["messages"])
+    chat_history = "".join(
+        [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
+    )
 
     # Get response from API
     with st.spinner("Translating Dubs language (Woof Woof!)"):
         assistant_message_placeholder = st.chat_message("assistant").empty()
         full_response = ""
-        for response_chunk in call_api(chat_history):
-            full_response += response_chunk
-            assistant_message_placeholder.write(full_response)
-
-        # Append the assistant's response to the chat history
-        st.session_state["messages"].append({"role": "assistant", "content": full_response})
+        try:
+            response = requests.post(
+                SPACE_URL,
+                json={"prompt": chat_history},
+                headers={"Authorization": f"Bearer {API_KEY}"},
+                timeout=60,
+            )
+            response.raise_for_status()
+            full_response = response.json().get("response", "Sorry, something went wrong!")
+        except requests.exceptions.RequestException as e:
+            full_response = f"Error: {e}"
+
+        assistant_message_placeholder.write(full_response)
+
+    # Append the assistant's response to the chat history
+    st.session_state["messages"].append({"role": "assistant", "content": full_response})
+
+    # Auto-save the updated chat history
+    save_chat_history(st.session_state["session_name"], st.session_state["messages"])