Update app.py

app.py

@@ -34,12 +34,21 @@ data.to_sql(table_name, conn, index=False, if_exists='replace')
 # SQL table metadata (for validation and schema)
 valid_columns = list(data.columns)
 
+# Debug: Display valid columns for user to verify
+st.write(f"Valid columns: {valid_columns}")
+
 # 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)
+sql_agent = create_sql_agent(
+    OpenAI(temperature=0),
+    db=db,
+    verbose=True,
+    max_iterations=15,  # Increased iteration limit
+    max_execution_time=60  # Set timeout limit to 60 seconds
+)
 
 # Step 5: Use FAISS with RAG for context retrieval
 embeddings = OpenAIEmbeddings()
@@ -53,9 +62,12 @@ rag_chain = RetrievalQA.from_chain_type(llm=OpenAI(temperature=0), retriever=ret
 # Step 6: 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
+    parsed = sqlparse.parse(query)
+    for token in parsed[0].tokens:
+        if token.ttype is None:  # If it's a column name
+            column_name = str(token).strip()
+            if column_name not in valid_columns:
+                return False
     return True
 
 def validate_sql_with_sqlparse(query):
@@ -73,6 +85,8 @@ if user_prompt:
 
         # Step 9: Generate SQL query using SQL agent
         generated_sql = sql_agent.run(f"{user_prompt} {context}")
+        
+        # Debug: Display generated SQL query for inspection
         st.write(f"Generated SQL Query: {generated_sql}")
 
         # Step 10: Validate SQL query