Update app.py

app.py

@@ -23,7 +23,7 @@ except OSError as e:
 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)
+st.image(IMAGE_PATH_2, width=200)
 
 # -------------------------
 #    Utility Functions
@@ -114,28 +114,39 @@ for message in st.session_state["messages"]:
 # -------------------------
 #    Streaming Logic
 # -------------------------
-def stream_response(prompt_text, api_key):
+def stream_response(chat_history, api_key):
     """
     Stream text from the HF Inference Endpoint (or any streaming API).
     Yields each chunk of text as it arrives.
     """
     try:
+        # Prepare payload for the POST request
+        payload = {
+            "inputs": chat_history,
+            "parameters": {
+                "temperature": 0.7,
+                "top_p": 0.9,
+                "max_new_tokens": 200,
+            },
+            "options": {
+                "stream": True
+            },
+        }
+
         # POST request with stream=True to get partial chunks
         response = requests.post(
             SPACE_URL,
-            json={"prompt": prompt_text},  # Adjust this to match your endpoint’s payload
+            json=payload,
             headers={"Authorization": f"Bearer {api_key}"},
             stream=True
         )
         response.raise_for_status()
 
-        # The endpoint presumably returns lines of JSON. Adjust parsing if your endpoint differs:
+        # Process and yield streamed chunks of text
         for line in response.iter_lines():
             if line:
                 data = json.loads(line.decode("utf-8"))
-                # Example format: data might be [{"generated_text": "..."}]
-                # Adjust if your endpoint returns different JSON keys
-                chunk = data[0].get("generated_text", "")
+                chunk = data.get("generated_text", "")
                 yield chunk
 
     except requests.exceptions.Timeout:
@@ -157,23 +168,9 @@ if prompt := st.chat_input():
         st.chat_message("user").write(prompt)
 
         # 2) Build combined chat history for the model prompt
-        #    or whatever format your endpoint needs. Example:
         chat_history = "".join(
             [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
         )
 
         # 3) Create a placeholder for the assistant’s streamed response
-        with st.spinner("Dubs is thinking... Woof Woof! 🐾"):
-            assistant_message_placeholder = st.chat_message("assistant", avatar=Dubs_PATH).empty()
-
-            full_response = ""
-            # 4) Stream chunks from the API
-            for chunk in stream_response(chat_history, HF_API_KEY):
-                full_response += chunk
-                # Continuously update the placeholder with the partial response
-                assistant_message_placeholder.write(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"])
+        with st.spinner("Dubs is thinking... Wo