Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,133 +1,115 @@
 import streamlit as st
 import requests
 import uuid
-import time
-import os
-
-# --- CONFIGURATION ---
-# Replace this with your exact direct backend Space URL (no trailing slash)
-API_URL = "https://lightrt-text2sql-backend.hf.space" 
 
+# 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.")
 
-# --- 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 ---
-if "thread_id" not in st.session_state:
-    st.session_state.thread_id = str(uuid.uuid4())
+# 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]
+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 "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 = []
 
-# --- SIDEBAR: DATABASE CONNECTION ---
+# 3. SIDEBAR: DATABASE CONNECTION (The Handoff to FastAPI)
 with st.sidebar:
     st.header("⚙️ Database Setup")
     
-    db_input = st.text_input(
-        "Enter Database URL:",
-        disabled=st.session_state.is_db_connected
-    )
+    # Text input for the database URL
+    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", 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..."):
+        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
+                    }
+                    
                     try:
-                        payload = {"connection_url": db_input, "user_id": st.session_state.user_id}
-                        
-                        # Added HEADERS here
+                        # Make sure your backend endpoint is actually /upload_url 
+                        # (Check your FastAPI main.py to confirm this endpoint name!)
                         response = requests.post(
-                            f"{API_URL}/upload_url", 
-                            json=payload, 
-                            headers=HEADERS,
-                            timeout=30
+                            "https://lightrt-text2sql-backend.hf.space/upload_url", 
+                            json=payload_data
                         )
                         
                         if response.status_code == 200:
                             st.session_state.is_db_connected = True
                             st.session_state.connection_url = db_input
-                            
-                            # Keeping your buffer for Qdrant storage compilation
-                            time.sleep(15) 
-                            
                             st.success("Database connected securely!")
-                            st.rerun() 
-                        elif response.status_code == 429:
-                            st.error("🚨 Hugging Face is rate-limiting the connection. Please wait 2-3 minutes before clicking again.")
+                            st.rerun() # Refresh to update the sidebar UI
                         else:
-                            st.error(f"Failed to connect: {response.status_code} - {response.text}")
+                            st.error(f"Connection failed: {response.text}")
+                            
                     except requests.exceptions.ConnectionError:
-                        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.")
+                        st.error("Cannot connect to backend. Is FastAPI running?")
+            else:
+                st.warning("Please enter a valid URL first.")
     else:
         st.success("✅ Connected to Database")
         st.caption(f"URL: {st.session_state.connection_url}")
-        
-        if st.button("Disconnect & Reset", use_container_width=True):
+        if st.button("Disconnect & Reset"):
             st.session_state.clear()
             st.rerun()
 
-# --- MAIN CHAT INTERFACE ---
-st.title("🗣️ Text2SQL Agent")
-
-if not st.session_state.is_db_connected:
-    st.info("👈 Please connect your database in the sidebar to begin analyzing data.")
-else:
-    # Display message history safely
-    for msg in st.session_state.chat_history:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"])
+# 4. CHAT HISTORY RENDERING
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
 
-    # Chat Input Box
-    if user_query := st.chat_input("Ask a question about your data..."):
+# 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..."):
         
-        # 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)
-
-            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
-                        }
-                        
-                        # Added HEADERS and explicit timeout handling
-                        response = requests.post(
-                            f"{API_URL}/chat", 
-                            json=payload, 
-                            headers=HEADERS,
-                            timeout=60
-                        )
-                        
-                        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.")
+        # Immediately render the user's new message to the UI
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+            
+        # Show a loading indicator while we wait for FastAPI
+        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
+                )
+                
+                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)
+                    
+                    # 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}")
+                    
+            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.")