Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,134 +1,126 @@
 import streamlit as st
 import requests
 import uuid
+import time
 import os
-from dotenv import load_dotenv
 
-load_dotenv()
+# --- CONFIGURATION ---
+# Your FastAPI server URL on Hugging Face
+API_URL = "https://lightrt-text2sql-backend.hf.space" 
 
+st.set_page_config(page_title="Text2SQL Agent", page_icon="🤖", layout="centered")
+
+# --- HUGGING FACE FIREWALL BYPASS (CHROME SPOOFING) ---
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# We are disguising the request to look exactly like a Google Chrome browser
 HEADERS = {
     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
     "Accept": "application/json",
     "Content-Type": "application/json"
 }
 
-# Add the token if it exists
 if HF_TOKEN:
     HEADERS["Authorization"] = f"Bearer {HF_TOKEN}"
 
-# 1. PAGE CONFIGURATION
-st.set_page_config(page_title="Text2SQL Agent", page_icon="🤖", layout="centered")
-st.title("🗣️ Text2SQL Agent")
-st.markdown("Connect your database and ask questions about your data in plain English.")
+# (Optional) A quick UI check so you know the token loaded in the cloud
+if not HF_TOKEN:
+    st.sidebar.error("⚠️ HF_TOKEN missing from Space Secrets!")
+else:
+    st.sidebar.success("🔑 Token Loaded! Ready for cloud connection.")
 
-# 2. SESSION STATE INITIALIZATION (The Memory Bank)
-if "user_id" not in st.session_state:
-    st.session_state.user_id = "tenant_" + str(uuid.uuid4())[:8]
+
+# --- SESSION STATE INITIALIZATION ---
 if "thread_id" not in st.session_state:
     st.session_state.thread_id = str(uuid.uuid4())
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+if "user_id" not in st.session_state:
+    st.session_state.user_id = "tenant_" + str(uuid.uuid4())[:8]
 if "is_db_connected" not in st.session_state:
     st.session_state.is_db_connected = False
 if "connection_url" not in st.session_state:
     st.session_state.connection_url = ""
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
 
-# 3. SIDEBAR: DATABASE CONNECTION (The Handoff to FastAPI)
+# --- SIDEBAR: DATABASE CONNECTION ---
 with st.sidebar:
     st.header("⚙️ Database Setup")
     
-    # Text input for the database URL
-    db_input = st.text_input("Enter Database URL:", disabled=st.session_state.is_db_connected)
+    db_input = st.text_input(
+        "Enter Database URL:",
+        disabled=st.session_state.is_db_connected
+    )
     
     if not st.session_state.is_db_connected:
-        if st.button("Connect & Initialize"):
-            if db_input:
-                with st.spinner("Connecting to backend..."):
-                    payload_data = {
-                        "connection_url": db_input, 
-                        "user_id": st.session_state.user_id
-                    }
-                    
+        if st.button("Connect & Initialize", type="primary", use_container_width=True):
+            if not db_input:
+                st.error("Please enter a valid URL.")
+            else:
+                with st.spinner("Building embeddings and initializing agent..."):
                     try:
-                        # Make sure your backend endpoint is actually /upload_url 
-                        # (Check your FastAPI main.py to confirm this endpoint name!)
-                        response = requests.post(
-                            "https://lightrt-text2sql-backend.hf.space/upload_url", 
-                            json=payload_data,
-                            headers=HEADERS
-                        )
+                        payload = {"connection_url": db_input, "user_id": st.session_state.user_id}
+                        
+                        # Added headers=HEADERS here!
+                        response = requests.post(f"{API_URL}/upload_url", json=payload, headers=HEADERS)
                         
                         if response.status_code == 200:
                             st.session_state.is_db_connected = True
                             st.session_state.connection_url = db_input
+                            
+                            time.sleep(15) 
+                            
                             st.success("Database connected securely!")
