Upload 2 files

UI.py

@@ -1,121 +1,102 @@
-import streamlit as st
-from engine import process_question,download_transcript_txt,download_transcript_json
-
-
-
-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")
-
-# =========================
-# Session State
-# =========================
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# =========================
-# Show 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 = 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}
-    )
-    with st.chat_message("user"):
-        st.markdown(user_input)
-
-    # Call backend
-    with st.spinner("Thinking..."):
-        try:
-            result = process_question(user_input)
-        except Exception as e:
-            result = {
-                "status": "error",
-                "message": str(e)
-            }
-
-    # =========================
-    # Build Assistant Reply
-    # =========================
-    reply = ""
-
-    if result.get("status") == "ok":
-
-        # 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"
-
-        # Table output
-        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"
-
-        # 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')}"
-
-    # =========================
-    # Display Assistant Message
-    # =========================
-    st.session_state.messages.append(
-        {"role": "assistant", "content": reply}
-    )
-
-    with st.chat_message("assistant"):
-        st.markdown(reply, unsafe_allow_html=True)
-# =========================
-# Download Conversation
-# =========================
-st.divider()
-st.subheader("📥 Download Conversation")
-
-# Create 3 columns: spacer | TXT | JSON
-col_spacer, col_txt, col_json = st.columns([6, 2, 2])
-
-with col_txt:
-    txt = download_transcript_txt()
-    st.download_button(
-        label="📄 TXT",
-        data=txt,
-        file_name="chat_transcript.txt",
-        mime="text/plain",
-        use_container_width=True
-    )
-
-with col_json:
-    js = download_transcript_json()
-    st.download_button(
-        label="🧾 JSON",
-        data=js,
-        file_name="chat_transcript.json",
-        mime="application/json",
-        use_container_width=True
-    )
-
+import streamlit as st
+from engine import process_question,download_transcript_txt,download_transcript_json
+
+
+
+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")
+
+# =========================
+# Session State
+# =========================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# =========================
+# Show 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 = 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}
+    )
+    with st.chat_message("user"):
+        st.markdown(user_input)
+
+    # Call backend
+    with st.spinner("Thinking..."):
+        try:
+            result = process_question(user_input)
+        except Exception as e:
+            result = {
+                "status": "error",
+                "message": str(e)
+            }
+
+    # =========================
+    # Build Assistant Reply
+    # =========================
+    reply = ""
+
+    if result.get("status") == "ok":
+
+        # 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"
+
+        # Table output
+        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"
+
+        # 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')}"
+
+    # =========================
+    # Display Assistant Message
+    # =========================
+    st.session_state.messages.append(
+        {"role": "assistant", "content": reply}
+    )
+
+    with st.chat_message("assistant"):
+        st.markdown(reply, unsafe_allow_html=True)
+        
+        # Download button inside chat bubble
+        st.download_button(
+        label="⬇️ Download conversation",
+        data=download_transcript_txt(),
+        file_name="chat_transcript.txt",
+        mime="text/plain",
+        use_container_width=True
+    )
+

engine.py

@@ -7,6 +7,16 @@ from datetime import datetime
 
 TRANSCRIPT = []
 
+def log_interaction(user_q, sql=None, result=None, error=None):
+    TRANSCRIPT.append({
+        "timestamp": datetime.utcnow().isoformat(),
+        "question": user_q,
+        "sql": sql,
+        "result_preview": result[:10] if isinstance(result, list) else result,
+        "error": error
+    })
+
+
 
 # =========================
 # SETUP
@@ -628,15 +638,6 @@ def build_table_summary(table_name):
 
 def process_question(question):
     
-    def log_interaction(user_q, sql=None, result=None, error=None):
-        TRANSCRIPT.append({
-        "timestamp": datetime.utcnow().isoformat(),
-        "question": user_q,
-        "sql": sql,
-        "result_preview": result[:10] if isinstance(result, list) else result,
-        "error": error
-    })
-
     global LAST_PROMPT_TYPE, LAST_SUGGESTED_DATE
 
     q = question.strip().lower()
@@ -758,11 +759,27 @@ def process_question(question):
     sql = validate_sql(sql)
     cols, rows = run_query(sql)
     
+    # ✅ LOG ONCE (THIS FIXES YOUR DOWNLOAD ISSUE)
     log_interaction(
-    user_q=question,
-    sql=sql,
-    result=rows
-)
+        user_q=question,
+        sql=sql,
+        result=rows
+    )
+
+    if not rows:
+        return {
+            "status": "ok",
+            "message": friendly("No records found."),
+            "data": []
+        }
+
+    return {
+        "status": "ok",
+        "sql": sql,
+        "columns": cols,
+        "data": rows
+    }
+
 
 
     # ----------------------------------