Update app.py

app.py

@@ -11,7 +11,10 @@ EOS_TOKEN = "<|end|>"
 CHAT_HISTORY_DIR = "chat_histories"  # Directory to save chat histories
 
 # Ensure the directory exists
-os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
+try:
+    os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
+except OSError as e:
+    st.error(f"Failed to create chat history directory: {e}")
 
 # Streamlit Configurations (must come first)
 st.set_page_config(page_title="🐶 DUBSChat", layout="wide")
@@ -22,16 +25,23 @@ 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")
-    with open(file_path, "w") as f:
-        json.dump(messages, f)
+    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)
-    with open(file_path, "r") as f:
-        return json.load(f)
+    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():
     """
@@ -71,8 +81,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?"}
     ]
+if "session_name" not in st.session_state:
     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")
@@ -87,36 +97,54 @@ for message in st.session_state["messages"]:
 
 # User Input
 if prompt := st.chat_input():
-    # 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"]]
-    )
-
-    # Get response from API
-    with st.spinner("Translating Dubs language (Woof Woof!)"):
-        assistant_message_placeholder = st.chat_message("assistant").empty()
-        full_response = ""
-        try:
-            response = requests.post(
-                SPACE_URL,
-                json={"prompt": chat_history},
-                headers={"Authorization": f"Bearer {dubs_key}"},
-                timeout=120,
-            )
-            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"])
+    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("Translating Dubs language (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()
+
+                    # Read and yield the response incrementally
+                    for chunk in response.iter_content(chunk_size=256):
+                        if chunk:
+                            yield chunk.decode("utf-8")
+                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"])