Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,16 +2,24 @@ import streamlit as st
 import requests
 import uuid
 import time
+import os
 
 # --- CONFIGURATION ---
-API_URL = "https://lightrt-text2sql-backend.hf.space" # Your FastAPI server URL
+# Replace this with your exact direct backend Space URL (no trailing slash)
+API_URL = "https://lightrt-text2sql-backend.hf.space" 
 
-st.set_page_config(page_title="Text@SQL Agent", page_icon="🤖", layout="centered")
+st.set_page_config(page_title="Text2SQL Agent", page_icon="🤖", layout="centered")
+
+# --- HEADERS TO BYPASS HUGGING FACE PROXY FILTERS ---
+# Adding a definitive User-Agent prevents Hugging Face from flagging server-to-server calls as scraping bots
+HEADERS = {
+    "User-Agent": "Text2SQL-Frontend-Service",
+    "Content-Type": "application/json"
+}
 
 # --- SESSION STATE INITIALIZATION ---
-# This ensures variables survive when Streamlit re-renders the page
 if "thread_id" not in st.session_state:
-    st.session_state.thread_id = str(uuid.uuid4()) # Unique session ID for LangGraph memory
+    st.session_state.thread_id = str(uuid.uuid4())
 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:
@@ -25,7 +33,6 @@ if "chat_history" not in st.session_state:
 with st.sidebar:
     st.header("⚙️ Database Setup")
     
-    # If already connected, disable the input to enforce ONE database connection
     db_input = st.text_input(
         "Enter Database URL:",
         disabled=st.session_state.is_db_connected
@@ -38,31 +45,37 @@ with st.sidebar:
             else:
                 with st.spinner("Building embeddings and initializing agent..."):
                     try:
-                        # 1. Hit your FastAPI upload endpoint
                         payload = {"connection_url": db_input, "user_id": st.session_state.user_id}
-                        response = requests.post(f"{API_URL}/upload_url", json=payload)
+                        
+                        # Added HEADERS here
+                        response = requests.post(
+                            f"{API_URL}/upload_url", 
+                            json=payload, 
+                            headers=HEADERS,
+                            timeout=30
+                        )
                         
                         if response.status_code == 200:
-                            # 2. Lock the connection and unlock the chat
                             st.session_state.is_db_connected = True
                             st.session_state.connection_url = db_input
                             
-                            # Because your FastAPI upload uses BackgroundTasks, it returns instantly. 
-                            # We add a 2-second UI buffer here so the Qdrant embeddings have time to finish 
-                            # before the user fires off their first chat question.
+                            # Keeping your buffer for Qdrant storage compilation
                             time.sleep(15) 
                             
                             st.success("Database connected securely!")
-                            st.rerun() # Refresh UI to unlock the chat window
+                            st.rerun() 
+                        elif response.status_code == 429:
+                            st.error("🚨 Hugging Face is rate-limiting the connection. Please wait 2-3 minutes before clicking again.")
                         else:
-                            st.error(f"Failed to connect: {response.text}")
+                            st.error(f"Failed to connect: {response.status_code} - {response.text}")
                     except requests.exceptions.ConnectionError:
-                        st.error("🚨 Cannot connect to backend. Is FastAPI running?")
+                        st.error("🚨 Cannot connect to backend. Verify your API_URL is correct and the Space is running.")
+                    except requests.exceptions.Timeout:
+                        st.error("🚨 Connection timed out. The backend server took too long to reply.")
     else:
         st.success("✅ Connected to Database")
         st.caption(f"URL: {st.session_state.connection_url}")
         
-        # Add a reset button just in case they want to start completely over
         if st.button("Disconnect & Reset", use_container_width=True):
             st.session_state.clear()
             st.rerun()
@@ -70,45 +83,51 @@ with st.sidebar:
 # --- MAIN CHAT INTERFACE ---
 st.title("🗣️ Text2SQL Agent")
 
-# The Lock: Do not render the chat if DB is not connected
 if not st.session_state.is_db_connected:
     st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
 else:
-    # 1. Display previous chat messages from session state
+    # Display message history safely
     for msg in st.session_state.chat_history:
         with st.chat_message(msg["role"]):
             st.markdown(msg["content"])
 
-    # 2. The Chat Input box
+    # Chat Input Box
     if user_query := st.chat_input("Ask a question about your data..."):
         
-        # Immediately display the user's question in the UI
-        st.session_state.chat_history.append({"role": "user", "content": user_query})
-        with st.chat_message("user"):
-            st.markdown(user_query)
+        # Guard against empty submissions triggering extra calls
+        if user_query.strip() != "":
+            st.session_state.chat_history.append({"role": "user", "content": user_query})
+            with st.chat_message("user"):
+                st.markdown(user_query)
 
-        # 3. Call the LangGraph Backend
-        with st.chat_message("assistant"):
-            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
-                    }
-                    
-                    response = requests.post(f"{API_URL}/chat", json=payload)
-                    
-                    if response.status_code == 200:
-                        # Extract the final_result from your FastAPI JSON response
-                        answer = response.json().get("response", "No response found.")
-                        st.markdown(answer)
+            with st.chat_message("assistant"):
+                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 assistant's answer to the UI history
-                        st.session_state.chat_history.append({"role": "assistant", "content": answer})
-                    else:
-                        st.error(f"Agent Error: {response.text}")
+                        # Added HEADERS and explicit timeout handling
+                        response = requests.post(
+                            f"{API_URL}/chat", 
+                            json=payload, 
+                            headers=HEADERS,
+                            timeout=60
+                        )
                         
-                except requests.exceptions.ConnectionError:
-                    st.error("🚨 Connection dropped. Ensure FastAPI is running.")
+                        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 limit tripped on Hugging Face. Give it a brief pause before submitting another question.")
+                        else:
+                            st.error(f"Agent Error: {response.status_code} - {response.text}")
+                            
+                    except requests.exceptions.ConnectionError:
+                        st.error("🚨 Connection dropped. Ensure backend Space is still awake.")
+                    except requests.exceptions.Timeout:
+                        st.error("🚨 The LLM execution timed out waiting for an answer.")