Update app.py

app.py

@@ -37,85 +37,121 @@ index = faiss.IndexFlatL2(embedding_dim)
 def preprocess_query(query):
     tokens = query.lower().split()
     epilepsy_keywords = ["seizure", "epilepsy", "convulsion", "neurology", "brain activity"]
-    healthcare_keywords = ["headache", "fever", "blood pressure", "diabetes", "cough", "flu", "nutrition", "mental health", "pain", "legs", "body pain"] # Added "pain" and "legs" and "body pain"
 
     is_epilepsy_related = any(k in tokens for k in epilepsy_keywords)
-    is_healthcare_related = any(k in tokens for k in healthcare_keywords) and not is_epilepsy_related
 
-    return tokens, is_epilepsy_related, is_healthcare_related
+    return tokens, is_epilepsy_related
 
 # Function to Generate Response with Chat History
 def generate_response(user_query, chat_history):
-    tokens, is_epilepsy_related, is_healthcare_related = preprocess_query(user_query)
-
-    # Greeting Responses
-    greetings = ["hello", "hi", "hey"]
-    if any(word in tokens for word in greetings):
-        return "👋 Hello! How can I assist you today?"
-
-    # If Healthcare Related but Not Epilepsy - Provide General Wellness Tips
-    if is_healthcare_related:
-        general_health_tips = (
-            "For general health and well-being:\n"
-            "- 💧 Stay hydrated by drinking plenty of water throughout the day.\n"
-            "- 🍎 Maintain a balanced diet rich in fruits, vegetables, and whole grains.\n"
-            "- 🚶‍♀️ Incorporate regular physical activity into your daily routine.\n"
-            "- 😴 Ensure you get adequate sleep to allow your body to rest and recover.\n"
-            "- 🧘 Practice stress-reducing activities such as deep breathing or meditation.\n"
-            "- 🩺 **Important:** These tips are for general wellness. Always consult a healthcare professional for any specific health concerns or before making significant changes to your health regimen."
-        )
-        return (
-            f"**NeuroGuard:** 🩺 It sounds like your question '{user_query}' is about general health. While I specialize in epilepsy, here are some general wellness tips that might be helpful:\n\n"
-            f"{general_health_tips}"
-        )
-
-    # If Not Healthcare or Epilepsy Related
-    if not is_epilepsy_related:
-        return (
-            f"**NeuroGuard:** 💡 Your query '{user_query}' doesn't seem to be directly related to epilepsy or general healthcare. \n\n"
-            "For general information:\n"
-            "- 📚 Always seek information from reliable sources and consult experts when needed."
-        )
-
-    # Try Getting Medical Insights from PubMedBERT
+    # Grammatical Correction using LLaMA (Hidden from User)
     try:
