Mikasa06/AI_Agent_ChartMaster

1from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool2import datetime3import requests4import pytz5import yaml6from tools.final_answer import FinalAnswerTool7import matplotlib.pyplot as plt8import io9from PIL import Image as PILImage10from IPython.display import display, Image11import numpy as np12 13 14from Gradio_UI import GradioUI15 16# Below is an example of a tool that does nothing. Amaze us with your creativity !17@tool18def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return type19    #Keep this format for the description / args / args description but feel free to modify the tool20    """A tool that does nothing yet 21    Args:22        arg1: the first argument23        arg2: the second argument24    """25    return "What magic will you build ?"26 27@tool28def get_current_time_in_timezone(timezone: str) -> str:29    """A tool that fetches the current local time in a specified timezone.30    Args:31        timezone: A string representing a valid timezone (e.g., 'America/New_York').32    """33    try:34        # Create timezone object35        tz = pytz.timezone(timezone)36        # Get current time in that timezone37        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")38        return f"The current local time in {timezone} is: {local_time}"39    except Exception as e:40        return f"Error fetching time for timezone '{timezone}': {str(e)}"41 42 43 44 45 46# Bar chart tool47 48@tool49def generate_bar_chart(50    x_values: list[str], y_values: list[int], title: str = 'Bar Chart',51    x_label: str = 'X-Axis', y_label: str = 'Y-Axis', show_labels: bool = False52) -> PILImage:53    """Generates a bar chart from the provided x and y values and returns a PILImage object.54    Args:55        x_values: A list of string values for the x-axis.56        y_values: A list of numerical values for the y-axis.57        title: Title for the bar plot.58        x_label: Label for the x-axis.59        y_label: Label for the y-axis.60        show_labels: Whether to display value labels on top of the bars.61    Returns:62        A PIL Image object containing the generated bar chart.63    """64    try:65        if len(x_values) != len(y_values):66            raise ValueError("x_values and y_values must have the same length.")67        68        plt.figure(figsize=(8, 6))69        bars = plt.bar(x_values, y_values, color=plt.cm.Paired.colors, edgecolor='black')70        71        if show_labels:72            for bar in bars:73                plt.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),74                         f'{bar.get_height()}', ha='center', va='bottom',75                         fontsize=10, fontweight='bold')76        77        plt.xlabel(x_label)78        plt.ylabel(y_label)79        plt.title(title)80        81        # Save the plot to a BytesIO buffer instead of a file82        img_buffer = io.BytesIO()83        plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')84        plt.close()  # Close the plot to free memory85        86        img_buffer.seek(0)  # Move to the beginning of the buffer87        return PILImage.open(img_buffer)  # Return the image object88    except Exception as e:89        print(f"Error generating bar chart: {str(e)}")90        return None91 92 93# Scatter plot tool94@tool95def generate_scatter_plot_with_labels(96    x_values: list[float], y_values: list[float], labels: list[str],97    title: str = 'Scatter Plot', x_label: str = 'X-Axis', y_label: str = 'Y-Axis'98) -> PILImage:99    """Generates a scatter plot from the provided x and y values and labels each point. Returns a PILImage object.100    Args:101        x_values: A list of numerical values for the x-axis.102        y_values: A list of numerical values for the y-axis.103        labels: A list of labels corresponding to each point.104        title: Title for the scatter plot.105        x_label: Label for the x-axis.106        y_label: Label for the y-axis.107    Returns:108        A PIL Image object containing the scatter plot.109    """110    try:111        if len(x_values) != len(y_values) or len(x_values) != len(labels):112            raise ValueError("x_values, y_values, and labels must all have the same length.")113        114        plt.figure(figsize=(6, 4))115        plt.scatter(x_values, y_values, color='blue', edgecolors='black', marker='o')116        117        for i, label in enumerate(labels):118            plt.text(x_values[i], y_values[i], label, fontsize=9, ha='right', color='red')119        120        plt.title(title)121        plt.xlabel(x_label)122        plt.ylabel(y_label)123        124        img_buffer = io.BytesIO()125        plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')126        plt.close()127        128        img_buffer.seek(0)129        return PILImage.open(img_buffer)130    except Exception as e:131        print(f"Error generating scatter plot with labels: {str(e)}")132        return None133 134# Line plot tool135@tool136def generate_line_plot(137    x_values: list[float], y_values: list[float],138    title: str = 'Line Plot', x_label: str = 'X-Axis', y_label: str = 'Y-Axis'139) -> PILImage:140    """Generates a line plot from the provided x and y values. Returns a PILImage object.141    Args:142        x_values: A list of numerical values for the x-axis.143        y_values: A list of numerical values for the y-axis.144        title: Title for the line plot.145        x_label: Label for the x-axis.146        y_label: Label for the y-axis.147    Returns:148        A PIL Image object containing the line plot.149    """150    try:151        if len(x_values) != len(y_values):152            raise ValueError("x_values and y_values must have the same length.")153        154        plt.figure(figsize=(6, 4))155        plt.plot(x_values, y_values, marker='o', linestyle='-', color='b')156        plt.title(title)157        plt.xlabel(x_label)158        plt.ylabel(y_label)159        plt.grid(True)160        161        img_buffer = io.BytesIO()162        plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')163        plt.close()164        165        img_buffer.seek(0)166        return PILImage.open(img_buffer)167    except Exception as e:168        print(f"Error generating line plot: {str(e)}")169        return None170 171# Histogram tool172@tool173def generate_histogram(174    data: list[float], bins: int = 10,175    title: str = 'Histogram', x_label: str = 'Values', y_label: str = 'Frequency'176) -> PILImage:177    """Generates a histogram from the provided data and returns a PILImage object.178    Args:179        data: A list of numerical values.180        bins: Number of bins for the histogram.181        title: Title for the histogram.182        x_label: Label for the x-axis.183        y_label: Label for the y-axis.184    Returns:185        A PIL Image object containing the histogram plot.186    """187    try:188        if not data:189            raise ValueError("Data list is empty.")190        191        plt.figure(figsize=(6, 4))192        plt.hist(data, bins=bins, color='purple', edgecolor='black', alpha=0.7)193        194        plt.title(title)195        plt.xlabel(x_label)196        plt.ylabel(y_label)197        198        img_buffer = io.BytesIO()199        plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')200        plt.close()201        202        img_buffer.seek(0)203        return PILImage.open(img_buffer)204    except Exception as e:205        print(f"Error generating histogram: {str(e)}")206        return None207 208# Pie chart tool209@tool210def generate_pie_chart(211    labels: list[str], values: list[float], title: str = 'Pie Chart'212) -> PILImage:213    """Generates a pie chart from the provided labels and values with dynamic colors. Returns a PILImage object.214    Args:215        labels: A list of category labels.216        values: A list of numerical values for each category.217        title: Title for the pie chart.218    Returns:219        A PIL Image object containing the pie chart.220    """221    try:222        if len(labels) != len(values):223            raise ValueError("Labels and values must have the same length.")224        if any(v < 0 for v in values):225            raise ValueError("Values must be non-negative.")226        if sum(values) == 0:227            raise ValueError("Sum of values must be greater than zero.")228        229        cmap = plt.get_cmap('tab10')230        colors = [cmap(i / len(values)) for i in range(len(values))]231        232        fig, ax = plt.subplots(figsize=(6, 6))233        ax.pie(values, labels=labels, autopct='%1.1f%%', startangle=140,234               colors=colors, textprops={'fontsize': 10})235        ax.set_title(title, fontsize=12)236        237        img_buffer = io.BytesIO()238        fig.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')239        plt.close(fig)240        241        img_buffer.seek(0)242        return PILImage.open(img_buffer)243    except Exception as e:244        print(f"Error generating pie chart: {str(e)}")245        return None246 247# Box plot tool248@tool249def generate_box_plot(250    data: list[list[float]], labels: list[str] = None,251    title: str = 'Box Plot', y_label: str = 'Values'252) -> PILImage:253    """Generates a box plot from the provided data and returns a PILImage object.254    Args:255        data: A list of numerical lists representing different categories.256        labels: A list of labels corresponding to each dataset (optional).257        title: Title for the box plot.258        y_label: Label for the y-axis.259    Returns:260        A PIL Image object containing the box plot.261    """262    try:263        if not all(isinstance(category, list) and all(isinstance(x, (int, float)) for x in category) for category in data):264            raise ValueError("Data must be a list of numerical lists.")265        if labels and len(labels) != len(data):266            raise ValueError("Labels length must match the number of data categories.")267        268        plt.figure(figsize=(8, 6))269        plt.boxplot(data)270        271        if labels:272            plt.xticks(range(1, len(labels) + 1), labels, rotation=20)273        274        plt.title(title)275        plt.ylabel(y_label)276        277        img_buffer = io.BytesIO()278        plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')279        plt.close()280        281        img_buffer.seek(0)282        return PILImage.open(img_buffer)283    except Exception as e:284        print(f"Error generating box plot: {str(e)}")285        return None286 287 288 289 290final_answer = FinalAnswerTool()291 292# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:293# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 294 295model = HfApiModel(296max_tokens=2096,297temperature=0.5,298# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',299model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded300custom_role_conversions=None,301)302 303 304# Import tool from Hub305image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)306 307with open("prompts.yaml", 'r') as stream:308    prompt_templates = yaml.safe_load(stream)309    310agent = CodeAgent(311    model=model,312    tools=[final_answer,generate_bar_chart,generate_scatter_plot_with_labels,generate_line_plot,generate_histogram,generate_pie_chart,generate_box_plot], ## add your tools here (don't remove final answer)313    max_steps=6,314    verbosity_level=2,315    grammar=None,316    planning_interval=None,317    name=None,318    description=None,319    prompt_templates=prompt_templates320)321 322GradioUI(agent).launch()323 324