-                            st.rerun() # Refresh to update the sidebar UI
+                            st.rerun()
+                        elif response.status_code == 429:
+                            st.error("🚨 Rate Limited! Waiting for HF firewall cooldown.")
                         else:
-                            st.error(f"Connection failed: {response.text}")
-                            
+                            st.error(f"Failed to connect: {response.text}")
                     except requests.exceptions.ConnectionError:
-                        st.error("Cannot connect to backend. Is FastAPI running?")
-            else:
-                st.warning("Please enter a valid URL first.")
+                        st.error("🚨 Cannot connect to backend. Is FastAPI running?")
     else:
         st.success("✅ Connected to Database")
         st.caption(f"URL: {st.session_state.connection_url}")
-        if st.button("Disconnect & Reset"):
+        
+        if st.button("Disconnect & Reset", use_container_width=True):
             st.session_state.clear()
             st.rerun()
 
-# 4. CHAT HISTORY RENDERING
-for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
+# --- MAIN CHAT INTERFACE ---
+st.title("🗣️ Text2SQL Agent")
 
-# 5. CHAT INPUT & BACKEND COMMUNICATION
-# Only show chat input if the database is connected
-if st.session_state.is_db_connected:
-    if prompt := st.chat_input("Ask a question about your data..."):
+if not st.session_state.is_db_connected:
+    st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
+else:
+    for msg in st.session_state.chat_history:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+
+    if user_query := st.chat_input("Ask a question about your data..."):
         
-        # Immediately render the user's new message to the UI
-        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.session_state.chat_history.append({"role": "user", "content": user_query})
         with st.chat_message("user"):
-            st.markdown(prompt)
-            
-        # Show a loading indicator while we wait for FastAPI
+            st.markdown(user_query)
+
         with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-            message_placeholder.markdown("*(Analyzing schema and generating SQL...)*")
-            
-            # Prepare the JSON payload for FastAPI
-            payload = {
-                "message": prompt,
-                "user_id": st.session_state.user_id,
-                "thread_id": st.session_state.thread_id,
-                "connection_url": st.session_state.connection_url
-            }
-            
-            try:
-                # Send the question to your LangGraph backend
-                chat_response = requests.post(
-                    "https://lightrt-text2sql-backend.hf.space/chat", 
-                    json=payload,
-                    headers=HEADERS
-                )
-                
-                if chat_response.status_code == 200:
-                    # Extract the answer from the JSON response
-                    answer = chat_response.json().get("response", "No response found.")
-                    
-                    # Update the UI placeholder with the actual answer
-                    message_placeholder.markdown(answer)
+            with st.spinner("Analyzing schema and generating SQL..."):
+                try:
+                    payload = {
+                        "message": user_query,
+                        "thread_id": st.session_state.thread_id,
+                        "user_id": st.session_state.user_id,
+                        "connection_url": st.session_state.connection_url
+                    }
                     
-                    # Save the AI's answer to the session state memory
-                    st.session_state.messages.append({"role": "assistant", "content": answer})
-                else:
-                    message_placeholder.error(f"Error: {chat_response.text}")
+                    # Added headers=HEADERS here!
+                    response = requests.post(f"{API_URL}/chat", json=payload, headers=HEADERS)
                     
-            except requests.exceptions.ConnectionError:
-                message_placeholder.error("Cannot connect to backend. Is FastAPI running?")
-else:
-    st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
+                    if response.status_code == 200:
+                        answer = response.json().get("response", "No response found.")
+                        st.markdown(answer)
+                        st.session_state.chat_history.append({"role": "assistant", "content": answer})
+                    elif response.status_code == 429:
+                        st.error("🚨 Rate Limited by Hugging Face firewall.")
+                    else:
+                        st.error(f"Agent Error: {response.text}")
+                        
+                except requests.exceptions.ConnectionError:
+                    st.error("🚨 Connection dropped. Ensure FastAPI is running.")