Update app.py

app.py

@@ -2,18 +2,19 @@ import os
 import streamlit as st
 import pandas as pd
 import sqlite3
-from langchain import OpenAI, LLMChain, PromptTemplate
+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
-from sql_metadata import Parser
 
 # OpenAI API key (ensure it is securely stored)
-openai_api_key = os.getenv("OPENAI_API_KEY")
-
-# Initialize conversation history
-if 'conversation' not in st.session_state:
-    st.session_state.conversation = []  # Store previous conversation messages
+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"])
@@ -25,94 +26,66 @@ else:
     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)
+# 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)
 
-# Display the conversation thread
-st.markdown("### Conversation Thread:")
-for message in st.session_state.conversation:
-    if message.startswith("User:"):
-        st.markdown(f"<p style='color:blue'><strong>{message}</strong></p>", unsafe_allow_html=True)
-    else:
-        st.markdown(f"<p style='color:green'><strong>{message}</strong></p>", unsafe_allow_html=True)
+# Step 3: Set up the SQL Database for LangChain
+db = SQLDatabase.from_uri('sqlite:///:memory:')
+db.raw_connection = conn  # Use the in-memory connection for LangChain
+
+# Step 4: Create the SQL agent with the correct parameter name
+sql_agent = create_sql_agent(OpenAI(temperature=0), db=db, verbose=True)
+
+# Step 5: 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()
 
-# Step 3: Define SQL validation helpers
+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 6: 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:
-            return False, f"Invalid column detected: {column}"
-    return True, None
+    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 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 message:")
-
+# Step 7: 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:
-    # Add user prompt to conversation history
-    st.session_state.conversation.append(f"User: {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
-            columns_response = f"The columns are: {', '.join(valid_columns)}"
-            st.session_state.conversation.append(f"Bot: {columns_response}")
+        # Step 8: Retrieve context using RAG
+        context = rag_chain.run(user_prompt)
+        st.write(f"Retrieved Context: {context}")
+
+        # Step 9: Generate SQL query using SQL agent
+        generated_sql = sql_agent.run(f"{user_prompt} {context}")
+        st.write(f"Generated SQL Query: {generated_sql}")
+
+        # Step 10: 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:
-            # Generate SQL query based on user input
-            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.session_state.conversation.append(f"Bot: Generated SQL Query:\n{generated_sql}")
+            # Step 11: Execute SQL query
+            result = pd.read_sql(generated_sql, conn)
+            st.write("Query Results:")
+            st.dataframe(result)
 
-            # Step 7: Validate SQL query
-            if not validate_sql_with_sqlparse(generated_sql):
-                error_message = "Generated SQL is not valid."
-                st.session_state.conversation.append(f"Bot: {error_message}")
-            elif not validate_sql(generated_sql, valid_columns)[0]:
-                invalid_column_message = "Generated SQL references invalid columns."
-                st.session_state.conversation.append(f"Bot: {invalid_column_message}")
-            else:
-                # Step 8: Execute SQL query
-                result = pd.read_sql_query(generated_sql, conn)
-                st.session_state.conversation.append("Bot: Here are the results of your query:")
-                st.session_state.conversation.append(result.to_string(index=False))  # Add query result as string
-                
     except Exception as e:
         logging.error(f"An error occurred: {e}")
-        error_message = f"Error: {e}"
-        st.session_state.conversation.append(f"Bot: {error_message}")
-
-# Persist the conversation after each message
+        st.write(f"Error: {e}")