Update app.py

app.py

@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from dotenv import load_dotenv
 import anthropic
+import ast
 
 # Load environment variables
 load_dotenv()
@@ -118,20 +119,32 @@ if uploaded_file is not None:
             else:
                 parsed_query = ask_gpt(f"Convert this question into a Pandas operation: {query}")
             
-            # Execute the query
-            try:
-                result = eval(parsed_query, {"df": df, "pd": pd})  # Ensuring df is correctly referenced
-                st.write("### Result:")
-                st.write(result if isinstance(result, pd.DataFrame) else str(result))
-                
-                # If numerical data, show a visualization
-                if isinstance(result, pd.Series) and result.dtype in ["int64", "float64"]:
-                    fig, ax = plt.subplots()
-                    result.plot(kind="bar", ax=ax)
-                    st.pyplot(fig)
-                
-            except Exception as e:
-                st.error(f"Error executing query: {str(e)}")
+            # Validate and clean query
+            parsed_query = parsed_query.strip().replace("`", "")
+            st.write(f"Parsed Query: `{parsed_query}`")
+            
+            # Check for column existence if query involves a column
+            for col in column_names:
+                if col in parsed_query and col not in df.columns:
+                    st.error(f"Error: Column '{col}' not found in the uploaded file.")
+                    break
+            else:
+                # Execute the query
+                try:
+                    result = eval(parsed_query, {"df": df, "pd": pd})  # Ensuring df is correctly referenced
+                    st.write("### Result:")
+                    st.write(result if isinstance(result, pd.DataFrame) else str(result))
+                    
+                    # If numerical data, show a visualization
+                    if isinstance(result, pd.Series) and result.dtype in ["int64", "float64"]:
+                        fig, ax = plt.subplots()
+                        result.plot(kind="bar", ax=ax)
+                        st.pyplot(fig)
+                    
+                except SyntaxError as e:
+                    st.error(f"Syntax Error in parsed query: {str(e)}")
+                except Exception as e:
+                    st.error(f"Error executing query: {str(e)}")
     
     # Memory for context retention
     if "query_history" not in st.session_state: