Add plotting tools

Gradio_UI.py

@@ -271,6 +271,7 @@ class GradioUI:
             - List all the countries that are there.
             - How many distinct schools are there in each country? Show as a table. Sort by the country names.
             - Show the average download speed and latency of each school in Kenya.
+            - Show the count of schools in each country as a bar diagram.
             '''
             gr.Markdown(title)
             

README.md

@@ -11,4 +11,22 @@ license: apache-2.0
 short_description: 'QoScope '
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# QoScope 🔭
+            
+QoS data analysis thousands of school networks under UNICEF/Giga
+
+---
+
+
+# Installation
+
+```bash
+pip install -r requirements.txt
+python app.py
+```
+
+Add `HF_TOKEN` environment variable, e.g., in the `.env` file.
+
+---
+
+QoScope is released under an Apache 2.0 license.

app.py

@@ -1,14 +1,15 @@
-import tempfile
-
-import matplotlib.pyplot as plt
-import numpy as np
+"""
+QoScope agent with tools.
+"""
 import yaml
-
 import sqlalchemy as sqa
 from Gradio_UI import GradioUI
-from smolagents import CodeAgent, HfApiModel, load_tool, tool, Tool
+from smolagents import CodeAgent, HfApiModel, tool
+
+from tools.diagrams import PlotTool
 from tools.final_answer import FinalAnswerTool
 
+
 DATABASE_PATH = 'sqlite:///resampled_daily_avg.sqlite'
 
 db_engine = sqa.create_engine(DATABASE_PATH, echo=False)
@@ -20,7 +21,8 @@ def run_sql_query(sql_query: str) -> str:
     Run a SQL query on a table in a SQLite database and return the output/result as a string.
     The output contains one row of result(s) in each line.
     The output is simple text without any Markdown styling. An agent should take this plain output
-    and format appropriately when required, e.g., as a Markdown table or summarizing the results in words.
+    and format appropriately when required, e.g., as a Markdown table or summarizing the results
+    in words.
 
     The only available table in the SQLite database is `school_measurements`.
     The table's CREATE statement (DDL) is as follows:
@@ -37,8 +39,8 @@ def run_sql_query(sql_query: str) -> str:
         iso3_format TEXT
     )
 
-    IMPORTANT: You will ONLY execute SELECT queries.
-    You will NEVER execute any other types of SQL queries, e.g., INSERT, DELETE, DROP, and so on, which changes the database in any way.
+    IMPORTANT: You will ONLY execute SELECT queries. You will NEVER execute any other types of
+    SQL queries, e.g., INSERT, DELETE, DROP, and so on, which changes the database in any way.
 
     Args:
         sql_query: An appropriate, correct SQL query for a SQLite database
@@ -54,118 +56,11 @@ def run_sql_query(sql_query: str) -> str:
         rows = con.execute(sqa.text(sql_query))
         for row in rows:
             # Each row is a tuple
-            print(f'\n\n>>>{row=}')
             output += '\n' + str(row)
 
     return output
 
 
-# @tool
-def plot_line_diagram(
-        x_values,
-        y_values_list,
-        labels=None,
-        title='Line Diagram',
-        xlabel='X-axis',
-        ylabel='Y-axis'
-):
-    """
-    Plots a line diagram with one or more y-values.
-
-    :param x_values: List of x-values.
-    :param y_values_list: List of lists containing y-values. Each inner list represents a separate line.
-    :param labels: List of labels for each line (optional).
-    :param title: Title of the plot (default: 'Line Diagram').
-    :param xlabel: Label for the X-axis (default: 'X-axis').
-    :param ylabel: Label for the Y-axis (default: 'Y-axis').
-    """
-    plt.figure(figsize=(10, 6))
-
-    for i, y_values in enumerate(y_values_list):
-        label = labels[i] if labels and i < len(labels) else f'Line {i + 1}'
-        plt.plot(x_values, y_values, label=label)
-
-    plt.title(title)
-    plt.xlabel(xlabel)
-    plt.ylabel(ylabel)
-    plt.legend()
-    plt.grid(True)
-    plt.show()
-
-
-@tool
-def plot_bar_diagram(
-        x_values: list,
-        y_values_list: list[list[int | float]],
-        labels: list[str] | None = None,
-        title: str = 'Bar Diagram',
-        xlabel: str = 'X-axis',
-        ylabel: str = 'Y-axis'
-) -> str:
-    """
-    Plot a bar diagram with one or more y-values and save the image to a temporary file.
-    Return the path to the saved image file. The path can be used to display the image.
-
-    Args:
-        x_values: List of x-values.
-        y_values_list: List of lists containing y-values. Each inner list represents a separate set of bars.
-        labels: List of labels for each set of bars (optional).
-        title: Title of the plot (default: 'Bar Diagram').
-        xlabel: Label for the X-axis (default: 'X-axis').
-        ylabel: Label for the Y-axis (default: 'Y-axis').
-
-    Returns:
-        Path to the saved image file.
-    """
-    bar_width = 0.2
-    n = len(y_values_list)
-
-    # Set positions of bars on X axis
-    r = [np.arange(len(x_values))]
-    for i in range(1, n):
-        r.append([x + bar_width for x in r[i - 1]])
-
-    plt.figure(figsize=(10, 6))
-
-    for i, y_values in enumerate(y_values_list):
-        label = labels[i] if labels and i < len(labels) else f'Set {i + 1}'
-        plt.bar(r[i], y_values, width=bar_width, label=label)
-
-    # Adding xticks
-    plt.xlabel(xlabel)
-    plt.ylabel(ylabel)
-    plt.title(title)
-    plt.xticks(
-        [r + bar_width * (n - 1) / 2 for r in range(len(x_values))],
-        x_values,
-        rotation=45,
-        ha='right'
-    )
-    plt.legend()
-    plt.grid(True, axis='y')
-    plt.tight_layout()
-
-    # Save the plot as an image file in a temporary directory
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-    plt.savefig(temp_file.name)
-    plt.close()
-
-    return temp_file.name
-
-
-# class TextToImageTool(Tool):
-#     description = "This tool creates an image according to a prompt, which is a text description."
-#     name = "image_generator"
-#     inputs = {"prompt": {"type": "string", "description": "The image generator prompt. Don't hesitate to add details in the prompt to make the image look better, like 'high-res, photorealistic', etc."}}
-#     output_type = "image"
-#     model_sdxl = "black-forest-labs/FLUX.1-schnell"
-#     client = InferenceClient(model_sdxl)
-#
-#
-#     def forward(self, prompt):
-#         return self.client.text_to_image(prompt)
-
-
 ### Main block ###
 
 final_answer = FinalAnswerTool()
@@ -176,14 +71,14 @@ code_model = HfApiModel(
     custom_role_conversions=None,
 )
 
-with open('prompts.yaml', 'r') as stream:
+with open('prompts.yaml', 'r', encoding='utf-8') as stream:
     prompt_templates = yaml.safe_load(stream)
 
 agent = CodeAgent(
     model=code_model,
     tools=[
         run_sql_query,
-        plot_bar_diagram,
+        PlotTool(),
         final_answer,
     ],
     max_steps=6,

requirements.txt

@@ -4,3 +4,4 @@ pandas
 SQLAlchemy
 gradio
 matplotlib
+Pillow

tools/diagrams.py

@@ -0,0 +1,201 @@
+import io
+import tempfile
+
+import numpy as np
+from PIL import Image
+from PIL.ImageFile import ImageFile
+from matplotlib import pyplot as plt
+from smolagents.tools import Tool
+
+
+def _plot_line_diagram(
+        x_values: list,
+        y_values_list: list[list[int | float]],
+        labels: list[str] | None = None,
+        title: str = 'Bar Diagram',
+        xlabel: str = 'X-axis',
+        ylabel: str = 'Y-axis'
+) -> str:
+    """
+    Plot a line diagram with one or more y-values and save the image to a temporary file.
+    Return the path to the saved image file.
+
+    :param x_values: List of x-values.
+    :param y_values_list: List of lists containing y-values.
+     Each inner list represents a separate line.
+    :param labels: List of labels for each line (optional).
+    :param title: Title of the plot (default: 'Line Diagram').
+    :param xlabel: Label for the X-axis (default: 'X-axis').
+    :param ylabel: Label for the Y-axis (default: 'Y-axis').
+    :return: Path to the saved image file.
+    """
+
+    plt.figure(figsize=(10, 6))
+
+    for i, y_values in enumerate(y_values_list):
+        label = labels[i] if labels and i < len(labels) else f'Line {i+1}'
+        plt.plot(x_values, y_values, label=label)
+
+    plt.title(title)
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.legend()
+    plt.grid(True)
+
+    # Save the plot as an image file in a temporary directory
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    plt.savefig(temp_file.name)
+    plt.close()
+
+    return temp_file.name
+
+
+def _plot_bar_diagram(
+        x_values: list,
+        y_values_list: list[list[int | float]],
+        labels: list[str] | None = None,
+        title: str = 'Bar Diagram',
+        xlabel: str = 'X-axis',
+        ylabel: str = 'Y-axis'
+) -> str:
+    """
+    Plot a bar diagram with one or more y-values and save the image to a temporary file.
+    Return the path to the saved image file.
+
+    :param x_values: List of x-values.
+    :param y_values_list: List of lists containing y-values.
+     Each inner list represents a separate set of bars.
+    :param labels: List of labels for each set of bars (optional).
+    :param title: Title of the plot (default: 'Bar Diagram').
+    :param xlabel: Label for the X-axis (default: 'X-axis').
+    :param ylabel: Label for the Y-axis (default: 'Y-axis').
+    :return: Path to the saved image file.
+    """
+
+    bar_width = 0.2
+    n = len(y_values_list)
+
+    # Set positions of bars on X axis
+    r = [np.arange(len(x_values))]
+    for i in range(1, n):
+        r.append([x + bar_width for x in r[i - 1]])
+
+    plt.figure(figsize=(10, 6))
+
+    for i, y_values in enumerate(y_values_list):
+        label = labels[i] if labels and i < len(labels) else f'Set {i + 1}'
+        plt.bar(r[i], y_values, width=bar_width, label=label)
+
+    # Adding xticks
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.title(title)
+    plt.xticks(
+        [r + bar_width * (n - 1) / 2 for r in range(len(x_values))],
+        x_values,
+        rotation=45,
+        ha='right'
+    )
+    plt.legend()
+    plt.grid(True, axis='y')
+    plt.tight_layout()
+
+    # Save the plot as an image file in a temporary directory
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+    plt.savefig(temp_file.name)
+    plt.close()
+
+    return temp_file.name
+
+
+class PlotTool(Tool):
+    """
+    A tool to plot bar and line diagrams and return the images.
+    """
+
+    name = 'plot_bar_line_diagrams'
+    description = (
+        'Plot a bar or line diagram with one or more y-values and save the image.'
+        ' Return the saved image file as `ImageFile`.'
+        ' An agent must take this `ImageFile` and display the image.'
+    )
+    inputs = {
+        'plot_type': {
+            'type': 'string',
+            'description': 'The type of plot. Only two values are valid: `bar` and `line`'
+        },
+        'x_values': {'type': 'array', 'description': 'List of x-values.'},
+        'y_values_list': {
+            'type': 'array',
+            'description': (
+                'A list of lists containing y-values (numbers).'
+                ' Each inner list represents a separate set of bars.'
+                ' The input type is list[list[int | float]]'
+            )
+        },
+        'labels': {
+            'type': 'array',
+            'nullable': True,
+            'description': (
+                'A list of labels for each set of bars (optional). Defaults to `None`.'
+                ' If provided, the length of `labels` must be equal to the length of `x_values`.'
+            )
+        },
+        'title': {
+            'type': 'string',
+            'nullable': True,
+            'description': 'Title of the plot (default: "Bar Diagram")'
+        },
+        'xlabel': {
+            'type': 'string',
+            'nullable': True,
+            'description': 'Label for the X-axis (default: "X-axis")'
+        },
+        'ylabel': {
+            'type': 'string',
+            'nullable': True,
+            'description': 'Label for the Y-axis (default: "Y-axis").'
+        },
+    }
+    output_type = 'image'
+
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def forward(
+            self,
+            plot_type: str,
+            x_values: list,
+            y_values_list: list[list[int | float]],
+            labels: list[str] | None = None,
+            title: str = 'Bar Diagram',
+            xlabel: str = 'X-axis',
+            ylabel: str = 'Y-axis'
+    ) -> ImageFile:
+
+        if plot_type == 'bar':
+            img_file_name = _plot_bar_diagram(
+                x_values=x_values,
+                y_values_list=y_values_list,
+                labels=labels,
+                xlabel=xlabel,
+                ylabel=ylabel,
+                title=title
+            )
+        elif plot_type == 'line':
+            img_file_name = _plot_line_diagram(
+                x_values=x_values,
+                y_values_list=y_values_list,
+                labels=labels,
+                xlabel=xlabel,
+                ylabel=ylabel,
+                title=title
+            )
+        else:
+            img_file_name = None
+
+        if img_file_name:
+            with open(img_file_name, 'rb') as in_file:
+                return Image.open(io.BytesIO(in_file.read()))
+        else:
+            return None