Update graph_tool.py

graph_tool.py

@@ -4,12 +4,13 @@ import matplotlib.pyplot as plt
 import numpy as np
 import json
 
-def generate_plot(plot_config_json):
+def generate_plot(input: str, return_html: bool = True):
     """
-    Generates a plot (bar, line, or pie) and returns it as an HTML-formatted Base64-encoded image string.
+    Generates a plot (bar, line, or pie) and returns it as an HTML-formatted Base64-encoded image string,
+    or just the Base64 string if return_html=False.
     
     Args:
-        plot_config_json (str): JSON string containing all plot configuration:
+        input (str): JSON string containing all plot configuration:
             {
                 "data": {"key1": value1, "key2": value2, ...},
                 "labels": ["label1", "label2", ...],
@@ -18,13 +19,14 @@ def generate_plot(plot_config_json):
                 "x_label": "X Axis Label" (optional),
                 "y_label": "Y Axis Label" (optional)
             }
+        return_html (bool): Whether to return full HTML <img> tag (default True).
         
     Returns:
-        str: An HTML img tag with Base64-encoded plot image.
+        str: Either an HTML <img> tag or a raw Base64-encoded image string.
     """
     try:
         # Parse the main JSON configuration
-        config = json.loads(plot_config_json)
+        config = json.loads(input)
     except json.JSONDecodeError as e:
         return f'<p style="color:red;">Error parsing JSON configuration: {e}</p>'
     
@@ -43,6 +45,9 @@ def generate_plot(plot_config_json):
     if not isinstance(data, dict):
         return '<p style="color:red;">Data must be a dictionary with keys as labels and values as numbers.</p>'
     
+    if not data:
+        return '<p style="color:red;">Data dictionary is empty.</p>'
+    
     if not isinstance(labels, list):
         return '<p style="color:red;">Labels must be a list.</p>'
     
@@ -59,10 +64,15 @@ def generate_plot(plot_config_json):
         except (ValueError, TypeError):
             return '<p style="color:red;">All data values must be numeric.</p>'
         
+        # Replace x-axis labels if provided and length matches
+        if labels and len(labels) == len(x_data):
+            x_data = labels
+        
         if plot_type == 'bar':
             bars = ax.bar(x_data, y_data)
             ax.set_xlabel(x_label)
             ax.set_ylabel(y_label)
+            ax.set_ylim(bottom=0)
             
             # Add value labels on top of bars
             for bar, value in zip(bars, y_data):
@@ -74,15 +84,21 @@ def generate_plot(plot_config_json):
             ax.plot(x_data, y_data, marker='o', linewidth=2, markersize=6)
             ax.set_xlabel(x_label)
             ax.set_ylabel(y_label)
+            ax.set_ylim(bottom=0)
             ax.grid(True, alpha=0.3)
             
         elif plot_type == 'pie':
             # For pie charts, use labels parameter if provided, otherwise use data keys
-            pie_labels = labels if len(labels) == len(y_data) else x_data
-            wedges, texts, autotexts = ax.pie(y_data, labels=pie_labels, autopct='%1.1f%%', 
-                                            startangle=90, textprops={'fontsize': 10})
+            pie_labels = labels if len(labels) == len(y_data) else list(data.keys())
+            wedges, texts, autotexts = ax.pie(
+                y_data,
+                labels=pie_labels,
+                autopct='%1.1f%%', 
+                startangle=90,
+                textprops={'fontsize': 10}
+            )
             ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle
-            # Don't set x/y labels for pie charts as they don't make sense
+            # Don't set x/y labels for pie charts
             
         else:
             return f'<p style="color:red;">Invalid plot_type: {plot_type}. Choose "bar", "line", or "pie".</p>'
@@ -94,33 +110,37 @@ def generate_plot(plot_config_json):
         
         # Save plot to a BytesIO buffer in memory
         buf = BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight', dpi=150, 
-                   facecolor='white', edgecolor='none')
+        plt.savefig(
+            buf, 
+            format='png', 
+            bbox_inches='tight', 
+            dpi=150, 
+            facecolor='white', 
+            edgecolor='none'
+        )
         plt.close(fig)  # Close the plot to free up memory
         
         # Encode the image data to a Base64 string
         img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
         
-        # Return HTML img tag with proper styling
-        return f'''
-        <div style="text-align: center; margin: 20px 0;">
-            <img src="data:image/png;base64,{img_base64}" 
-                 style="max-width: 100%; height: auto; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);" 
-                 alt="{title}" />
-        </div>
-        '''
+        if return_html:
+            # More descriptive alt text for accessibility
+            alt_text = title
+            if plot_type == "pie":
+                alt_text += " - " + ", ".join(f"{lbl}: {val}" for lbl, val in zip(pie_labels, y_data))
+            else:
+                alt_text += " - " + ", ".join(f"{lbl}: {val}" for lbl, val in zip(x_data, y_data))
+            
+            return f'''
+            <div style="text-align: center; margin: 20px 0;">
+                <img src="data:image/png;base64,{img_base64}" 
+                     style="max-width: 100%; height: auto; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);" 
+                     alt="{alt_text}" />
+            </div>
+            '''
+        else:
+            return img_base64
          
     except Exception as e:
         plt.close('all')  # Clean up any open figures
         return f'<p style="color:red;">Error generating plot: {str(e)}</p>'
-
-# Example usage:
-# plot_config = {
-#     "data": {"A": 10, "B": 20, "C": 15},
-#     "labels": ["Category A", "Category B", "Category C"],
-#     "plot_type": "bar",
-#     "title": "Sample Bar Chart",
-#     "x_label": "Categories",
-#     "y_label": "Values"
-# }
-# result = generate_plot_single_input(json.dumps(plot_config))