Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,100 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from langchain_community.utilities import SQLDatabase
+from langchain.chat_models import ChatOpenAI
+from langchain.agents import create_sql_agent
+from langchain_groq import ChatGroq
+from langchain_community.agent_toolkits import SQLDatabaseToolkit
+
+import tempfile
+import sqlite3
+import pandas as pd
+
+
+def is_valid_sqlite(file_path):
+    try:
+        with sqlite3.connect(file_path) as conn:
+            conn.execute("SELECT name FROM sqlite_master LIMIT 1;")
+        return True
+    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}")
+
+        if llm_provider == 'OPENAI':
+            llm = ChatOpenAI(api_key=api_key, model=model_name)
+        elif llm_provider == 'OPEN_ROUTER':
+            llm = ChatOpenAI(api_key=api_key, base_url='https://openrouter.ai/api/v1', model=model_name)
+        elif llm_provider == 'GROQ':
+            llm = ChatGroq(api_key=api_key, model=model_name)
+        else:
+            return "Unsupported LLM provider selected."
+
+        toolkit = SQLDatabaseToolkit(llm=llm, db=db)
+        db_chain = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
+        return db_chain.run(query)
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+
+def show_tables_as_df(db_path):
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+    tables = cursor.fetchall()
+
+    if tables:
+        for table_name in tables:
+            table = table_name[0]
+            st.subheader(f"Table: {table}")
+            df = pd.read_sql_query(f"SELECT * FROM {table} LIMIT 10", conn)
+            st.dataframe(df)
+    else:
+        st.write("No tables found in database.")
+
+    conn.close()
+
+
+# Streamlit UI
+st.title('🗃️ Chat with SQLite Database')
+
+st.write("This app lets you interact with a SQLite database using natural language queries powered by LLMs.")
+
+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 uploaded_file is not None:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as tmpfile:
+            tmpfile.write(uploaded_file.read())
+            tmp_db_path = tmpfile.name
+
+        if not is_valid_sqlite(tmp_db_path):
+            st.error("The uploaded file is not a valid SQLite database.")
+        else:
+            st.success("Valid SQLite database uploaded!")
+
+            # Show tables as pandas DataFrames
+            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)
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+except Exception as e:
+    st.error(f"Error: {str(e)}")