Upload graph_tool.py

graph_tool.py

@@ -0,0 +1,109 @@
+#graph_tool.py
+
+import base64
+import io
+import json
+from typing import Dict, List, Literal, Tuple
+
+import matplotlib.pyplot as plt
+from langchain_core.tools import tool
+
+# 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) from a dictionary of data and returns it
+    as a base64 encoded PNG image artifact.
+
+    Args:
+        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:
+        A tuple containing:
+        - A string message confirming the plot was generated.
+        - A dictionary artifact with the base64 encoded image string and its format.
+    """
+    # --- 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
+    
+    try:
+        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))
+
+        if plot_type == 'bar':
+            # 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)
+            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')
+
+        elif plot_type == 'line':
+            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':
+            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:
+            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)
+        
+        img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+
+        # --- 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')
+        content = f"An unexpected error occurred while generating the plot: {str(e)}"
+        artifact = {"error": str(e)}
+        return content, artifact