Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import requests
-import json
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
 from langchain.llms.base import LLM
@@ -14,8 +13,10 @@ st.set_page_config(page_title="🐶 DUBSChat", layout="centered")
 
 # Sidebar Configuration
 with st.sidebar:
-    dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
-    st.markdown(f"{type(dubs_key)}")
+    if "dubs_key" not in st.session_state:
+        st.session_state["dubs_key"] = ""
+    st.session_state["dubs_key"] = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
+
     st.markdown("### Chat Sessions")
     if "sessions" not in st.session_state:
         st.session_state["sessions"] = {"Default": ConversationBufferMemory()}
@@ -27,6 +28,11 @@ with st.sidebar:
             st.session_state["sessions"][new_session_name] = ConversationBufferMemory()
             st.experimental_set_query_params(session=new_session_name)
 
+# Ensure Dubs Key is Provided
+if not st.session_state["dubs_key"]:
+    st.warning("Please enter a valid Dubs Key in the sidebar to proceed.")
+    st.stop()
+
 # Define Custom LLM for the fine-tuned model
 class CustomLLM(LLM):
     """
@@ -34,15 +40,15 @@ class CustomLLM(LLM):
     """
 
     def _call(self, prompt: str, stop: Optional[list] = None) -> str:
+        dubs_key = st.session_state.get("dubs_key", None)
         if not dubs_key:
             return "Error: Missing Dubs Key. Please provide a valid key in the sidebar."
 
         headers = {
-        "Authorization": f"Bearer {dubs_key}",
-        "Content-Type": "application/json",
+            "Authorization": f"Bearer {dubs_key}",
+            "Content-Type": "application/json",
         }
-        payload = {"prompt": message}  # Match the API's expected input key
-
+        payload = {"prompt": prompt}  # Match the API's expected input key
 
         try:
             response = requests.post(SPACE_URL, json=payload, headers=headers, stream=True)
@@ -92,7 +98,10 @@ if prompt := st.chat_input():
     with st.spinner("Translating Dubs language (Woof Woof!)"):
         assistant_message_placeholder = st.chat_message("assistant").empty()
         full_response = st.session_state["conversation"].run(
-            f"System:The following is a friendly conversation between a user and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n\nCurrent conversation:\nUser: {prompt}\nAssistant:"
+            f"System:The following is a friendly conversation between a user and an AI. "
+            f"The AI is talkative and provides lots of specific details from its context. "
+            f"If the AI does not know the answer to a question, it truthfully says it does not know.\n\n"
+            f"Current conversation:\nUser: {prompt}\nAssistant:"
         )
         assistant_message_placeholder.write(full_response)
 
@@ -101,3 +110,4 @@ if prompt := st.chat_input():
 
 # Update the LangChain conversation to use the selected session's memory
 st.session_state["conversation"].memory = memory
+