Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image, ImageDraw
 import google.generativeai as genai
 import traceback
 
-def process_file(file, instructions, api_key):
+def process_file(file, instructions, api_key, graph_type):
     try:
         # Initialize Gemini
         genai.configure(api_key=api_key)
@@ -24,9 +24,10 @@ def process_file(file, instructions, api_key):
             Data columns: {list(df.columns)}
             Data shape: {df.shape}
             Instructions: {instructions}
+            Graph Type: {graph_type}
             
-            Based on this, create 3 appropriate visualizations that provide meaningful insights. For each visualization:
-            1. Choose the most suitable plot type (bar, line, scatter, hist, pie, heatmap)
+            Based on this, create 3 appropriate visualizations of the specified type that provide meaningful insights. For each visualization:
+            1. Use the specified plot type: {graph_type}
             2. Determine appropriate data aggregation (e.g., top 5 categories, monthly averages)
             3. Select relevant columns for x-axis, y-axis, and any additional dimensions (color, size)
             4. Provide a clear, concise title that explains the insight
@@ -36,9 +37,9 @@ def process_file(file, instructions, api_key):
             
             Return your response as a Python list of dictionaries:
             [
-                {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}},
-                {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}},
-                {{"title": "...", "plot_type": "...", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}}
+                {{"title": "...", "plot_type": "{graph_type}", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}},
+                {{"title": "...", "plot_type": "{graph_type}", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}},
+                {{"title": "...", "plot_type": "{graph_type}", "x": "...", "y": "...", "agg_func": "...", "top_n": ..., "additional": {{"color": "...", "size": "..."}}}}
             ]
         """)
 
@@ -71,18 +72,18 @@ def process_file(file, instructions, api_key):
             if 'top_n' in plot and plot['top_n']:
                 plot_df = plot_df.nlargest(plot['top_n'], plot['y'])
             
-            if plot['plot_type'] == 'bar':
+            if graph_type == 'bar':
                 plot_df.plot(kind='bar', x=plot['x'], y=plot['y'], ax=ax)
-            elif plot['plot_type'] == 'line':
+            elif graph_type == 'line':
                 plot_df.plot(kind='line', x=plot['x'], y=plot['y'], ax=ax)
-            elif plot['plot_type'] == 'scatter':
+            elif graph_type == 'scatter':
                 plot_df.plot(kind='scatter', x=plot['x'], y=plot['y'], ax=ax, 
                              c=plot['additional'].get('color'), s=plot_df[plot['additional'].get('size', 'y')])
-            elif plot['plot_type'] == 'hist':
+            elif graph_type == 'hist':
                 plot_df[plot['x']].hist(ax=ax, bins=20)
-            elif plot['plot_type'] == 'pie':
+            elif graph_type == 'pie':
                 plot_df.plot(kind='pie', y=plot['y'], labels=plot_df[plot['x']], ax=ax, autopct='%1.1f%%')
-            elif plot['plot_type'] == 'heatmap':
+            elif graph_type == 'heatmap':
                 pivot_df = plot_df.pivot(index=plot['x'], columns=plot['additional']['color'], values=plot['y'])
                 ax.imshow(pivot_df, cmap='YlOrRd')
                 ax.set_xticks(range(len(pivot_df.columns)))
@@ -91,7 +92,7 @@ def process_file(file, instructions, api_key):
                 ax.set_yticklabels(pivot_df.index)
             
             ax.set_title(plot['title'])
-            if plot['plot_type'] != 'pie':
+            if graph_type != 'pie':
                 ax.set_xlabel(plot['x'])
                 ax.set_ylabel(plot['y'])
             plt.tight_layout()
@@ -121,13 +122,14 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
         instructions = gr.Textbox(label="Analysis Instructions", placeholder="Describe the analysis you want...")
     
     api_key = gr.Textbox(label="Gemini API Key", type="password")
+    graph_type = gr.Dropdown(choices=["bar", "line", "scatter", "hist", "pie", "heatmap"], label="Graph Type")
     submit = gr.Button("Generate Insights", variant="primary")
     
     output_images = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
 
     submit.click(
         process_file,
-        inputs=[file, instructions, api_key],
+        inputs=[file, instructions, api_key, graph_type],
         outputs=output_images
     )