Update app.py

app.py

@@ -21,7 +21,7 @@ if "follow_up" not in st.session_state:
     st.session_state.follow_up = ""  # Stores follow-up question
 
 if "last_topic" not in st.session_state:
-    st.session_state.last_topic = ""  # Tracks last discussed topic
+    st.session_state.last_topic = ""  # Stores last user topic
 
 # --- Set Up Model & API Functions ---
 model_id = "mistralai/Mistral-7B-Instruct-v0.3"
@@ -51,17 +51,9 @@ def analyze_sentiment(user_text):
     return result['label']
 
 def predict_action(user_text):
-    """
-    Determines the topic of the user's message.
-    """
     if "NASA" in user_text or "space" in user_text:
         return "nasa_info"
-    elif "quark" in user_text or "physics" in user_text or "quantum" in user_text:
-        return "physics"
-    elif "AI" in user_text or "machine learning" in user_text:
-        return "AI"
-    else:
-        return "general_query"
+    return "general_query"
 
 def generate_follow_up(user_text):
     """
@@ -80,7 +72,7 @@ def generate_follow_up(user_text):
 
 def get_response(system_message, chat_history, user_text, max_new_tokens=256):
     """
-    Generates HAL's response and follow-up, ensuring no duplicate queries or misplaced follow-ups.
+    Generates HAL's response, making it more conversational and engaging.
     """
     sentiment = analyze_sentiment(user_text)
     action = predict_action(user_text)
@@ -88,43 +80,38 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256):
     if action == "nasa_info":
         nasa_url, nasa_title, nasa_explanation = get_nasa_apod()
         response = f"**{nasa_title}**\n\n{nasa_explanation}"
+        chat_history.append({'role': 'user', 'content': user_text})
         chat_history.append({'role': 'assistant', 'content': response})
+        
         follow_up = generate_follow_up(user_text)
+        chat_history.append({'role': 'assistant', 'content': follow_up})
         return response, follow_up, chat_history, nasa_url
 
     hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.9)
 
     prompt = PromptTemplate.from_template(
-        "[INST] {system_message}\n\n"
-        "Current Conversation:\n{chat_history}\n\n"
-        "User: {user_text}.\n [/INST]\n"
-        "AI: Keep responses conversational and engaging. Start with a friendly phrase like "
-        "'Certainly!', 'Of course!', or 'Great question!' before answering."
-        " Keep responses concise but engaging.\nHAL:"
+        (
+            "[INST] {system_message}"
+            "\nCurrent Conversation:\n{chat_history}\n\n"
+            "\nUser: {user_text}.\n [/INST]"
+            "\nAI: Keep responses conversational and engaging. Start with a friendly phrase like "
+            "'Certainly!', 'Of course!', or 'Great question!' before answering."
+            " Keep responses concise but engaging."
+            "\nHAL:"
+        )
     )
 
     chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
     response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
     response = response.split("HAL:")[-1].strip()
 
-    # ✅ Avoid duplicate user messages in history
-    if not chat_history or chat_history[-1]["content"] != user_text:
-        chat_history.append({'role': 'user', 'content': user_text})
-
+    chat_history.append({'role': 'user', 'content': user_text})
     chat_history.append({'role': 'assistant', 'content': response})
 
-    # ✅ Avoid repeating follow-ups when topic changes
-    current_topic = action
-    if current_topic != st.session_state.last_topic:
-        st.session_state.follow_up = ""
-    else:
-        follow_up = generate_follow_up(user_text)
-        chat_history.append({'role': 'assistant', 'content': follow_up})
-        st.session_state.follow_up = follow_up
+    follow_up = generate_follow_up(user_text)
+    chat_history.append({'role': 'assistant', 'content': follow_up})
 
-    st.session_state.last_topic = current_topic
-
-    return response, st.session_state.follow_up, chat_history, None
+    return response, follow_up, chat_history, None
 
 # --- Chat UI ---
 st.title("🚀 HAL - Your NASA AI Assistant")
@@ -135,10 +122,41 @@ if st.sidebar.button("Reset Chat"):
     st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
     st.session_state.response_ready = False
     st.session_state.follow_up = ""
-    st.session_state.last_topic = ""
-    st.rerun()
-
-# --- Chat History Display ---
+    st.experimental_rerun()
+
+# Custom Chat Styling
+st.markdown("""
+    <style>
+    .user-msg {
+        background-color: #0078D7;
+        color: white;
+        padding: 10px;
+        border-radius: 10px;
+        margin-bottom: 5px;
+        width: fit-content;
+        max-width: 80%;
+    }
+    .assistant-msg {
+        background-color: #333333;
+        color: white;
+        padding: 10px;
+        border-radius: 10px;
+        margin-bottom: 5px;
+        width: fit-content;
+        max-width: 80%;
+    }
+    .container {
+        display: flex;
+        flex-direction: column;
+        align-items: flex-start;
+    }
+    @media (max-width: 600px) {
+        .user-msg, .assistant-msg { font-size: 16px; max-width: 100%; }
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+# --- Chat History Display (Ensures All Messages Are Visible) ---
 st.markdown("<div class='container'>", unsafe_allow_html=True)
 
 for message in st.session_state.chat_history:
@@ -150,12 +168,11 @@ for message in st.session_state.chat_history:
 st.markdown("</div>", unsafe_allow_html=True)
 
 # --- Single Input Box for Both Initial and Follow-Up Messages ---
-user_input = st.chat_input("Type your message here...")
+user_input = st.chat_input("Type your message here...")  # Uses Enter to submit
 
 if user_input:
-    # ✅ Prevent duplicate user messages
-    if not st.session_state.chat_history or st.session_state.chat_history[-1]["content"] != user_input:
-        st.session_state.chat_history.append({'role': 'user', 'content': user_input})
+    # Save user message in chat history
+    st.session_state.chat_history.append({'role': 'user', 'content': user_input})  
 
     # Generate HAL's response
     response, follow_up, st.session_state.chat_history, image_url = get_response(
@@ -165,14 +182,16 @@ if user_input:
     )
 
     st.session_state.chat_history.append({'role': 'assistant', 'content': response})
+
     st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {response}</div>", unsafe_allow_html=True)
 
     if image_url:
         st.image(image_url, caption="NASA Image of the Day")
 
     st.session_state.follow_up = follow_up
-    st.session_state.response_ready = True  
+    st.session_state.response_ready = True  # Enables follow-up response cycle
 
 if st.session_state.response_ready and st.session_state.follow_up:
+    st.session_state.chat_history.append({'role': 'assistant', 'content': st.session_state.follow_up})
     st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {st.session_state.follow_up}</div>", unsafe_allow_html=True)
     st.session_state.response_ready = False