Update app.py

app.py

@@ -5,10 +5,10 @@ import os
 import datetime
 
 # Constants
-SPACE_URL = "https://abanm-dubs.hf.space/api/generate"
-API_KEY = "s3cr3t_k3y"  # Replace with your actual API key
+API_URL = "https://z7svds7k42bwhhgm.us-east-1.aws.endpoints.huggingface.cloud"
+HF_API_KEY = os.getenv("HF_API_KEY")  # Retrieve the Hugging Face API key from system variables
 EOS_TOKEN = "<|end|>"
-CHAT_HISTORY_DIR = "chat_histories"  # Directory to save chat histories
+CHAT_HISTORY_DIR = "chat_histories"
 IMAGE_PATH = "DubsChat.png" 
 IMAGE_PATH_2 = "Reboot AI.png"
 Dubs_PATH = "Dubs.png" 
@@ -20,17 +20,26 @@ 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 save_chat_history(session_name, messages):
+def query(payload):
     """
-    Save the chat history to a JSON file.
+    Query the Hugging Face endpoint.
     """
+    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()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return {"error": str(e)}
+
+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:
@@ -38,11 +47,7 @@ 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:
@@ -51,18 +56,11 @@ 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,11 +70,7 @@ 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"):
@@ -86,8 +80,6 @@ with st.sidebar:
     else:
         st.write("No past sessions available.")
 
-
-
 # Chat History Initialization
 if "messages" not in st.session_state:
     st.session_state["messages"] = [
@@ -98,9 +90,7 @@ 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
@@ -108,15 +98,13 @@ 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 dubs_key:
-        st.warning("Please provide a valid Dubs Key.")
+    if not HF_API_KEY:
+        st.warning("Please set the Hugging Face API key in the system environment.")
     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)
 
@@ -125,42 +113,19 @@ if prompt := st.chat_input():
             [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
         )
 
-        # Stream response
+        # Query the model
         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
+            payload = {
+                "inputs": chat_history,
+                "parameters": {
+                    "max_new_tokens": 150
+                }
+            }
+            response = query(payload)
+            assistant_response = response.get("generated_text", "Error: Unable to fetch 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"])