Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import pandas as pd | |
| import sqlite3 | |
| from langchain import OpenAI, LLMChain, PromptTemplate | |
| from langchain_community.utilities import SQLDatabase | |
| import sqlparse | |
| import logging | |
| from sql_metadata import Parser # Added import | |
| # OpenAI API key (ensure it is securely stored) | |
| openai_api_key = os.getenv("OPENAI_API_KEY") | |
| # Step 1: Upload CSV data file (or use default) | |
| csv_file = st.file_uploader("Upload your CSV file", type=["csv"]) | |
| if csv_file is None: | |
| data = pd.read_csv("default_data.csv") # Use default CSV if no file is uploaded | |
| st.write("Using default data.csv file.") | |
| else: | |
| data = pd.read_csv(csv_file) | |
| st.write(f"Data Preview ({csv_file.name}):") | |
| st.dataframe(data.head()) | |
| # Step 2: Load CSV data into a persistent SQLite database | |
| db_file = 'my_database.db' | |
| conn = sqlite3.connect(db_file) | |
| table_name = csv_file.name.split('.')[0] if csv_file else "default_table" | |
| data.to_sql(table_name, conn, index=False, if_exists='replace') | |
| # SQL table metadata (for validation and schema) | |
| valid_columns = list(data.columns) | |
| st.write(f"Valid columns: {valid_columns}") | |
| # Step 3: Define SQL validation helpers | |
| def validate_sql(query, valid_columns): | |
| """Validates the SQL query by ensuring it references only valid columns.""" | |
| parser = Parser(query) | |
| columns_in_query = parser.columns | |
| for column in columns_in_query: | |
| if column not in valid_columns: | |
| st.write(f"Invalid column detected: {column}") | |
| return False | |
| return True | |
| def validate_sql_with_sqlparse(query): | |
| """Validates SQL syntax using sqlparse.""" | |
| parsed_query = sqlparse.parse(query) | |
| return len(parsed_query) > 0 | |
| # Step 4: Set up the LLM Chain to generate SQL queries | |
| template = """ | |
| You are an expert data scientist. Given a natural language question, the name of the table, and a list of valid columns, generate a valid SQL query that answers the question. | |
| Question: {question} | |
| Table name: {table_name} | |
| Valid columns: {columns} | |
| SQL Query: | |
| """ | |
| prompt = PromptTemplate(template=template, input_variables=['question', 'table_name', 'columns']) | |
| sql_generation_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt) | |
| # Step 5: Generate SQL query based on user input | |
| user_prompt = st.text_input("Enter your natural language prompt:") | |
| if user_prompt: | |
| try: | |
| # Step 6: Adjust the logic to handle "what are the columns" query | |
| if "columns" in user_prompt.lower(): | |
| # Custom logic to return columns | |
| st.write(f"The columns are: {', '.join(valid_columns)}") | |
| else: | |
| columns = ', '.join(valid_columns) | |
| generated_sql = sql_generation_chain.run({'question': user_prompt, 'table_name': table_name, 'columns': columns}) | |
| # Debug: Display generated SQL query for inspection | |
| st.write(f"Generated SQL Query:\n{generated_sql}") | |
| # Step 7: Validate SQL query | |
| if not validate_sql_with_sqlparse(generated_sql): | |
| st.write("Generated SQL is not valid.") | |
| elif not validate_sql(generated_sql, valid_columns): | |
| st.write("Generated SQL references invalid columns.") | |
| else: | |
| # Step 8: Execute SQL query | |
| result = pd.read_sql_query(generated_sql, conn) | |
| st.write("Query Results:") | |
| st.dataframe(result) | |
| except Exception as e: | |
| logging.error(f"An error occurred: {e}") | |
| st.write(f"Error: {e}") | |