Spaces:
Running
Running
| import os | |
| import streamlit as st | |
| import pandas as pd | |
| import sqlite3 | |
| import openai | |
| from langchain import OpenAI | |
| from langchain_community.agent_toolkits.sql.base import create_sql_agent | |
| from langchain_community.utilities import SQLDatabase | |
| from langchain_community.document_loaders import CSVLoader | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_community.embeddings import OpenAIEmbeddings | |
| from langchain.chains import RetrievalQA | |
| import sqlparse | |
| import logging | |
| # 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 SQLite database with dynamic table name | |
| conn = sqlite3.connect(':memory:') # Use an in-memory SQLite database | |
| 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) | |
| # Step 3: Use a SQL Agent and setup LangChain's SQL Database connection | |
| db = SQLDatabase.from_uri('sqlite:///:memory:') | |
| db.raw_connection = conn # Use the in-memory connection for LangChain | |
| sql_agent = create_sql_agent(OpenAI(temperature=0), db, verbose=True) | |
| # Step 4: Use FAISS with RAG for context retrieval | |
| embeddings = OpenAIEmbeddings() | |
| loader = CSVLoader(file_path=csv_file.name if csv_file else "default_data.csv") | |
| documents = loader.load() | |
| vector_store = FAISS.from_documents(documents, embeddings) | |
| retriever = vector_store.as_retriever() | |
| rag_chain = RetrievalQA.from_chain_type(llm=OpenAI(temperature=0), retriever=retriever) | |
| # Step 5: Define SQL validation helpers | |
| def validate_sql(query, valid_columns): | |
| """Validates the SQL query by ensuring it references only valid columns.""" | |
| for column in valid_columns: | |
| if column not in query: | |
| 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 6: Generate SQL query based on user input and run it with LangChain SQL Agent | |
| user_prompt = st.text_input("Enter your natural language prompt:") | |
| if user_prompt: | |
| try: | |
| # Step 7: Retrieve context using RAG | |
| context = rag_chain.run(user_prompt) | |
| st.write(f"Retrieved Context: {context}") | |
| # Step 8: Generate SQL query using SQL agent | |
| generated_sql = sql_agent.run(f"{user_prompt} {context}") | |
| st.write(f"Generated SQL Query: {generated_sql}") | |
| # Step 9: 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 10: Execute SQL query | |
| result = pd.read_sql(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}") | |