Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,13 +5,13 @@ from langchain.agents import create_sql_agent
 from langchain_groq import ChatGroq
 from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from dotenv import load_dotenv
+import tempfile
 import sqlite3
 import pandas as pd
 import os
 
 load_dotenv()
 
-
 def is_valid_sqlite(file_path):
     try:
         with sqlite3.connect(file_path) as conn:
@@ -20,7 +20,6 @@ def is_valid_sqlite(file_path):
     except sqlite3.DatabaseError:
         return False
 
-
 def text_to_sql(query: str, db_path: str, llm_provider: str, api_key: str, model_name: str):
     try:
         db = SQLDatabase.from_uri(f"sqlite:///{db_path}")
@@ -41,7 +40,6 @@ def text_to_sql(query: str, db_path: str, llm_provider: str, api_key: str, model
     except Exception as e:
         return f"Error: {str(e)}"
 
-
 def show_tables_as_df(db_path):
     conn = sqlite3.connect(db_path)
     cursor = conn.cursor()
@@ -61,46 +59,40 @@ def show_tables_as_df(db_path):
 
 
 # Streamlit UI
-st.title('🗃️ Chat with SQLite Database')
+st.title('🗃️ Chat with YOUR SQLite Database')
 
-st.write("Interact with your SQLite databases stored on server using LLM-powered queries.")
+st.write("Upload your `.db` file and interact using natural language queries powered by LLMs.")
 
-# List of .db files in "data" folder
-db_folder = "data"
-os.makedirs(db_folder, exist_ok=True)
-db_files = [f for f in os.listdir(db_folder) if f.endswith('.db')]
-
-selected_db = st.selectbox("Select Database", options=db_files)
+uploaded_file = st.file_uploader("Upload SQLite Database (.db file)", type=["db"])
 
 llm_provider = st.radio("Choose LLM Provider", options=['OPEN_ROUTER', 'GROQ', 'OPENAI'])
 model_name = st.text_input("Enter the Model Name", value='nousresearch/deephermes-3-mistral-24b-preview:free')
 api_key = st.text_input("Enter Your API Key", type="password")
 query = st.text_area("Enter Your Query")
 
-try:
-    if selected_db:
-        tmp_db_path = os.path.join(db_folder, selected_db)
+if uploaded_file:
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".db", dir="/tmp") as tmpfile:
+        tmpfile.write(uploaded_file.read())
+        tmp_db_path = tmpfile.name
 
-        if not is_valid_sqlite(tmp_db_path):
-            st.error("The selected file is not a valid SQLite database.")
-        else:
-            st.success("Valid SQLite database selected!")
-
-            # Show tables
-            st.info("Displaying first 10 rows from each table:")
-            show_tables_as_df(tmp_db_path)
-
-            if st.button("RUN Query"):
-                if not api_key or not model_name:
-                    st.error("Please provide API key and model name.")
-                elif not query.strip():
-                    st.error("Please enter a query.")
-                else:
-                    st.info(f"Running query on `{selected_db}`...")
-                    result = text_to_sql(query, tmp_db_path, llm_provider, api_key, model_name)
-                    st.success("Query Result:")
-                    st.write(result)
-
-except Exception as e:
-    st.error(f"Error: {str(e)}")
+    if not is_valid_sqlite(tmp_db_path):
+        st.error("Uploaded file is not a valid SQLite database.")
+    else:
+        st.success(f"Valid database `{uploaded_file.name}` uploaded!")
+
+        st.info("Displaying first 10 rows from each table:")
+        show_tables_as_df(tmp_db_path)
+
+        if st.button("RUN Query"):
+            if not api_key or not model_name:
+                st.error("Please provide API key and model name.")
+            elif not query.strip():
+                st.error("Please enter a query.")
+            else:
+                st.info(f"Running query on `{uploaded_file.name}`...")
+                result = text_to_sql(query, tmp_db_path, llm_provider, api_key, model_name)
+                st.success("Query Result:")
+                st.write(result)
+else:
+    st.info("Please upload a SQLite `.db` file to begin.")