Update app.py

app.py

@@ -9,9 +9,9 @@ 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"
-IMAGE_PATH = "DubsChat.png" 
+IMAGE_PATH = "DubsChat.png"
 IMAGE_PATH_2 = "Reboot AI.png"
-Dubs_PATH = "Dubs.png" 
+Dubs_PATH = "Dubs.png"
 
 # Ensure the directory exists
 try:
@@ -34,10 +34,12 @@ def query(payload):
     }
     try:
         response = requests.post(API_URL, headers=headers, json=payload)
-        response.raise_for_status()
+        response.raise_for_status()  # Raise an error for bad responses (4xx, 5xx)
         return response.json()
     except requests.exceptions.RequestException as e:
-        return {"error": str(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")
@@ -124,7 +126,11 @@ if prompt := st.chat_input():
                 }
             }
             response = query(payload)
-            assistant_response = response.get("generated_text", "Error: Unable to fetch response.")
+            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})