Update app.py

app.py

@@ -86,18 +86,23 @@ def ensure_english(text):
 
 # ✅ Ensure Every Response Has a Follow-Up Question
 def generate_follow_up(user_text):
-    """Generates a follow-up question to guide the user toward related topics or next steps."""
+    """Generates a clean follow-up question to guide the user toward related topics or next steps."""
     prompt_text = (
         f"Given the user's question: '{user_text}', generate a SHORT follow-up question "
         "suggesting either a related topic or asking if they need further help. "
         "Example: 'Would you like to explore quantum superposition or ask about another physics concept?' "
         "Keep it concise and engaging."
     )
+    
     hf = get_llm_hf_inference(max_new_tokens=40, temperature=0.8)
     output = hf.invoke(input=prompt_text).strip()
     
-    # Fallback in case of an empty response
-    return output if output else "Would you like to explore another related topic or ask about something else?"
+    # ✅ Remove unnecessary characters (like backticks and misplaced formatting)
+    cleaned_output = re.sub(r"```|''|\"", "", output).strip()
+
+    # ✅ Fallback in case the response is empty or invalid
+    return cleaned_output if cleaned_output else "Would you like to explore another related topic or ask about something else?"
+
 
 # ✅ Main Response Function
 def get_response(system_message, chat_history, user_text, max_new_tokens=512):
@@ -149,6 +154,16 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=512):
 
     return response, follow_up, chat_history, None
 
+# ✅ Ensure response is displayed
+if response:
+    st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {response}</div>", unsafe_allow_html=True)
+
+# ✅ Save and display follow-up question separately
+if follow_up:  # 🔍 Here is the `if follow_up:` section
+    st.session_state.chat_history.append({'role': 'assistant', 'content': follow_up})
+    st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {follow_up}</div>", unsafe_allow_html=True)
+
+
 # ✅ Streamlit UI
 st.title("🚀 HAL - NASA AI Assistant")