Update app.py

app.py

@@ -4,7 +4,7 @@ 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.agent_toolkits.sql.base import SQLDatabaseToolkit, create_sql_agent  # SQL toolkit import
 from langchain_community.utilities import SQLDatabase
 from langchain_community.document_loaders import CSVLoader
 from langchain_community.vectorstores import FAISS
@@ -34,12 +34,17 @@ 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
+# 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
-sql_agent = create_sql_agent(OpenAI(temperature=0), db, verbose=True)
 
-# Step 4: Use FAISS with RAG for context retrieval
+# Step 4: Create the SQL toolkit
+sql_toolkit = SQLDatabaseToolkit(db)
+
+# Step 5: Create the SQL agent using the toolkit
+sql_agent = create_sql_agent(OpenAI(temperature=0), toolkit=sql_toolkit, verbose=True)
+
+# Step 6: 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()
@@ -48,7 +53,7 @@ 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
+# Step 7: 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,25 +66,25 @@ def validate_sql_with_sqlparse(query):
     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
+# Step 8: 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
+        # Step 9: Retrieve context using RAG
         context = rag_chain.run(user_prompt)
         st.write(f"Retrieved Context: {context}")
 
-        # Step 8: Generate SQL query using SQL agent
+        # Step 10: 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
+        # Step 11: 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
+            # Step 12: Execute SQL query
             result = pd.read_sql(generated_sql, conn)
             st.write("Query Results:")
             st.dataframe(result)