Update src/streamlit_app.py

src/streamlit_app.py

@@ -9,7 +9,7 @@ import torch
 st.markdown(
     """
     <style>
-    /* اجعل كل النصوص من اليمين إلى اليسار */
+    /* RTL styling for inputs/buttons */
     .stTextArea, .stTextInput, .stButton {
         direction: rtl;
         text-align: right;
@@ -54,7 +54,36 @@ def classify_question(text):
     return LABELS[pred]
 
 # =======================
-# Chatbot Response
+# LLM first-aid double-check
+# =======================
+def llm_first_aid_check(query, api_key):
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    prompt = f"""
+هل السؤال التالي يندرج ضمن الإسعافات الأولية؟ أجب فقط بـ "نعم" أو "لا":
+{query}
+"""
+    data = {
+        "model": "mistralai/mistral-7b-instruct",
+        "messages": [
+            {"role": "system", "content": "أنت مساعد ذكي يساعد في تصنيف الأسئلة الطبية."},
+            {"role": "user", "content": prompt}
+        ],
+        "temperature": 0,
+        "max_tokens": 10
+    }
+    response = requests.post(url, headers=headers, json=data)
+    if response.status_code == 200:
+        answer = response.json()["choices"][0]["message"]["content"].strip()
+        return "نعم" in answer
+    else:
+        return False
+
+# =======================
+# Chatbot response
 # =======================
 def get_mistral_response(prompt, api_key):
     url = "https://openrouter.ai/api/v1/chat/completions"
@@ -100,18 +129,27 @@ if st.button("اسأل"):
         # Add user message
         st.session_state.messages.append({"role": "user", "content": user_input})
 
-        # Classification
+        # First classifier check
         with st.spinner("جار تحليل سؤالك..."):
             category = classify_question(user_input)
 
         if category == "LABEL_1":
-            st.success("تم التعرف على سؤال إسعافات أولية. نحصل الآن على الإرشادات ✅")
-            # Get bot response
-            with st.spinner("جارٍ التواصل مع المساعد الذكي..."):
-                answer = get_mistral_response(user_input, api_key)
-                st.session_state.messages.append({"role": "bot", "content": answer})
+            # LLM double-check
+            with st.spinner("جارٍ التحقق من صحة السؤال..."):
+                is_first_aid = llm_first_aid_check(user_input, api_key)
+
+            if is_first_aid:
+                st.success("تم التأكد من أن السؤال يتعلق بالإسعافات الأولية ✅")
+                # Get bot response
+                with st.spinner("جارٍ التواصل مع المساعد الذكي..."):
+                    answer = get_mistral_response(user_input, api_key)
+                    st.session_state.messages.append({"role": "bot", "content": answer})
+            else:
+                st.warning("يبدو أن السؤال لا يتعلق بالإسعافات الأولية ❌")
+                st.session_state.messages.append({"role": "bot", "content": "عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية."})
         else:
-            st.session_state.messages.append({"role": "bot", "content": "عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية❌"})
+            st.error("عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية❌")
+            st.session_state.messages.append({"role": "bot", "content": "عذرًا، يمكنني الإجابة فقط على الأسئلة المتعلقة بالإسعافات الأولية."})
 
 # =======================
 # Display chat