Update app.py

app.py

@@ -13,6 +13,7 @@ IMAGE_PATH = "DubsChat.png"
 IMAGE_PATH_2 = "Reboot AI.png"
 Dubs_PATH = "Dubs.png"
 
+
 # Ensure the directory exists
 try:
     os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
@@ -20,28 +21,17 @@ except OSError as e:
     st.error(f"Failed to create chat history directory: {e}")
 
 # Streamlit Configurations
-st.set_page_config(page_title="DUBSChat", page_icon=IMAGE_PATH, layout="wide")
+st.set_page_config(page_title="DUBSChat",page_icon=IMAGE_PATH, layout="wide")
+
+
 st.logo(IMAGE_PATH_2)
 
+
 # Utility Functions
-def query(payload):
+def save_chat_history(session_name, messages):
     """
-    Query the Hugging Face endpoint.
+    Save the chat history to a JSON file.
     """
-    headers = {
-        "Authorization": f"Bearer {HF_API_KEY}",
-        "Content-Type": "application/json",
-    }
-    try:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        response.raise_for_status()  # Raise an error for bad responses (4xx, 5xx)
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        return {"error": f"Request error: {e}"}
-    except json.JSONDecodeError:
-        return {"error": "Invalid JSON response from server"}
-
-def save_chat_history(session_name, messages):
     file_path = os.path.join(CHAT_HISTORY_DIR, f"{session_name}.json")
     try:
         with open(file_path, "w") as f:
@@ -49,7 +39,11 @@ def save_chat_history(session_name, messages):
     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:
@@ -58,11 +52,18 @@ def load_chat_history(file_name):
         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:
+
+        # Reset chat
     if st.button("New Chat"):
         st.session_state["messages"] = [
             {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
@@ -72,7 +73,11 @@ with st.sidebar:
         save_chat_history(st.session_state["session_name"], st.session_state["messages"])
         st.success("Chat reset and new session started.")
 
+    dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
     saved_sessions = get_saved_sessions()
+    
+    
+    # Display saved sessions
     if saved_sessions:
         selected_session = st.radio("Past Sessions:", saved_sessions)
         if st.button("Load Session"):
@@ -82,6 +87,8 @@ with st.sidebar:
     else:
         st.write("No past sessions available.")
 
+
+
 # Chat History Initialization
 if "messages" not in st.session_state:
     st.session_state["messages"] = [
@@ -92,7 +99,9 @@ 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.image(IMAGE_PATH, width=250)
+
+st.image(IMAGE_PATH,width=250)
+
 st.markdown("Empowering you with a Sustainable AI")
 
 # Display Chat History
@@ -100,13 +109,15 @@ 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", avatar=Dubs_PATH).write(message["content"])
+        st.chat_message("assistant",avatar=Dubs_PATH).write(message["content"])
+        
 
 # User Input
 if prompt := st.chat_input():
-    if not HF_API_KEY:
-        st.warning("Please set the Hugging Face API key in the system environment.")
+    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)
 
@@ -115,23 +126,42 @@ if prompt := st.chat_input():
             [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
         )
 
-        # Query the model
+        # Stream response
         with st.spinner("Dubs is thinking... Woof Woof! 🐾"):
             assistant_message_placeholder = st.chat_message("assistant").empty()
 
-            payload = {
-                "inputs": chat_history,
-                "parameters": {
-                    "max_new_tokens": 150
-                }
-            }
-            response = query(payload)[0]["generated_text"]
-            if "error" in response:
-                assistant_response = response["error"]
-            else:
-                assistant_response = response.get("generated_text", "Error: No generated text in response.")
-            
-            assistant_message_placeholder.write(assistant_response)
-
-        st.session_state["messages"].append({"role": "assistant", "content": assistant_response})
-        save_chat_history(st.session_state["session_name"], st.session_state["messages"])
+            # Function to stream the response
+            def stream_response():
+                try:
+                    response = requests.post(
+                        SPACE_URL,
+                        json={"prompt": chat_history},
+                        headers={"Authorization": f"Bearer {HF_API_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"])