Update app.py

app.py

@@ -1,95 +1,113 @@
 import os
-import re
-import requests
-import torch
 import streamlit as st
-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  # Ensure this package is installed
+from langdetect import detect
+import torch
 
-# ✅ Check for GPU or Default to CPU
+# Check if GPU is available but don't load anything yet
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"✅ Using device: {device}")  # Debugging info
-
-# ✅ Environment Variables
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN is None:
-    raise ValueError("HF_TOKEN is not set. Please add it to your environment variables.")
-
-NASA_API_KEY = os.getenv("NASA_API_KEY")
-if NASA_API_KEY is None:
-    raise ValueError("NASA_API_KEY is not set. Please add it to your environment variables.")
-
-# ✅ Set Up Streamlit
 st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="🚀")
 
-# ✅ Initialize Session State Variables (Ensuring Chat History Persists)
+# Initialize session state variables
 if "chat_history" not in st.session_state:
-    st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
+    st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you with NASA-related information today?"}]
 
-# ✅ 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.3):
-    return HuggingFaceEndpoint(
-        repo_id=model_id,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,  # 🔥 Lowered temperature for more factual and structured responses
-        token=HF_TOKEN,
-        task="text-generation",
-        device=-1 if device == "cpu" else 0  # ✅ Force CPU (-1) or GPU (0)
-    )
+if "model_loaded" not in st.session_state:
+    st.session_state.model_loaded = False
 
-# ✅ Ensure English Responses
+# Load environment variables
+def load_api_keys():
+    hf_token = os.getenv("HF_TOKEN")
+    nasa_api_key = os.getenv("NASA_API_KEY")
+    
+    missing_keys = []
+    if not hf_token:
+        missing_keys.append("HF_TOKEN")
+    if not nasa_api_key:
+        missing_keys.append("NASA_API_KEY")
+    
+    return hf_token, nasa_api_key, missing_keys
+
+# Lazy-load the model only when needed
+def load_model():
+    with st.spinner("Loading AI model... This may take a moment."):
+        try:
+            from langchain_huggingface import HuggingFaceEndpoint
+            from langchain_core.prompts import PromptTemplate
+            from langchain_core.output_parsers import StrOutputParser
+            
+            hf_token, _, _ = load_api_keys()
+            
+            # Use a smaller model if you're having resource issues
+            llm = HuggingFaceEndpoint(
+                repo_id="meta-llama/Llama-2-7b-chat-hf",  # Consider a smaller model like "distilroberta-base"
+                max_new_tokens=800,
+                temperature=0.3,
+                token=hf_token,
+                task="text-generation",
+                device=-1 if device == "cpu" else 0
+            )
+            st.session_state.model_loaded = True
+            st.session_state.llm = llm
+            st.session_state.prompt = PromptTemplate.from_template(
+                "[INST] You are HAL, a NASA AI assistant with deep knowledge of space, astronomy, and NASA missions. "
+                "Answer concisely and accurately.\n\n"
+                "CONTEXT:\n{chat_history}\n"
+                "\nLATEST USER INPUT:\nUser: {user_text}\n"
+                "[END CONTEXT]\n"
+                "Assistant:"
+            )
+            return True
+        except Exception as e:
+            st.error(f"Error loading model: {str(e)}")
+            return False
+
+# Ensure English responses
 def ensure_english(text):
     try:
-        detected_lang = detect(text)
-        if detected_lang != "en":
-            return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
+        if text and len(text) > 5:  # Only check if there's meaningful text
+            detected_lang = detect(text)
+            if detected_lang != "en":
+                return "⚠️ Sorry, I only respond in English. Can you rephrase your question?"
+        return text
     except:
-        return "⚠️ Language detection failed. Please ask your question again."
-    return text
-
-# ✅ Main Response Function (Fixing Repetition & Context)
-def get_response(system_message, chat_history, user_text, max_new_tokens=800):
-    # ✅ Ensure conversation history is included correctly
-    filtered_history = "\n".join(
-        f"{msg['role'].capitalize()}: {msg['content']}"
-        for msg in chat_history[-5:]  # ✅ Only keep the last 5 exchanges to prevent overflow
-    )
- 
-    prompt = PromptTemplate.from_template(
-        "[INST] You are a highly knowledgeable AI assistant. Answer concisely, avoid repetition, and structure responses well."
-        "\n\nCONTEXT:\n{chat_history}\n"
-        "\nLATEST USER INPUT:\nUser: {user_text}\n"
-        "\n[END CONTEXT]\n"
-        "Assistant:"
-    )
+        return text  # Return original if detection fails
 
-    # ✅ Invoke Hugging Face Model
-    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.3)  # 🔥 Lowered temperature
-    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=filtered_history))
+# Get response from the model
+def get_response(user_text):
+    if not st.session_state.model_loaded:
+        if not load_model():
+            return "Sorry, I'm having trouble loading. Please try again or check your environment setup."
     
-    # Clean up the response - remove any "HAL:" prefix if present
-    response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
-    response = ensure_english(response)
-
-    if not response:
-        response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?"
-
-    # ✅ Update conversation history
-    chat_history.append({'role': 'user', 'content': user_text})
-    chat_history.append({'role': 'assistant', 'content': response})
-
-    # ✅ Keep only last 10 exchanges to prevent unnecessary repetition
-    return response, chat_history[-10:]
-
-# ✅ Streamlit UI
-st.title("🚀 HAL - NASA AI Assistant")
-
-# ✅ Justify all chatbot responses
+    try:
+        # Prepare conversation history
+        filtered_history = "\n".join(
+            f"{msg['role'].capitalize()}: {msg['content']}"
+            for msg in st.session_state.chat_history[-5:]
+        )
+        
+        from langchain_core.output_parsers import StrOutputParser
+        
+        # Create and invoke the chat pipeline
+        chat = st.session_state.prompt | st.session_state.llm.bind(skip_prompt=True) | StrOutputParser()
+        
+        response = chat.invoke({
+            "user_text": user_text,
+            "chat_history": filtered_history
+        })
+        
+        # Clean up response
+        response = response.split("HAL:")[-1].strip() if "HAL:" in response else response.strip()
+        response = ensure_english(response)
+        
+        if not response:
+            response = "I'm sorry, but I couldn't generate a response. Can you rephrase your question?"
+        
+        return response
+        
+    except Exception as e:
+        return f"I encountered an error: {str(e)}. Please try again with a different question."
+
+# UI Styling
 st.markdown("""
     <style>
     .user-msg, .assistant-msg {
@@ -100,29 +118,38 @@ st.markdown("""
         max-width: 80%;
         text-align: justify;
     }
-    .user-msg { background-color: #696969; color: white; }
+    .user-msg { background-color: #696969; color: white; margin-left: auto; }
     .assistant-msg { background-color: #333333; color: white; }
-    .container { display: flex; flex-direction: column; align-items: flex-start; }
+    .container { display: flex; flex-direction: column; }
     @media (max-width: 600px) { .user-msg, .assistant-msg { font-size: 16px; max-width: 100%; } }
     </style>
 """, unsafe_allow_html=True)
 
-# ✅ Chat UI
-user_input = st.chat_input("Type your message here...")
-
-if user_input:
-    # Get response and update chat history
-    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
-    )
+# Main UI
+st.title("🚀 HAL - NASA AI Assistant")
 
-# ✅ Display chat history (ONLY display from history, not separately)
-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)
+# Check for API keys before allowing interaction
+hf_token, nasa_api_key, missing_keys = load_api_keys()
+if missing_keys:
+    st.error(f"Missing environment variables: {', '.join(missing_keys)}. Please set them to use this application.")
+else:
+    # Chat interface
+    user_input = st.chat_input("Ask me about NASA, space missions, or astronomy...")
+    
+    if user_input:
+        # Add user message to history
+        st.session_state.chat_history.append({"role": "user", "content": user_input})
+        
+        # Get AI response
+        with st.spinner("Thinking..."):
+            response = get_response(user_input)
+            st.session_state.chat_history.append({"role": "assistant", "content": response})
+    
+    # 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)