Update app.py

app.py

@@ -79,7 +79,7 @@ class DataAnalysisAgent(CodeAgent):
 
         Use the provided tools to analyze this specific dataset and return detailed results.
         """
-        return super().run(enhanced_prompt, data=self.dataset, **kwargs) # Pass data as argument
+        return super().run(enhanced_prompt, data=self.dataset, **kwargs)  # Pass data as argument
 
 
 @tool
@@ -204,16 +204,15 @@ def suggest_features(data: pd.DataFrame) -> str:
 @tool
 def describe_data(data: pd.DataFrame) -> str:
     """Generates a comprehensive descriptive statistics report for the entire DataFrame.
-    
+
     Args:
-       data: A pandas DataFrame containing the dataset to analyze.
-    
+        data: A pandas DataFrame containing the dataset to analyze.
+
     Returns:
-      str: String representation of the descriptive statistics
+        str: String representation of the descriptive statistics
     """
-    
-    return data.describe(include='all').to_string()
 
+    return data.describe(include="all").to_string()
 
 
 @tool
@@ -232,19 +231,19 @@ def execute_code(code_string: str, data: pd.DataFrame) -> str:
 
         # Execute the code with the passed variables
         exec(code_string, local_vars)
-        
+
         if "result" in local_vars:
-           if isinstance(local_vars["result"], (pd.DataFrame, pd.Series)):
-              return local_vars["result"].to_string()
-           elif isinstance(local_vars["result"], plt.Figure):
-             buf = io.BytesIO()
-             local_vars["result"].savefig(buf, format='png')
-             plt.close(local_vars["result"])
-             return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
-           else:
-            return str(local_vars["result"])
+            if isinstance(local_vars["result"], (pd.DataFrame, pd.Series)):
+                return local_vars["result"].to_string()
+            elif isinstance(local_vars["result"], plt.Figure):
+                buf = io.BytesIO()
+                local_vars["result"].savefig(buf, format="png")
+                plt.close(local_vars["result"])
+                return f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"
+            else:
+                return str(local_vars["result"])
         else:
-             return "Code executed successfully, but no variable called 'result' was assigned."
+            return "Code executed successfully, but no variable called 'result' was assigned."
 
     except Exception as e:
         return f"Error executing code: {str(e)}"
@@ -279,10 +278,10 @@ def main():
     if "agent" not in st.session_state:
         st.session_state["agent"] = None
     if "custom_code" not in st.session_state:
-        st.session_state['custom_code'] = ""
+        st.session_state["custom_code"] = ""
 
     uploaded_file = st.file_uploader("Choose a CSV, Excel, or JSON file", type=["csv", "xlsx", "xls", "json"])
-    
+
     if uploaded_file:
         with st.spinner("Loading and processing your data..."):
             data = load_data(uploaded_file)
@@ -297,7 +296,7 @@ def main():
                         analyze_categorical_columns,
                         suggest_features,
                         describe_data,
-                        execute_code
+                        execute_code,
                     ],
                     model=GroqLLM(),
                     additional_authorized_imports=["pandas", "numpy", "matplotlib", "seaborn"],
@@ -316,7 +315,7 @@ def main():
                 "Correlation Analysis",
                 "Categorical Analysis",
                 "Feature Engineering",
-                 "Data Description",
+                "Data Description",
                 "Custom Code",
                 "Custom Question",
             ],
@@ -364,29 +363,29 @@ def main():
                     "of the data."
                 )
                 st.write(result)
-                
+
         elif analysis_type == "Custom Code":
-            st.session_state['custom_code'] = st_ace(
+            st.session_state["custom_code"] = st_ace(
                 placeholder="Enter your Python code here...",
                 language="python",
                 theme="github",
                 key="code_editor",
-                value=st.session_state['custom_code'],
+                value=st.session_state["custom_code"],
             )
             if st.button("Run Code"):
-                 with st.spinner("Executing custom code..."):
+                with st.spinner("Executing custom code..."):
                     result = st.session_state["agent"].run(
-                       f"Execute the following code and return any 'result' variable"
-                       f"```python\n{st.session_state['custom_code']}\n```"
-                   )
+                        f"Execute the following code and return any 'result' variable"
+                        f"```python\n{st.session_state['custom_code']}\n```"
+                    )
                     if isinstance(result, str) and result.startswith("data:image"):
                         st.image(f"{result}")
                     else:
-                       st.write(result)
-                  
+                        st.write(result)
+
         elif analysis_type == "Custom Question":
-             question = st.text_input("What would you like to know about your data?")
-             if question:
+            question = st.text_input("What would you like to know about your data?")
+            if question:
                 with st.spinner("Analyzing..."):
                     result = st.session_state["agent"].run(question, stream=True)  # Pass stream argument here
                     st.write(result)