Upload UI.py

UI.py

@@ -6,64 +6,85 @@ st.set_page_config(page_title="Hospital AI Assistant", layout="wide")
 st.title("🏥 Hospital AI Assistant")
 st.caption("Ask questions about patients, conditions, visits, medications, labs")
 
-# Initialize chat history
+# =========================
+# Session State
+# =========================
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat history
+# =========================
+# Show Chat History
+# =========================
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+        st.markdown(msg["content"], unsafe_allow_html=True)
 
-# Chat input
+# =========================
+# Chat Input
+# =========================
 user_input = st.chat_input("Ask a question about hospital data...")
 
 if user_input:
     # Show user message
-    st.session_state.messages.append({"role": "user", "content": user_input})
+    st.session_state.messages.append(
+        {"role": "user", "content": user_input}
+    )
     with st.chat_message("user"):
         st.markdown(user_input)
 
-    # Call AI engine directly
+    # Call backend
     with st.spinner("Thinking..."):
         try:
             result = process_question(user_input)
         except Exception as e:
-            result = {"status": "error", "message": str(e)}
+            result = {
+                "status": "error",
+                "message": str(e)
+            }
+
+    # =========================
+    # Build Assistant Reply
+    # =========================
+    reply = ""
 
-    # Build assistant reply
     if result.get("status") == "ok":
-        reply = ""
 
-        # Time note (if any)
+        # Message (important for time-based or empty data)
+        if result.get("message"):
+            reply += f"❗ {result['message']}\n\n"
+
+        # Time note
         if result.get("note"):
-            reply += f"🕒 *{result['note']}*\n\n"
+            reply += f"🕒 {result['note']}\n\n"
 
-        # Data table
+        # Table output
         if result.get("data"):
             columns = result.get("columns", [])
             data = result["data"]
 
-            table_md = "| " + " | ".join(columns) + " |\n"
-            table_md += "| " + " | ".join(["---"] * len(columns)) + " |\n"
+            reply += "### Result\n"
+            reply += "| " + " | ".join(columns) + " |\n"
+            reply += "| " + " | ".join(["---"] * len(columns)) + " |\n"
 
             for row in data[:10]:
-                table_md += "| " + " | ".join(str(x) for x in row) + " |\n"
+                reply += "| " + " | ".join(str(x) for x in row) + " |\n"
 
-            reply += table_md
-        else:
-            reply += result.get("message", "No data found.")
-
-        # SQL toggle
-        reply += "\n\n---\n"
-        reply += "<details><summary><b>Generated SQL</b></summary>\n\n"
-        reply += f"```sql\n{result['sql']}\n```"
-        reply += "\n</details>"
+        # SQL (ONLY if present)
+        if result.get("sql"):
+            reply += "\n---\n"
+            reply += "<details><summary><b>Generated SQL</b></summary>\n\n"
+            reply += f"```sql\n{result['sql']}\n```"
+            reply += "\n</details>"
 
     else:
         reply = f"❌ {result.get('message', 'Something went wrong')}"
 
-    # Show assistant message
-    st.session_state.messages.append({"role": "assistant", "content": reply})
+    # =========================
+    # Display Assistant Message
+    # =========================
+    st.session_state.messages.append(
+        {"role": "assistant", "content": reply}
+    )
+
     with st.chat_message("assistant"):
-        st.markdown(reply)
+        st.markdown(reply, unsafe_allow_html=True)