Update app.py

app.py

@@ -9,6 +9,7 @@ SPACE_URL = "https://abanm-dubs.hf.space/api/generate"
 API_KEY = "s3cr3t_k3y"  # Replace with your actual API key
 EOS_TOKEN = "<|end|>"
 CHAT_HISTORY_DIR = "chat_histories"  # Directory to save chat histories
+IMAGE_PATH = "/Problems.png"  # Static path to the corner image
 
 # Ensure the directory exists
 try:
@@ -16,10 +17,10 @@ try:
 except OSError as e:
     st.error(f"Failed to create chat history directory: {e}")
 
-# Streamlit Configurations (must come first)
+# Streamlit Configurations
 st.set_page_config(page_title="🐶 DUBSChat", layout="wide")
 
-# Add custom CSS for background and positioning
+# Add custom CSS for background and corner image positioning
 st.markdown(
     """
     <style>
@@ -38,25 +39,142 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-# Display the image in the corner
-image_path = "/Problems.png"  # Static path to the image
+# Display the static image in the corner
 st.markdown(
-    f'<img src="{image_path}" class="corner-image">',
+    f'<img src="{IMAGE_PATH}" class="corner-image">',
     unsafe_allow_html=True,
 )
+
+# Utility Functions
+def save_chat_history(session_name, messages):
+    """
+    Save the chat history to a JSON file.
+    """
+    file_path = os.path.join(CHAT_HISTORY_DIR, f"{session_name}.json")
+    try:
+        with open(file_path, "w") as f:
+            json.dump(messages, f)
+    except IOError as e:
+        st.error(f"Failed to save chat history: {e}")
+
+
+def load_chat_history(file_name):
+    """
+    Load the chat history from a JSON file.
+    """
+    file_path = os.path.join(CHAT_HISTORY_DIR, file_name)
+    try:
+        with open(file_path, "r") as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        st.error("Failed to load chat history. Starting with a new session.")
+        return []
+
+
+def get_saved_sessions():
+    """
+    Get the list of saved chat sessions.
+    """
+    return [f.replace(".json", "") for f in os.listdir(CHAT_HISTORY_DIR) if f.endswith(".json")]
+
+
 # Sidebar Configuration
 with st.sidebar:
     st.title("📂 Dubs Recall")
+    saved_sessions = get_saved_sessions()
     dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
 
+    # Display saved sessions
+    if saved_sessions:
+        selected_session = st.radio("Past Sessions:", saved_sessions)
+        if st.button("Load Session"):
+            st.session_state["messages"] = load_chat_history(f"{selected_session}.json")
+            st.session_state["session_name"] = selected_session
+            st.success(f"Loaded session: {selected_session}")
+    else:
+        st.write("No past sessions available.")
+
+    # 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?"}
+        ]
+        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"])
+        st.success("Chat reset and new session started.")
+
 # Chat History Initialization
 if "messages" not in st.session_state:
     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?"}
     ]
+if "session_name" not in st.session_state:
+    st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
 
+# Main Chat UI
 st.title("🐶 DUBSChat")
 st.markdown("Empowering you with a Sustainable AI")
 
-# Rest of the Chat App Code Here...
+# Display Chat History
+for message in st.session_state["messages"]:
+    if message["role"] == "user":
+        st.chat_message("user").write(message["content"])
+    elif message["role"] == "assistant":
+        st.chat_message("assistant").write(message["content"])
+
+# User Input
+if prompt := st.chat_input():
+    if not dubs_key:
+        st.warning("Please provide a valid Dubs Key.")
+    else:
+        # Append the user's message to the chat history
+        st.session_state["messages"].append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
+
+        # Serialize chat history for API
+        chat_history = "".join(
+            [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
+        )
+
+        # Stream response
+        with st.spinner("Dubs is thinking... Woof Woof! 🐾"):
+            assistant_message_placeholder = st.chat_message("assistant").empty()
+
+            # Function to stream the response
+            def stream_response():
+                try:
+                    response = requests.post(
+                        SPACE_URL,
+                        json={"prompt": chat_history},
+                        headers={"Authorization": f"Bearer {dubs_key}"},
+                        stream=True,  # Enable streaming
+                        timeout=120,
+                    )
+                    response.raise_for_status()
+
+                    # Process and yield only the response field
+                    for line in response.iter_lines():
+                        if line:
+                            data = json.loads(line.decode("utf-8"))
+                            yield data.get("response", "")
+                except requests.exceptions.Timeout:
+                    yield "The request timed out. Please try again later."
+                except requests.exceptions.RequestException as e:
+                    yield f"Error: {e}"
+                except json.JSONDecodeError:
+                    yield "Error decoding server response."
+
+            # Update assistant's message incrementally
+            full_response = ""
+            for chunk in stream_response():
+                full_response += chunk
+                assistant_message_placeholder.write(full_response)
+
+        # Append the assistant's final 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"])
+