Update app.py

app.py

@@ -23,8 +23,10 @@ except OSError as e:
 # Streamlit Configurations
 st.set_page_config(page_title="DUBSChat", page_icon=IMAGE_PATH, layout="wide")
 
-# If you are using a custom "logo" method:
-st.logo(IMAGE_PATH_2)
+# NOTE: st.logo() is not a built-in Streamlit function.
+# Replace or remove if needed. For a simple image display, use:
+# st.image(IMAGE_PATH_2)
+st.logo(IMAGE_PATH_2)  # If you have a custom method for logos, otherwise replace with st.image()
 
 # -------------------------
 #    Utility Functions
@@ -67,8 +69,14 @@ 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."},
-            {"role": "assistant", "content": "Hello! How can I assist you today?"}
+            {
+                "role": "system", 
+                "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."
+            },
+            {
+                "role": "assistant", 
+                "content": "Hello! How can I assist you today?"
+            }
         ]
         st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
         save_chat_history(st.session_state["session_name"], st.session_state["messages"])
@@ -90,8 +98,14 @@ with st.sidebar:
 # -------------------------
 if "messages" not in st.session_state:
     st.session_state["messages"] = [
-        {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
-        {"role": "assistant", "content": "Hello! How can I assist you today?"}
+        {
+            "role": "system", 
+            "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."
+        },
+        {
+            "role": "assistant", 
+            "content": "Hello! How can I assist you today?"
+        }
     ]
 if "session_name" not in st.session_state:
     st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
@@ -135,6 +149,10 @@ def stream_response(prompt_text, api_key):
         for response in stream:
             if response.token.special:
                 continue
+            # Check for stop sequence and break if encountered
+            if response.token.text in gen_kwargs["stop_sequences"]:
+                break
+
             partial_text += response.token.text
             yield response.token.text
     except Exception as e:
@@ -146,28 +164,26 @@ def stream_response(prompt_text, api_key):
 prompt = st.chat_input()
 
 if prompt:
-        # 1) Add the user's message to session state
-        st.session_state["messages"].append({"role": "user", "content": prompt})
-        st.chat_message("user").write(prompt)
+    # 1) Add the user's message to session state
+    st.session_state["messages"].append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
 
-        # 2) Build combined chat history for the model prompt
-        chat_history = "".join(
-            [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
-        )
+    # 2) Build combined chat history for the model prompt
+    chat_history = "".join(
+        [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
+    )
 
-    with st.spinner("Dubs is thinking... Woof Woof! 🐾"):    
+    # 3) Generate the response under a spinner
+    with st.spinner("Dubs is thinking... Woof Woof! 🐾"):
         with st.chat_message("assistant", avatar=DUBS_PATH):
-                full_response = ""
-                placeholder = st.empty()
+            full_response = ""
+            placeholder = st.empty()
             response = stream_response(chat_history, HF_API_KEY)
             for item in response:
                 full_response += item
                 placeholder.markdown(full_response)
-            placeholder.markdown(full_response)
-
-        # 5) Save the final assistant message in session state
-        st.session_state["messages"].append({"role": "assistant", "content": full_response})
-        # 6) Persist updated chat history
-        save_chat_history(st.session_state["session_name"], st.session_state["messages"])
-
 
+    # 4) Save the final assistant message in session state
+    st.session_state["messages"].append({"role": "assistant", "content": full_response})
+    # 5) Persist updated chat history
+    save_chat_history(st.session_state["session_name"], st.session_state["messages"])