-        pubmedbert_embeddings = pubmedbert_pipeline(user_query)
-        embedding_mean = np.mean(pubmedbert_embeddings[0], axis=0)
-        index.add(np.array([embedding_mean]))
-        pubmedbert_insights = "**PubMedBERT Analysis:** ✅ Query is relevant to epilepsy research."
-    except Exception as e:
-        pubmedbert_insights = f"⚠️ Error during PubMedBERT analysis: {e}"
-
-    # Use LLaMA for Final Response Generation with Chat History Context
-    try:
-        prompt_history = ""
-        if chat_history:
-            prompt_history += "**Chat History:**\n"
-            for message in chat_history:
-                prompt_history += f"{message['role'].capitalize()}: {message['content']}\n"
-            prompt_history += "\n"
-
-        prompt = f"""
-        {prompt_history}
-        **User Query:** {user_query}
-        **Instructions:** Provide a concise, structured, and human-friendly response specifically about epilepsy or seizures, considering the conversation history if available.
+        correction_prompt = f"""
+        Correct the following user query for grammar and spelling errors, but keep the original intent intact.
+        Do not add or remove any information, just fix the grammar.
+        User Query: {user_query}
+        Corrected Query:
         """
-
-        chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
+        grammar_completion = client.chat.completions.create(
+            messages=[{"role": "user", "content": correction_prompt}],
             model="llama-3.3-70b-versatile",
             stream=False,
         )
-        model_response = chat_completion.choices[0].message.content.strip()
+        corrected_query = grammar_completion.choices[0].message.content.strip()
+        # If correction fails or returns empty, use original query
+        if not corrected_query:
+            corrected_query = user_query
     except Exception as e:
-        model_response = f"⚠️ Error generating response with LLaMA: {e}"
+        corrected_query = user_query # Fallback to original query if correction fails
+        print(f"⚠️ Grammar correction error: {e}") # Optional: Log the error for debugging
+
+    tokens, is_epilepsy_related = preprocess_query(corrected_query) # Use corrected query for processing
+
+    # Greeting Responses
+    greetings = ["hello", "hi", "hey"]
+    if any(word in tokens for word in greetings):
+        return "👋 Hello! How can I assist you today?"
+
+    # If Epilepsy Related - Use Epilepsy Focused Response
+    if is_epilepsy_related:
+        # Try Getting Medical Insights from PubMedBERT
+        try:
+            pubmedbert_embeddings = pubmedbert_pipeline(corrected_query) # Use corrected query for PubMedBERT
+            embedding_mean = np.mean(pubmedbert_embeddings[0], axis=0)
+            index.add(np.array([embedding_mean]))
+            pubmedbert_insights = "**PubMedBERT Analysis:** ✅ Query is relevant to epilepsy research."
+        except Exception as e:
+            pubmedbert_insights = f"⚠️ Error during PubMedBERT analysis: {e}"
+
+        # Use LLaMA for Final Response Generation with Chat History Context (Epilepsy Focus)
+        try:
+            prompt_history = ""
+            if chat_history:
+                prompt_history += "**Chat History:**\n"
+                for message in chat_history:
+                    prompt_history += f"{message['role'].capitalize()}: {message['content']}\n"
+                prompt_history += "\n"
+
+            epilepsy_prompt = f"""
+            {prompt_history}
+            **User Query:** {corrected_query} # Use corrected query for final response generation
+            **Instructions:** Provide a concise, structured, and human-friendly response specifically about epilepsy or seizures, considering the conversation history if available.
+            """
+
+            chat_completion = client.chat.completions.create(
+                messages=[{"role": "user", "content": epilepsy_prompt}],
+                model="llama-3.3-70b-versatile",
+                stream=False,
+            )
+            model_response = chat_completion.choices[0].message.content.strip()
+        except Exception as e:
+            model_response = f"⚠️ Error generating response with LLaMA: {e}"
+
+        return f"**NeuroGuard:** ✅ **Analysis:**\n{pubmedbert_insights}\n\n**Response:**\n{model_response}"
+
+
+    # If Not Epilepsy Related - Try to Answer as General Health Query
+    else:
+        # Try Getting Medical Insights from PubMedBERT (even for general health)
+        try:
+            pubmedbert_embeddings = pubmedbert_pipeline(corrected_query)
+            embedding_mean = np.mean(pubmedbert_embeddings[0], axis=0)
+            index.add(np.array([embedding_mean]))
+            pubmedbert_insights = "**PubMedBERT Analysis:**  PubMed analysis performed for health-related context." # General analysis message
+        except Exception as e:
+            pubmedbert_insights = f"⚠️ Error during PubMedBERT analysis: {e}"
+
+        # Use LLaMA for General Health Response Generation with Chat History Context
+        try:
+            prompt_history = ""
+            if chat_history:
+                prompt_history += "**Chat History:**\n"
+                for message in chat_history:
+                    prompt_history += f"{message['role'].capitalize()}: {message['content']}\n"
+                prompt_history += "\n"
+
+            general_health_prompt = f"""
+            {prompt_history}
+            **User Query:** {corrected_query}
+            **Instructions:** Provide a concise, structured, and human-friendly response to the general health query, considering the conversation history if available. If the query is clearly not health-related, respond generally.
+            """
+
+            chat_completion = client.chat.completions.create(
+                messages=[{"role": "user", "content": general_health_prompt}],
+                model="llama-3.3-70b-versatile",
+                stream=False,
+            )
+            model_response = chat_completion.choices[0].message.content.strip()
+        except Exception as e:
+            model_response = f"⚠️ Error generating response with LLaMA: {e}"
+
+        return f"**NeuroGuard:** ✅ **Analysis:**\n{pubmedbert_insights}\n\n**Response:**\n{model_response}"
 
-    return f"**NeuroGuard:** ✅ **Analysis:**\n{pubmedbert_insights}\n\n**Response:**\n{model_response}"
 
 # Streamlit UI Setup
-st.set_page_config(page_title="NeuroGuard: Epilepsy Chatbot", layout="wide")
-st.title("🧠 NeuroGuard: Epilepsy & Seizure Chatbot")
-st.write("💬 Ask me anything about epilepsy and seizures. I remember our conversation!")
+st.set_page_config(page_title="NeuroGuard: Epilepsy & Health Chatbot", layout="wide") # Updated title
+st.title("🧠 NeuroGuard: Epilepsy & Health Chatbot") # Updated title
+st.write("💬 Ask me anything about epilepsy, seizures, and general health. I remember our conversation!") # Updated description
 
 # Initialize Chat History in Session State
 if "chat_history" not in st.session_state: