Redesigned for LangGraph Agentic workflow

graph_tool.py

@@ -1,146 +1,107 @@
 import base64
-from io import BytesIO
-import matplotlib.pyplot as plt
-import numpy as np
+import io
 import json
+from typing import Dict, List, Literal, Tuple
+
+import matplotlib.pyplot as plt
+from langchain_core.tools import tool
 
-def generate_plot(input: str, return_html: bool = True):
+# Use the @tool decorator and specify the "content_and_artifact" response format.
+@tool(response_format="content_and_artifact")
+def generate_plot(
+    data: Dict[str, float],
+    plot_type: Literal["bar", "line", "pie"],
+    title: str = "Generated Plot",
+    labels: List[str] = None,
+    x_label: str = "",
+    y_label: str = ""
+) -> Tuple:
     """
-    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.
-    
+    Generates a plot (bar, line, or pie) from a dictionary of data and returns it
+    as a base64 encoded PNG image artifact.
+
     Args:
-        input (str): JSON string containing all plot configuration:
-            {
-                "data": {"key1": value1, "key2": value2, ...},
-                "labels": ["label1", "label2", ...],
-                "plot_type": "bar|line|pie",
-                "title": "Plot Title",
-                "x_label": "X Axis Label" (optional),
-                "y_label": "Y Axis Label" (optional)
-            }
-        return_html (bool): Whether to return full HTML <img> tag (default True).
-        
+        data (Dict[str, float]): A dictionary where keys are labels and values are the numeric data to plot.
+        plot_type (Literal["bar", "line", "pie"]): The type of plot to generate.
+        title (str): The title for the plot.
+        labels (List[str]): Optional list of labels to use for the x-axis or pie slices. If not provided, data keys are used.
+        x_label (str): The label for the x-axis (for bar and line charts).
+        y_label (str): The label for the y-axis (for bar and line charts).
+
     Returns:
-        str: Either an HTML <img> tag or a raw Base64-encoded image string.
+        A tuple containing:
+        - A string message confirming the plot was generated.
+        - A dictionary artifact with the base64 encoded image string and its format.
     """
-    try:
-        # Parse the main JSON configuration
-        config = json.loads(input)
-    except json.JSONDecodeError as e:
-        return f'<p style="color:red;">Error parsing JSON configuration: {e}</p>'
+    # --- Input Validation ---
+    if not isinstance(data, dict) or not data:
+        content = "Error: Data must be a non-empty dictionary."
+        artifact = {"error": content}
+        return content, artifact
     
-    # Extract parameters with defaults
     try:
-        data = config.get("data", {})
-        labels = config.get("labels", [])
-        plot_type = config.get("plot_type", "bar")
-        title = config.get("title", "Untitled Plot")
-        x_label = config.get("x_label", "")
-        y_label = config.get("y_label", "")
-    except Exception as e:
-        return f'<p style="color:red;">Error extracting configuration parameters: {e}</p>'
-    
-    # Validate inputs
-    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>'
+        y_data = [float(val) for val in data.values()]
+    except (ValueError, TypeError):
+        content = "Error: All data values must be numeric."
+        artifact = {"error": content}
+        return content, artifact
+
+    x_data = list(data.keys())
     
+    # --- Plot Generation ---
     try:
         fig, ax = plt.subplots(figsize=(10, 6))
-        
-        # Extract keys and values from the data dictionary
-        x_data = list(data.keys())
-        y_data = list(data.values())
-        
-        # Ensure y_data contains numeric values
-        try:
-            y_data = [float(val) for val in y_data]
-        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)
+            # Use provided labels if they match the data length, otherwise use data keys
+            bar_labels = labels if labels and len(labels) == len(x_data) else x_data
+            bars = ax.bar(bar_labels, 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):
                 height = bar.get_height()
-                ax.text(bar.get_x() + bar.get_width()/2., height,
-                       f'{value}', ha='center', va='bottom')
-                       
+                ax.text(bar.get_x() + bar.get_width()/2., height, f'{value}', ha='center', va='bottom')
+
         elif plot_type == 'line':
-            ax.plot(x_data, y_data, marker='o', linewidth=2, markersize=6)
+            line_labels = labels if labels and len(labels) == len(x_data) else x_data
+            ax.plot(line_labels, y_data, marker='o')
             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 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
-            
+            pie_labels = labels if labels and len(labels) == len(y_data) else list(data.keys())
+            ax.pie(y_data, labels=pie_labels, autopct='%1.1f%%', startangle=90)
+            ax.axis('equal')
+
         else:
-            return f'<p style="color:red;">Invalid plot_type: {plot_type}. Choose "bar", "line", or "pie".</p>'
-        
-        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)
-        
-        # Improve layout
+            content = f"Error: Invalid plot_type '{plot_type}'. Choose 'bar', 'line', or 'pie'."
+            artifact = {"error": content}
+            return content, artifact
+
+        ax.set_title(title, fontsize=14, fontweight='bold')
         plt.tight_layout()
+
+        # --- In-Memory Image Conversion ---
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', dpi=150)
+        plt.close(fig)
+        buf.seek(0)
         
-        # 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.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')
-        
-        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
-         
+
+        # --- Return Content and Artifact ---
+        content = f"Successfully generated a {plot_type} plot titled '{title}'."
+        artifact = {
+            "base64_image": img_base64,
+            "format": "png"
+        }
+        return content, artifact
+
     except Exception as e:
-        plt.close('all')  # Clean up any open figures
-        return f'<p style="color:red;">Error generating plot: {str(e)}</p>'
+        plt.close('all')
+        content = f"An unexpected error occurred while generating the plot: {str(e)}"
+        artifact = {"error": str(e)}
+        return content, artifact