Create app.py

app.py

@@ -0,0 +1,146 @@
+import streamlit as st
+from huggingface_hub import InferenceClient
+from langchain_core.output_parsers import StrOutputParser
+import os
+from dotenv import load_dotenv
+import pandas as pd
+import sqlite3
+import re
+
+load_dotenv()
+token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+api = InferenceClient(token=token)
+parser = StrOutputParser()
+
+# Streamlit app
+st.title("AiSQL: AI-Powered SQL Query Generator")
+
+# File uploader for CSV
+uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
+
+if uploaded_file:
+    # Read CSV into DataFrame
+    df = pd.read_csv(uploaded_file)
+    st.write("Uploaded Data:")
+    st.dataframe(df)
+    
+    # Normalize column names: replace spaces and special characters with underscores
+    df.columns = [re.sub(r'\W+', '_', col.strip()) for col in df.columns]
+    
+    st.write("Normalized Columns in the CSV:")
+    st.write(df.columns.tolist())
+    
+    # Create SQLite in-memory database
+    conn = sqlite3.connect(':memory:')
+    df.to_sql('data', conn, index=False, if_exists='replace')
+    
+    # Natural language query input
+    nl_query = st.text_area("Enter your query in natural language or in code:")
+    
+    if st.button("Run Query/Code"):
+        try:
+            # Generate SQL query using LLM
+            system_message = (
+                "You are an AI assistant that converts natural language queries into SQL queries based on the following table schema.\n"
+                f"Table name: data\n"
+                f"Columns: {', '.join(df.columns.tolist())}\n"
+                "Provide only the SQL query suggestion in code blocks without any explanations, comments, or other text."
+            )
+            messages = [
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": nl_query}
+            ]
+            llm = api.chat.completions.create(
+                model="Qwen/Qwen2.5-Coder-32B-Instruct",
+                max_tokens=150,
+                messages=messages
+            )
+            raw_response = llm.choices[0].message['content'].strip()
+            
+            # Remove code blocks if present
+            sql_query = re.sub(r'```sql\n?|\n?```', '', raw_response).strip()
+            
+            # Additional cleaning: Extract the first SQL statement
+            match = re.search(r'\b(SELECT|INSERT|UPDATE|DELETE|CREATE|DROP|ALTER)\b[\s\S]*?;', sql_query, re.IGNORECASE)
+            if match:
+                sql_query = match.group(0)
+            else:
+                st.error("Failed to extract a valid SQL query from the response.")
+                st.write("**Raw LLM Response:**")
+                st.write(raw_response)
+                st.stop()
+            
+            # Validate that the SQL query starts with a valid keyword
+            valid_sql_keywords = ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'DROP', 'ALTER']
+            if not any(sql_query.upper().startswith(keyword) for keyword in valid_sql_keywords):
+                st.error("The generated SQL query does not start with a valid SQL command.")
+                st.write("**Extracted SQL Query:**")
+                st.write(sql_query)
+                st.stop()
+            
+            st.markdown(f"**Generated SQL Query:** `{sql_query}`")
+            
+            # Execute SQL query
+            result = pd.read_sql_query(sql_query, conn)
+            st.write("Query Results:")
+            st.dataframe(result)
+        except Exception as e:
+            st.error(f"Error: {e}")
+    
+    # Generate query suggestions using LLM
+    if st.button("Show Query Suggestions"):
+        try:
+            system_message = (
+                "You are an AI assistant that provides SQL query suggestions based on the following table schema.\n"
+                f"Table name: data\n"
+                f"Columns: {', '.join(df.columns.tolist())}\n"
+                "Provide exactly 5 example SQL queries separated by semicolons without any explanations, comments, or code blocks."
+            )
+            suggestion_messages = [
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": "Provide SQL query suggestions."}
+            ]
+            suggestions_llm = api.chat.completions.create(
+                model="Qwen/Qwen2.5-Coder-32B-Instruct",
+                max_tokens=300,
+                messages=suggestion_messages
+            )
+            raw_suggestions = suggestions_llm.choices[0].message['content']
+            # Remove code blocks if present
+            suggestions = re.sub(r'```sql\n?|\n?```', '', raw_suggestions).strip()
+            
+            # Split multiple queries separated by semicolons
+            suggestions_list = [query.strip() for query in suggestions.split(';') if query.strip()]
+            
+            # Validate each suggestion starts with a valid SQL keyword
+            valid_sql_keywords = ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'DROP', 'ALTER']
+            valid_suggestions = []
+            for query in suggestions_list:
+                if any(query.upper().startswith(keyword) for keyword in valid_sql_keywords):
+                    valid_suggestions.append(query + ';')
+            
+            st.session_state['valid_suggestions'] = valid_suggestions
+            
+            if valid_suggestions:
+                formatted_suggestions = ';\n'.join(valid_suggestions)
+                st.write("**Query Suggestions:**")
+                st.code(formatted_suggestions, language='sql')
+                
+                # Optionally, allow users to select a suggestion to execute
+                if 'valid_suggestions' in st.session_state:
+                    selected_query = st.selectbox("Select a query to execute:", st.session_state['valid_suggestions'])
+                    if st.button("Execute Selected Query"):
+                        # Execute the selected query
+                        try:
+                            st.write(f"**Executing SQL Query:** `{selected_query}`")
+                            result = pd.read_sql_query(selected_query, conn)
+                            st.write("Query Results:")
+                            st.dataframe(result)
+                        except Exception as e:
+                            st.error(f"Error executing selected query: {e}")
+            else:
+                st.error("No valid SQL query suggestions were generated.")
+                st.write("**Raw Suggestions Response:**")
+                st.write(suggestions)
+        except Exception as e:
+            st.error(f"Error generating suggestions: {e}")