Balams/NL2SQL

1import streamlit as st2from phi.agent import Agent3from phi.model.google import Gemini4import duckdb5import kagglehub6import re7import pandas as pd8import os9 10# Load LLM11os.environ["GOOGLE_API_KEY"] = st.secrets["GOOGLE_API_KEY"]12model_name=Gemini(id="gemini-2.0-flash-exp", temperature=0)13 14def load_dataset():15    path=kagglehub.dataset_download("andrexibiza/grocery-sales-dataset")16    categories = pd.read_csv(path + "/categories.csv")17    cities = pd.read_csv(path + "/cities.csv")18    countries = pd.read_csv(path + "/countries.csv")19    customers = pd.read_csv(path + "/customers.csv")20    employees =  pd.read_csv(path + "/employees.csv")21    products = pd.read_csv(path + "/products.csv")22    sales = pd.read_csv(path + "/sales.csv", nrows=50000)23 24    con = duckdb.connect("sales.db")25    # Store DataFrames as tables in DuckDB26    con.execute("CREATE TABLE IF NOT EXISTS categories AS SELECT * FROM categories")27    con.execute("CREATE TABLE IF NOT EXISTS cities AS SELECT * FROM cities")28    con.execute("CREATE TABLE IF NOT EXISTS countries AS SELECT * FROM countries")29    con.execute("CREATE TABLE IF NOT EXISTS customers AS SELECT * FROM customers")30    con.execute("CREATE TABLE IF NOT EXISTS employees AS SELECT * FROM employees")31    con.execute("CREATE TABLE IF NOT EXISTS products AS SELECT * FROM products")32    con.execute("CREATE TABLE IF NOT EXISTS sales AS SELECT * FROM sales")33 34    con.execute("UPDATE sales SET TotalPrice =CAST(FLOOR((RANDOM() % 7500) * 4 + 4) AS INTEGER)")35    con.close()36    print("Data is successfully Loaded..")37 38 39def run_sql_query(sql):40    con = duckdb.connect("sales.db")41    result = con.execute(sql).df()42    con.close()43    return result44 45def get_tablenames():46    con = duckdb.connect("sales.db")47    # Get all table names48    tables = con.execute("SHOW TABLES").df()["name"].tolist()49    tables_dict = {}50    # Retrieve column names for each table51    for table in tables:52        columns = con.execute(f"DESCRIBE {table}").df()["column_name"].tolist()53        tables_dict[table] = columns  # Store in dictionary54    con.close()55    return tables_dict56 57def sql_code_generator(user_query, tables_names):58    agent = Agent(59        name='Sql Agent',60        model=model_name,61        description="You are senior sql developer.",62        instructions=[63            "Convert user query in to SQL questions",64            "Use this table information to generate your code"65            f"Table Details {tables_names}"66        ],67        show_tool_calls=True,68        debug_mode=True69 70    )71 72    response=agent.run(user_query)73    response=re.sub(r"```sql\s*|\s*```", "", response.content.strip())74 75    return response