Update app.py

app.py

@@ -1,19 +1,25 @@
 import os
 import gradio as gr
 import pandas as pd
+import matplotlib
+# Use a non-interactive backend for Matplotlib to prevent cloud server crashes
+matplotlib.use('Agg') 
+import matplotlib.pyplot as plt
 from google import genai
 
 # 1. Initialize the Gemini client securely from the environment
 api_key = os.environ.get('GEMINI_API_KEY')
 client = genai.Client(api_key=api_key)
 
-# 2. Define the core data analysis logic
-def analyze_data(uploaded_file, user_question):
-    # Check if a file was actually uploaded
+# 2. Define the core data and plotting logic
+def analyze_and_plot(uploaded_file, user_question):
     if uploaded_file is None:
-        return "⚠️ Please upload a CSV file first before asking a question."
+        return "⚠️ Please upload a CSV file first.", None
         
     try:
+        # Clear any previous figures to avoid overlapping graphs
+        plt.close('all')
+        
         # Read the uploaded CSV file using Pandas
         df = pd.read_csv(uploaded_file.name)
         
@@ -24,7 +30,7 @@ def analyze_data(uploaded_file, user_question):
         data_preview = df.head(3).to_string()
         data_types = df.dtypes.to_string()
         
-        # Build a robust prompt telling Gemini about the data structure
+        # System instructions forcing Gemini to return structured code if a chart is needed
         system_context = (
             f"You are an expert Python Data Analyst. The user has uploaded a dataset.\n"
             f"Dataset Statistics:\n"
@@ -32,43 +38,68 @@ def analyze_data(uploaded_file, user_question):
             f"- Total Columns: {col_count}\n"
             f"- Column Names: {column_names}\n\n"
             f"Data Types:\n{data_types}\n\n"
-            f"Preview of first 3 rows:\n{data_preview}\n\n"
-            f"Please answer the user's question by performing statistical logic or providing "
-            f"Python code snippets if requested. Always base your calculations on this data framework."
+            f"Preview:\n{data_preview}\n\n"
+            f"INSTRUCTIONS:\n"
+            f"1. Answer the user's question clearly.\n"
+            f"2. If the user asks for a chart, graph, or plot, you MUST include a clean Python code block "
+            f"that uses 'matplotlib.pyplot' to build it. Assume the dataframe is already loaded as a variable named 'df'.\n"
+            f"3. Crucial: The LAST line of your chart code block must save the image exactly as: plt.savefig('output_chart.png')\n"
+            f"4. Do not include plt.show(). Only save it to 'output_chart.png'."
         )
         
-        # Send the metadata context and the user query to Gemini
+        # Send the metadata context and user query to Gemini
         response = client.models.generate_content(
             model='gemini-2.5-flash',
             contents=[system_context, f"User Question: {user_question}"]
         )
-        return response.text
+        
+        text_response = response.text
+        chart_path = None
+        
+        # Extract and execute the python code if Gemini generated a graph script
+        if "```python" in text_response:
+            try:
+                # Isolate the code inside the markdown blocks
+                code_block = text_response.split("```python")[1].split("```")[0]
+                
+                # Create a safe execution environment containing the user's data
+                local_vars = {'df': df, 'plt': plt}
+                exec(code_block, {}, local_vars)
+                
+                # Check if the file was successfully written by the executed code
+                if os.path.exists("output_chart.png"):
+                    chart_path = "output_chart.png"
+            except Exception as e:
+                text_response += f"\n\n*(Note: There was a background error rendering the chart: {str(e)})*"
+                
+        return text_response, chart_path
         
     except Exception as e:
-        return f"❌ An error occurred processing the file: {str(e)}"
+        return f"❌ An error occurred processing the file: {str(e)}", None
 
-# 3. Design a beautiful, custom web dashboard using Gradio Blocks
+# 3. Design the custom layout using Gradio Blocks
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📊 Cloud AI Data Analyst")
-    gr.Markdown("Upload a spreadsheet dataset and ask the AI to clean it, summarize it, or uncover hidden insights.")
+    gr.Markdown("# 📊 Visual Cloud AI Data Analyst")
+    gr.Markdown("Upload a CSV file, ask a question, or tell the AI exactly what kind of chart to build.")
     
     with gr.Row():
         with gr.Column(scale=1):
             file_input = gr.File(label="Step 1: Upload CSV Dataset", file_types=[".csv"])
             question_input = gr.Textbox(
-                label="Step 2: What would you like to know about this data?", 
-                placeholder="e.g., Summarize the main trends, find missing values, or calculate averages..."
+                label="Step 2: Ask a question or request a graph", 
+                placeholder="e.g., 'Make a bar chart of sales by region' or 'Plot a histogram of ages'..."
             )
-            submit_btn = gr.Button("Analyze Data", variant="primary")
+            submit_btn = gr.Button("Analyze & Generate Chart", variant="primary")
             
         with gr.Column(scale=1):
-            output_text = gr.Markdown(label="AI Analysis Result")
+            output_text = gr.Markdown(label="AI Written Insights")
+            output_image = gr.Image(label="Generated Data Visualization", type="filepath")
             
     # Connect the inputs to the function when clicking the button
     submit_btn.click(
-        fn=analyze_data, 
+        fn=analyze_and_plot, 
         inputs=[file_input, question_input], 
-        outputs=output_text
+        outputs=[output_text, output_image]
     )
 
 # 4. Launch the public cloud server