Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import requests
+import json
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.llms.base import LLM
@@ -48,13 +49,25 @@ class CustomLLM(LLM):
             "Authorization": f"Bearer {dubs_key}",
             "Content-Type": "application/json",
         }
-        payload = {"prompt": prompt}  # Match the API's expected input key
+        payload = {"prompt": prompt}
 
         try:
             response = requests.post(SPACE_URL, json=payload, headers=headers, stream=True)
             response.raise_for_status()
-            result = response.json()
+
+            # Debugging: Log the raw response
+            raw_response = response.text
+            print("Raw Response:", raw_response)
+
+            # Handle non-JSON responses
+            try:
+                result = response.json()
+            except json.JSONDecodeError:
+                return f"Error decoding response: {raw_response}"
+
+            # Extract and return the API's response content
             return result.get("response", "No response from the model.")
+
         except requests.exceptions.HTTPError as e:
             if response.status_code == 401:
                 return "Error: Unauthorized. Please check your Dubs Key."
@@ -110,4 +123,3 @@ if prompt := st.chat_input():
 
 # Update the LangChain conversation to use the selected session's memory
 st.session_state["conversation"].memory = memory
-