Update app.py

app.py

@@ -79,23 +79,27 @@ def validate_sql_with_sqlparse(query):
 user_prompt = st.text_input("Enter your natural language prompt:")
 if user_prompt:
     try:
-        # Step 8: Retrieve context using RAG
-        context = rag_chain.run(user_prompt)
+        # Step 8: Add valid column names to the prompt
+        column_hints = f" Use only these columns: {', '.join(valid_columns)}"
+        prompt_with_columns = user_prompt + column_hints
+
+        # Step 9: Retrieve context using RAG
+        context = rag_chain.run(prompt_with_columns)
         st.write(f"Retrieved Context: {context}")
 
-        # Step 9: Generate SQL query using SQL agent
-        generated_sql = sql_agent.run(f"{user_prompt} {context}")
+        # Step 10: Generate SQL query using SQL agent
+        generated_sql = sql_agent.run(f"{prompt_with_columns} {context}")
         
         # Debug: Display generated SQL query for inspection
         st.write(f"Generated SQL Query: {generated_sql}")
 
-        # Step 10: 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 11: Execute SQL query
+            # Step 12: Execute SQL query
             result = pd.read_sql(generated_sql, conn)
             st.write("Query Results:")
             st.dataframe(result)