Update app.py

app.py

@@ -1,16 +1,17 @@
-import os  # Add this line at the top of your file
+import os
 import streamlit as st
 import pandas as pd
 import sqlite3
 import openai
-from transformers import pipeline  # Using Hugging Face pipeline for memory-efficient loading
 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")
@@ -27,33 +28,27 @@ else:
 
 # Step 2: Load CSV data into SQLite database with dynamic table name
 conn = sqlite3.connect(':memory:')  # Use an in-memory SQLite database
-
-# Dynamically name the table based on the uploaded file name or fallback to a default name
 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 smaller LLaMA for context retrieval (RAG)
-llama_pipeline = pipeline("text-generation", model="huggyllama/llama-2-3b-hf", device=0)  # Use smaller model
+# 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: Implement RAG with FAISS for vectorized document retrieval
-embeddings = OpenAIEmbeddings()  # You can use other embeddings if preferred
+# 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()
 
-# Use FAISS for retrieval and document search
 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: OpenAI for SQL query generation based on user prompt and context
-openai_llm = OpenAI(temperature=0)
-db = SQLDatabase.from_uri('sqlite:///:memory:')  # Create an SQLite database for LangChain
-db.raw_connection = conn  # Use the in-memory connection for LangChain
-sql_agent = create_sql_agent(openai_llm, db, verbose=True)
-
-# Step 6: Validate and Execute the SQL Query
+# 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:
@@ -61,36 +56,34 @@ def validate_sql(query, valid_columns):
             return False
     return True
 
-# Step 7: SQL Validation with `sqlparse`
 def validate_sql_with_sqlparse(query):
     """Validates SQL syntax using sqlparse."""
     parsed_query = sqlparse.parse(query)
     return len(parsed_query) > 0
 
-# Step 8: Get user prompt, retrieve context, and generate SQL query
+# 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 9: Retrieve relevant context using LLaMA RAG
-        rag_result = llama_pipeline(user_prompt, max_length=200)
-        st.write(f"Retrieved Context from LLaMA RAG: {rag_result}")
-
-        # Step 10: Generate SQL query with OpenAI based on user prompt and retrieved context
-        query_input = f"{user_prompt} {rag_result}"
-        generated_sql = sql_agent.run(query_input)
+        # 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 11: Validate the SQL query before execution
+        # 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 12: Execute the SQL query
+            # 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}")