Upload UI.py

UI.py

@@ -3,12 +3,15 @@ from engine import process_question
 from datetime import datetime
 
 # =========================
-# PAGE SETUP
+# PAGE CONFIG
 # =========================
-st.set_page_config(page_title="Hospital AI Assistant", layout="wide")
+st.set_page_config(
+    page_title="Text-to-SQL AI",
+    layout="wide"
+)
 
-st.title("🏥 Hospital AI Assistant")
-st.caption("Ask questions about patients, conditions, visits, medications, labs")
+st.title("🧠 Text-to-SQL Assistant")
+st.caption("Ask questions in natural language. I’ll generate SQL using your database metadata.")
 
 # =========================
 # SESSION STATE
@@ -20,19 +23,19 @@ if "transcript" not in st.session_state:
     st.session_state.transcript = []
 
 # =========================
-# SHOW CHAT HISTORY
+# CHAT HISTORY
 # =========================
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"], unsafe_allow_html=True)
 
 # =========================
-# CHAT INPUT
+# USER INPUT
 # =========================
-user_input = st.chat_input("Ask a question about hospital data...")
+user_input = st.chat_input("Ask something like: 'Show employees in IT department'")
 
 if user_input:
-    # Show user message
+    # Display user message
     st.session_state.messages.append(
         {"role": "user", "content": user_input}
     )
@@ -40,8 +43,8 @@ if user_input:
     with st.chat_message("user"):
         st.markdown(user_input)
 
-    # Call backend
-    with st.spinner("Thinking..."):
+    # Call engine
+    with st.spinner("Generating SQL..."):
         try:
             result = process_question(user_input)
         except Exception as e:
@@ -56,48 +59,26 @@ if user_input:
     reply = ""
 
     if result.get("status") == "ok":
-
         if result.get("message"):
-            reply += f"❗ {result['message']}\n\n"
-
-        if result.get("note"):
-            reply += f"🕒 {result['note']}\n\n"
-
-        if result.get("data"):
-            columns = result.get("columns", [])
-            data = result["data"]
-
-            reply += "### Result\n"
-            reply += "| " + " | ".join(columns) + " |\n"
-            reply += "| " + " | ".join(["---"] * len(columns)) + " |\n"
-
-            for row in data[:10]:
-                reply += "| " + " | ".join(str(x) for x in row) + " |\n"
+            reply += f"### ✅ Result\n{result['message']}\n\n"
 
         if result.get("sql"):
-            reply += "\n---\n"
-            reply += "<details><summary><b>Generated SQL</b></summary>\n\n"
+            reply += "### 🧾 Generated SQL\n"
             reply += f"```sql\n{result['sql']}\n```"
-            reply += "\n</details>"
 
     else:
-        reply = f"❌ {result.get('message', 'Something went wrong')}"
+        reply = f"❌ **Error:** {result.get('message')}"
 
-    # =========================
-    # SAVE TRANSCRIPT ✅ FIXED
-    # =========================
+    # Save transcript
     st.session_state.transcript.append({
         "timestamp": datetime.utcnow().isoformat(),
         "question": user_input,
-        "reply": reply,            # ✅ THIS WAS MISSING
+        "reply": reply,
         "sql": result.get("sql"),
-        "result_preview": result.get("data", [])[:10],
         "error": result.get("message") if result.get("status") != "ok" else None
     })
 
-    # =========================
-    # SHOW ASSISTANT MESSAGE
-    # =========================
+    # Show assistant reply
     st.session_state.messages.append(
         {"role": "assistant", "content": reply}
     )
@@ -106,34 +87,25 @@ if user_input:
         st.markdown(reply, unsafe_allow_html=True)
 
 # =========================
-# DOWNLOAD TRANSCRIPT
+# DOWNLOAD LOG
 # =========================
-def download_transcript_txt():
+def download_transcript():
     lines = []
-    for i, entry in enumerate(st.session_state.transcript, 1):
+    for i, t in enumerate(st.session_state.transcript, 1):
         lines.append(f"\n--- Query {i} ---")
-        lines.append(f"Time: {entry['timestamp']}")
-        lines.append(f"Question: {entry['question']}")
-
-        if entry.get("reply"):
-            lines.append("Reply:")
-            lines.append(entry["reply"])
-
-        if entry.get("sql"):
-            lines.append(f"SQL:\n{entry['sql']}")
-
-        if entry.get("error"):
-            lines.append(f"Error: {entry['error']}")
-
+        lines.append(f"Time: {t['timestamp']}")
+        lines.append(f"Question: {t['question']}")
+        lines.append(f"Reply:\n{t['reply']}")
+        if t.get("sql"):
+            lines.append(f"SQL:\n{t['sql']}")
     return "\n".join(lines)
 
-# Show download button only if data exists
 if st.session_state.transcript:
     st.divider()
     st.download_button(
-        label="⬇️ Download conversation",
-        data=download_transcript_txt(),
-        file_name="chat_transcript.txt",
+        "⬇️ Download Query Log",
+        data=download_transcript(),
+        file_name="text_to_sql_log.txt",
         mime="text/plain",
         use_container_width=True
     )