Update app.py

app.py

@@ -7,11 +7,11 @@ from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.prompts import PromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from transformers import pipeline
-from langdetect import detect
+from langdetect import detect  # Ensure this package is installed
 
 # ✅ Check for GPU or Default to CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"✅ Using device: {device}")
+print(f"✅ Using device: {device}")  # Debugging info
 
 # ✅ Environment Variables
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -25,23 +25,21 @@ if NASA_API_KEY is None:
 # ✅ Set Up Streamlit
 st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="🚀")
 
-# ✅ Ensure Session State Variables (Maintains Chat History)
+# ✅ Initialize Session State Variables (Ensuring Chat History Persists)
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
 if "response_ready" not in st.session_state:
     st.session_state.response_ready = False
-if "follow_up" not in st.session_state:
-    st.session_state.follow_up = ""
 
-# ✅ Initialize Hugging Face Model (CPU/GPU Compatible)
-def get_llm_hf_inference(model_id="meta-llama/Llama-2-7b-chat-hf", max_new_tokens=800, temperature=0.3):
+# ✅ Initialize Hugging Face Model (Explicitly Set to CPU/GPU)
+def get_llm_hf_inference(model_id="meta-llama/Llama-2-7b-chat-hf", max_new_tokens=800, temperature=0.8):
     return HuggingFaceEndpoint(
         repo_id=model_id,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         token=HF_TOKEN,
         task="text-generation",
-        device=-1 if device == "cpu" else 0
+        device=-1 if device == "cpu" else 0  # ✅ Force CPU (-1) or GPU (0)
     )
 
 # ✅ NASA API Function
@@ -57,7 +55,7 @@ def get_nasa_apod():
 sentiment_analyzer = pipeline(
     "sentiment-analysis",
     model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
-    device=-1 if device == "cpu" else 0
+    device=-1 if device == "cpu" else 0  # ✅ Force CPU (-1) or GPU (0)
 )
 
 def analyze_sentiment(user_text):
@@ -70,73 +68,56 @@ def predict_action(user_text):
         return "nasa_info"
     return "general_query"
 
-# ✅ Ensure English Responses (Fixed Detection Error)
+# ✅ Ensure English Responses
 def ensure_english(text):
-    """Ensures the response is in English, preventing false language detection errors."""
     try:
         detected_lang = detect(text)
-        if detected_lang == "en":
-            return text  # ✅ It's in English, return as-is
+        if detected_lang != "en":
+            return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
     except:
-        pass  # 🔥 Ignore detection errors, assume English
-    
-    return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
+        return "⚠️ Language detection failed. Please ask your question again."
+    return text
 
-# ✅ Main Response Function (Fixed History & Language Issues)
-def get_response(system_message, user_text, max_new_tokens=800):
-    """Generates a response and ensures conversation history is updated."""
-    
-    chat_history = st.session_state.chat_history  # ✅ Get Chat History
-
-    # ✅ Store User Input in Chat History BEFORE Generating Response
-    chat_history.append({'role': 'user', 'content': user_text})
-
-    # ✅ Detect Intent (NASA vs General AI chat)
+# ✅ Main Response Function (Follow-Up Question Removed)
+def get_response(system_message, chat_history, user_text, max_new_tokens=800):
     action = predict_action(user_text)
 
+    # ✅ Handle NASA-Specific Queries
     if action == "nasa_info":
         nasa_url, nasa_title, nasa_explanation = get_nasa_apod()
         response = f"**{nasa_title}**\n\n{nasa_explanation}"
-        follow_up = generate_follow_up(user_text)
-
-        # ✅ Append to chat history
+        chat_history.append({'role': 'user', 'content': user_text})
         chat_history.append({'role': 'assistant', 'content': response})
-        chat_history.append({'role': 'assistant', 'content': follow_up})
-        st.session_state.chat_history = chat_history
-        return response, follow_up, nasa_url
-
-    # ✅ Format Conversation History for Model Input
-    formatted_chat_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in chat_history)
+        return response, chat_history, nasa_url
 
     # ✅ Invoke Hugging Face Model
-    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.3)
+    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.9)
+
+    filtered_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in chat_history)
 
     prompt = PromptTemplate.from_template(
         "[INST] You are a helpful AI assistant.\n\nCurrent Conversation:\n{chat_history}\n\n"
         "User: {user_text}.\n [/INST]\n"
-        "AI: Provide a detailed explanation with depth. Use a conversational tone."
+        "AI: Provide a detailed explanation with depth. Use a conversational tone. "
         "🚨 Answer **only in English**."
+        "Ensure a friendly, engaging tone."
         "\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=formatted_chat_history))
+    response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=filtered_history))
     response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
 
-    # ✅ Prevent False Language Errors
     response = ensure_english(response)
 
     if not response:
         response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?"
 
-    follow_up = generate_follow_up(user_text)
-
-    # ✅ Append Responses to Chat History
-    chat_history.append({'role': 'assistant', 'content': response})
-    chat_history.append({'role': 'assistant', 'content': follow_up})
-    st.session_state.chat_history = chat_history
+    # ✅ Preserve conversation history
+    st.session_state.chat_history.append({'role': 'user', 'content': user_text})
+    st.session_state.chat_history.append({'role': 'assistant', 'content': response})
 
-    return response, follow_up, None
+    return response, chat_history
 
 # ✅ Streamlit UI
 st.title("🚀 HAL - NASA AI Assistant")
@@ -145,7 +126,7 @@ st.title("🚀 HAL - NASA AI Assistant")
 st.markdown("""
     <style>
     .user-msg, .assistant-msg {
-        padding: 10px;
+        padding: 11px;
         border-radius: 10px;
         margin-bottom: 5px;
         width: fit-content;
@@ -159,22 +140,29 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
+# ✅ Reset Chat Button
+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
 
-# ✅ Display Chat History
-for message in st.session_state.chat_history:
-    st.markdown(f"**{message['role'].capitalize()}**: {message['content']}")
-
-# ✅ Chat Input
+# ✅ Chat UI
 user_input = st.chat_input("Type your message here...")
 
 if user_input:
-    response, follow_up, image_url = get_response("You are a helpful AI assistant.", user_input)
+    response, st.session_state.chat_history = get_response(
+        system_message="You are a helpful AI assistant.",
+        user_text=user_input,
+        chat_history=st.session_state.chat_history
+    )
 
     if response:
-        st.markdown(f"**HAL**: {response}")
-
-    if follow_up:
-        st.markdown(f"**HAL**: {follow_up}")
+        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")
+# ✅ Display chat history
+st.markdown("<div class='container'>", unsafe_allow_html=True)
+for message in st.session_state.chat_history:
+    if message["role"] == "user":
+        st.markdown(f"<div class='user-msg'><strong>You:</strong> {message['content']}</div>", unsafe_allow_html=True)
+    else:
+        st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {message['content']}</div>", unsafe_allow_html=True)
+st.markdown("</div>", unsafe_allow_html=True)