feat: add rolling stats, K-Means clustering, and hypothesis testing

- Add rolling_stats_func: overlays rolling mean/std/min/max on a time
series using plotly.graph_objects; supports optional category grouping
- Add kmeans_clustering_func: StandardScaler + KMeans (sklearn), scatter
plot coloured by cluster, centroid summary table
- Add hypothesis_test_func: independent/one-sample t-test (scipy) and
chi-square test of independence; returns styled HTML results table
- Register all three in tool schemas, __init__.py, available_functions,
and system message; add scipy and scikit-learn to requirements.txt
- Refactor chart_tools.py: extract _LAYOUT_PARAM and _TRACE_STYLE_PARAM
constants to eliminate ~8x copy-paste; remove layout from required
arrays on line/bar/pie; fix hardcoded 'SQLite' references; standardise
all tool descriptions to consistent what/when/returns format
- Refactor stats_tools.py: same description standardisation
- Shorten tools_desc in system message from 13 lines to 2 sentences
(LLM already receives full schema descriptions separately)
- Fix copy-paste bug: sql example-question generator was using MongoDB
system persona instead of PostgreSQL
- Fix ROOT_URL defaulting to None when env var unset, which caused
chart iframe src to contain literal 'None' in the path

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

functions/__init__.py

@@ -1,15 +1,17 @@
 from .query_functions import graphql_schema_query, graphql_csv_query, query_func
 from .chart_functions import table_generation_func, scatter_chart_generation_func, \
     line_chart_generation_func, bar_chart_generation_func, pie_chart_generation_func, \
-    histogram_generation_func, box_chart_generation_func, correlation_heatmap_func, scatter_chart_fig
+    histogram_generation_func, box_chart_generation_func, correlation_heatmap_func, \
+    scatter_chart_fig, rolling_stats_func
 from .chat_functions import example_question_generator, chatbot_func
-from .stat_functions import regression_func, descriptive_stats_func
+from .stat_functions import regression_func, descriptive_stats_func, \
+    kmeans_clustering_func, hypothesis_test_func
 
 __all__ = [
     "query_func", "graphql_schema_query", "graphql_csv_query",
     "table_generation_func", "scatter_chart_generation_func", "line_chart_generation_func",
     "bar_chart_generation_func", "pie_chart_generation_func", "histogram_generation_func",
-    "box_chart_generation_func", "correlation_heatmap_func",
-    "regression_func", "descriptive_stats_func",
+    "box_chart_generation_func", "correlation_heatmap_func", "rolling_stats_func",
+    "regression_func", "descriptive_stats_func", "kmeans_clustering_func", "hypothesis_test_func",
     "scatter_chart_fig", "example_question_generator", "chatbot_func",
 ]

functions/chart_functions.py

@@ -9,7 +9,7 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-root_url = os.getenv("ROOT_URL")
+root_url = os.getenv("ROOT_URL", "")
 
 def llm_chart_data_scrub(data, layout):
    #Processing data to account for variation from LLM
@@ -427,6 +427,75 @@ def correlation_heatmap_func(session_hash, session_folder, columns: List[str]=[]
         return {"reply": f"There was an error generating the correlation heatmap. Error: {e}. You should probably try again."}
 
 
+def rolling_stats_func(x_column: str, y_column: str, session_hash, session_folder,
+                       window: int = 7, stats: List[str] = ["mean"],
+                       layout: List[dict] = [{}], category: str = "", **kwargs):
+    try:
+        import plotly.graph_objects as go
+
+        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
+        chart_path = f'{dir_path}/chart.html'
+        csv_query_path = f'{dir_path}/query.csv'
+
+        df = pd.read_csv(csv_query_path)
+
+        try:
+            df[x_column] = pd.to_datetime(df[x_column])
+        except Exception:
+            pass
+        df = df.sort_values(x_column)
+
+        valid_stats = {"mean", "std", "min", "max"}
+        selected_stats = [s for s in stats if s in valid_stats] or ["mean"]
+
+        fig = go.Figure()
+
+        groups = df[category].unique().tolist() if category and category in df.columns else [None]
+
+        for group in groups:
+            group_df = df[df[category] == group] if group is not None else df
+            prefix = f"{group} — " if group is not None else ""
+
+            fig.add_trace(go.Scatter(
+                x=group_df[x_column].values, y=group_df[y_column].values,
+                mode="lines", name=f"{prefix}{y_column} (raw)",
+                opacity=0.35, line=dict(width=1)
+            ))
+
+            rolling_obj = group_df[y_column].rolling(window)
+            for stat in selected_stats:
+                rolled = getattr(rolling_obj, stat)()
+                fig.add_trace(go.Scatter(
+                    x=group_df[x_column].values, y=rolled.values,
+                    mode="lines", name=f"{prefix}Rolling {stat.capitalize()} (w={window})",
+                    line=dict(width=2.5)
+                ))
+
+        fig.update_layout(
+            title=f"Rolling Statistics (window={window}) — {y_column}",
+            xaxis_title=x_column,
+            yaxis_title=y_column,
+            font=dict(family="Inter, system-ui, sans-serif"),
+            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        )
+
+        _, layout_dict = llm_chart_data_scrub({}, layout)
+        if layout_dict:
+            fig.update_layout(**layout_dict)
+
+        pio.write_html(fig, chart_path, full_html=False)
+
+        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
+        iframe = ('Please display this iframe: <div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="'
+                  + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>')
+        return {"reply": iframe}
+
+    except Exception as e:
+        print("ROLLING STATS ERROR")
+        print(e)
+        return {"reply": f"There was an error generating the rolling statistics chart. Error: {e}. You should probably try again."}
+
+
 def table_generation_func(session_hash, session_folder, **kwargs):
     print("TABLE GENERATION")
     try:

functions/chat_functions.py

@@ -27,7 +27,7 @@ def example_question_message(data_source, name, titles, schema):
                         data insights. Return nothing more than the array of questions because I need that specific data structure
                         to process your response. No other response type or data structure will work."""],
 
-        'sql' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
+        'sql' : [f"You are a helpful and knowledgeable agent who has access to a PostgreSQL database called {name}.",
                  f"""We have a PostgreSQL database with the following tables: {titles}. 
                         We also have an AI agent with access to the same database that will be performing data analysis.
                         Please return an array of seven strings, each one being a question for our data analysis agent
@@ -82,18 +82,11 @@ def system_message(data_source, titles, schema=""):
     print("TITLES")
     print(titles)
 
-    tools_desc = """
-                    You also have access to a table function, called table_generation_func, that takes query results from query.csv and returns a formatted inline table.
-                    You also have access to a scatter plot function, called scatter_chart_generation_func, that uses plotly dictionaries to generate a scatter plot and returns an iframe to display.
-                    You also have access to a line chart function, called line_chart_generation_func, that uses plotly dictionaries to generate a line chart and returns an iframe to display.
-                    You also have access to a bar chart function, called bar_chart_generation_func, that uses plotly dictionaries to generate a bar chart and returns an iframe to display.
-                    You also have access to a pie chart function, called pie_chart_generation_func, that uses plotly dictionaries to generate a pie chart and returns an iframe to display.
-                    You also have access to a histogram function, called histogram_generation_func, that uses plotly dictionaries to generate a histogram and returns an iframe to display.
-                    You also have access to a box plot function, called box_chart_generation_func, that shows the distribution of a numeric column optionally grouped by a category, and returns an iframe to display.
-                    You also have access to a correlation heatmap function, called correlation_heatmap_func, that computes pairwise correlations between numeric columns and returns a colour-coded heatmap iframe to display.
-                    You also have access to a descriptive statistics function, called descriptive_stats_func, that returns a summary statistics table (count, mean, std, min, quartiles, max) for numeric columns.
-                    You also have access to a linear regression function, called regression_func, that takes column names for independent and dependent variables and returns a regression summary and chart iframe.
-                    Could you please always display the generated charts, tables, and visualizations as part of your output?"""
+    tools_desc = (
+        " You have access to tools for querying the data source, generating charts and visualisations,"
+        " and performing statistical analyses — use them proactively whenever they would help answer the user's question."
+        " Always display any charts, tables, and visualisations inline in your responses by outputting the returned HTML verbatim."
+    )
 
     system_message_dict = {
         'file_upload': (
@@ -128,8 +121,9 @@ def chatbot_func(message, history, session_hash, data_source, titles, schema, *a
         table_generation_func, regression_func, descriptive_stats_func,
         scatter_chart_generation_func, line_chart_generation_func, bar_chart_generation_func,
         pie_chart_generation_func, histogram_generation_func,
-        box_chart_generation_func, correlation_heatmap_func,
+        box_chart_generation_func, correlation_heatmap_func, rolling_stats_func,
         query_func, graphql_schema_query, graphql_csv_query,
+        kmeans_clustering_func, hypothesis_test_func,
     )
     import tools.tools as tools
 
@@ -145,8 +139,11 @@ def chatbot_func(message, history, session_hash, data_source, titles, schema, *a
         "histogram_generation_func": histogram_generation_func,
         "box_chart_generation_func": box_chart_generation_func,
         "correlation_heatmap_func": correlation_heatmap_func,
+        "rolling_stats_func": rolling_stats_func,
         "regression_func": regression_func,
         "descriptive_stats_func": descriptive_stats_func,
+        "kmeans_clustering_func": kmeans_clustering_func,
+        "hypothesis_test_func": hypothesis_test_func,
     }
     
     if message_dict[session_hash][data_source] != None:

functions/stat_functions.py

@@ -5,12 +5,12 @@ from utils import TEMP_DIR
 import plotly.express as px
 import plotly.io as pio
 import os
-from functions import scatter_chart_fig
+from functions.chart_functions import scatter_chart_fig, llm_chart_data_scrub
 from dotenv import load_dotenv
 
 load_dotenv()
 
-root_url = os.getenv("ROOT_URL")
+root_url = os.getenv("ROOT_URL", "")
 
 def descriptive_stats_func(session_hash, session_folder, columns: List[str]=[], **kwargs):
     print("DESCRIPTIVE STATISTICS")
@@ -63,6 +63,190 @@ def descriptive_stats_func(session_hash, session_folder, columns: List[str]=[],
         return {"reply": f"There was an error generating descriptive statistics. Error: {e}. You should probably try again."}
 
 
+def kmeans_clustering_func(feature_columns: List[str], x_column: str, y_column: str,
+                           session_hash, session_folder, n_clusters: int = 3,
+                           layout: List[dict] = [{}], **kwargs):
+    print("KMEANS CLUSTERING")
+    try:
+        from sklearn.cluster import KMeans
+        from sklearn.preprocessing import StandardScaler
+        from html import escape
+
+        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
+        chart_path = f'{dir_path}/chart.html'
+        csv_query_path = f'{dir_path}/query.csv'
+
+        df = pd.read_csv(csv_query_path)
+
+        feature_df = df[feature_columns].select_dtypes(include='number').dropna()
+        if feature_df.shape[1] < 1:
+            return {"reply": "No numeric feature columns found for clustering. Please refine your query to include numeric columns."}
+
+        X_scaled = StandardScaler().fit_transform(feature_df)
+        labels = KMeans(n_clusters=n_clusters, random_state=42, n_init=10).fit_predict(X_scaled)
+
+        df_clustered = df.loc[feature_df.index].copy()
+        df_clustered['Cluster'] = [f'Cluster {l}' for l in labels]
+
+        fig = px.scatter(
+            df_clustered, x=x_column, y=y_column, color='Cluster',
+            title=f'K-Means Clustering (k={n_clusters})',
+        )
+        fig.update_layout(font=dict(family='Inter, system-ui, sans-serif'))
+
+        _, layout_dict = llm_chart_data_scrub({}, layout)
+        if layout_dict:
+            fig.update_layout(**layout_dict)
+
+        pio.write_html(fig, chart_path, full_html=False)
+
+        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
+        iframe = ('Please display this iframe: <div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="'
+                  + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>')
+
+        cluster_summary = df_clustered.groupby('Cluster')[feature_columns].mean().round(3)
+        header_cells = '<th style="background:#1e40af;">Cluster</th>' + ''.join(
+            f'<th>{escape(str(col))}</th>' for col in cluster_summary.columns
+        )
+        row_html = [
+            '<tr>'
+            + f'<td style="font-weight:600;color:#1e40af;background:#eff6ff;white-space:nowrap;">{escape(str(idx))}</td>'
+            + ''.join(f'<td>{escape(str(val))}</td>' for val in row)
+            + '</tr>'
+            for idx, row in cluster_summary.iterrows()
+        ]
+        style = (
+            '<style>'
+            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
+            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
+            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
+            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
+            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
+            '.vda-table tbody tr:last-child td{border-bottom:none;}'
+            '</style>'
+        )
+        summary_table = (
+            '<div class="vda-table-wrap"><table class="vda-table">'
+            f'<thead><tr>{header_cells}</tr></thead>'
+            '<tbody>' + '\n'.join(row_html) + '</tbody>'
+            '</table></div>'
+        )
+
+        return {"reply": f'{iframe}\n\n**Cluster Centroids (feature means per cluster):**\n{style}{summary_table}'}
+
+    except Exception as e:
+        print("KMEANS CLUSTERING ERROR")
+        print(e)
+        return {"reply": f"There was an error running K-Means clustering. Error: {e}. You should probably try again."}
+
+
+def hypothesis_test_func(test_type: str, column: str, session_hash, session_folder,
+                         column2: str = "", group_column: str = "",
+                         group_values: List[str] = [], pop_mean: float = 0.0, **kwargs):
+    print("HYPOTHESIS TEST")
+    try:
+        from scipy import stats
+        from html import escape
+
+        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
+        csv_query_path = f'{dir_path}/query.csv'
+        df = pd.read_csv(csv_query_path)
+
+        if test_type == "t_test_independent":
+            if not group_column or group_column not in df.columns:
+                return {"reply": "Please specify a valid group_column for the independent t-test."}
+            unique_groups = df[group_column].dropna().unique().tolist()
+            if group_values and len(group_values) == 2:
+                g1_label, g2_label = group_values[0], group_values[1]
+            elif len(unique_groups) == 2:
+                g1_label, g2_label = unique_groups[0], unique_groups[1]
+            else:
+                return {"reply": f"For an independent t-test, exactly 2 groups are needed. Found: {unique_groups}. Specify group_values with 2 entries."}
+
+            g1 = df[df[group_column] == g1_label][column].dropna()
+            g2 = df[df[group_column] == g2_label][column].dropna()
+            t_stat, p_value = stats.ttest_ind(g1, g2)
+
+            result_rows = [
+                ("Test", "Independent Samples T-Test"),
+                ("Column", column),
+                ("Group Column", group_column),
+                (f"Group 1", str(g1_label)),
+                (f"Group 2", str(g2_label)),
+                (f"Group 1 Mean (n={len(g1)})", f"{g1.mean():.4f}"),
+                (f"Group 2 Mean (n={len(g2)})", f"{g2.mean():.4f}"),
+                ("T-Statistic", f"{t_stat:.4f}"),
+                ("P-Value", f"{p_value:.6f}"),
+                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
+            ]
+            title = f"T-Test: {column} by {group_column}"
+
+        elif test_type == "t_test_one_sample":
+            sample = df[column].dropna()
+            t_stat, p_value = stats.ttest_1samp(sample, pop_mean)
+            result_rows = [
+                ("Test", "One-Sample T-Test"),
+                ("Column", column),
+                ("Hypothesized Mean (μ₀)", f"{pop_mean:.4f}"),
+                (f"Sample Mean (n={len(sample)})", f"{sample.mean():.4f}"),
+                ("Sample Std Dev", f"{sample.std():.4f}"),
+                ("T-Statistic", f"{t_stat:.4f}"),
+                ("P-Value", f"{p_value:.6f}"),
+                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
+            ]
+            title = f"One-Sample T-Test: {column} vs μ={pop_mean}"
+
+        elif test_type == "chi_square":
+            if not column2 or column2 not in df.columns:
+                return {"reply": "Please specify a valid column2 for the chi-square test."}
+            contingency = pd.crosstab(df[column], df[column2])
+            chi2, p_value, dof, _ = stats.chi2_contingency(contingency)
+            result_rows = [
+                ("Test", "Chi-Square Test of Independence"),
+                ("Column 1", column),
+                ("Column 2", column2),
+                ("Chi-Square Statistic", f"{chi2:.4f}"),
+                ("Degrees of Freedom", str(dof)),
+                ("P-Value", f"{p_value:.6f}"),
+                ("Significant at α=0.05", "Yes ✓" if p_value < 0.05 else "No ✗"),
+            ]
+            title = f"Chi-Square: {column} × {column2}"
+
+        else:
+            return {"reply": f"Unknown test_type '{test_type}'. Use one of: t_test_independent, t_test_one_sample, chi_square."}
+
+        style = (
+            '<style>'
+            '.vda-table-wrap{overflow-x:auto;margin:8px 0;border-radius:8px;border:1px solid #e5e7eb;}'
+            '.vda-table{width:100%;border-collapse:collapse;font-size:13px;font-family:Inter,system-ui,sans-serif;}'
+            '.vda-table thead th{background:#3B82F6;color:#fff;padding:9px 14px;text-align:left;white-space:nowrap;font-weight:600;}'
+            '.vda-table tbody td{padding:7px 14px;border-bottom:1px solid #f1f5f9;white-space:nowrap;}'
+            '.vda-table tbody tr:nth-child(even){background:#f8fafc;}'
+            '.vda-table tbody tr:last-child td{border-bottom:none;}'
+            '</style>'
+        )
+        header_cells = f'<th style="background:#1e40af;" colspan="2">{escape(title)}</th>'
+        row_html = [
+            '<tr>'
+            + f'<td style="font-weight:600;color:#1e40af;background:#eff6ff;white-space:nowrap;">{escape(label)}</td>'
+            + f'<td>{escape(value)}</td>'
+            + '</tr>'
+            for label, value in result_rows
+        ]
+        table = (
+            '<div class="vda-table-wrap"><table class="vda-table">'
+            f'<thead><tr>{header_cells}</tr></thead>'
+            '<tbody>' + '\n'.join(row_html) + '</tbody>'
+            '</table></div>'
+        )
+        return {"reply": style + table}
+
+    except Exception as e:
+        print("HYPOTHESIS TEST ERROR")
+        print(e)
+        return {"reply": f"There was an error running the hypothesis test. Error: {e}. You should probably try again."}
+
+
 def regression_func(independent_variables: List[str], dependent_variable: str, session_hash, session_folder, category: str='', **kwargs):
     print("LINEAR REGRESSION CALCULATION")
     print(independent_variables)

requirements.txt

@@ -14,3 +14,5 @@ pymongo_schema
 pandasql
 pluck-graphql
 certifi==2025.1.31
+scipy
+scikit-learn

tools/chart_tools.py

@@ -1,315 +1,308 @@
+# Shared parameter snippets reused across chart tool schemas.
+# Update here to change the description everywhere at once.
+
+_LAYOUT_PARAM = {
+    "type": "array",
+    "description": (
+        "Optional. An array containing a single JSON-formatted Plotly layout dictionary. "
+        "Use to set chart title, axis labels, colours, fonts, and other layout properties. "
+        "Example: [{\"title\": \"Monthly Sales\", \"xaxis\": {\"title\": \"Month\"}}]"
+    ),
+    "items": {"type": "string"},
+}
+
+_TRACE_STYLE_PARAM = {
+    "type": "array",
+    "description": (
+        "Optional. An array containing a single JSON-formatted Plotly trace styling dictionary. "
+        "Use to control visual properties such as line colour, opacity, and marker style. "
+        "Do NOT include 'x', 'y', or 'type' keys — those are set automatically from query.csv."
+    ),
+    "items": {"type": "string"},
+}
+
 chart_tool_schemas = [
     {
         "name": "scatter_chart_generation_func",
-        "description": """This is a scatter plot generation tool useful to generate scatter plots from queried data from our data source that we are querying.
-        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the scatter_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
-        to it for context if desired.""",
+        "description": (
+            "Generates a Plotly scatter plot from query.csv data. "
+            "Use when the user wants to visualise the relationship between two numeric columns, "
+            "create a bubble chart (via the size parameter), or overlay a trendline. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
-                "data": {
-                    "type": "array",
-                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
-                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
-                },
                 "x_column": {
                     "type": "array",
-                    "description": """An array of strings that correspond to the column names in our query.csv file that contain the x values of the graph. There can be more than one column
-                    that can each be plotted against the y_column, if needed.""",
-                    "items": {"type": "string"}
+                    "description": (
+                        "One or more column names from query.csv to plot on the x-axis. "
+                        "Multiple columns produce multiple series, each plotted against y_column."
+                    ),
+                    "items": {"type": "string"},
                 },
                 "y_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contain the y values of the graph."""
+                    "description": "Column name from query.csv to plot on the y-axis.",
                 },
                 "category": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data."""
+                    "description": "Optional column name used to colour-code points by a categorical grouping.",
                 },
                 "trendline": {
                     "type": "string",
-                    "description": """An optional field to specify the type of plotly trendline we wish to use in the scatter plot.
-                        This trendline value can be one of ['ols','lowess','rolling','ewm','expanding'].
-                        Do not send any values outside of this array as the function will fail.
-                        Infer this from the user's message."""
+                    "description": (
+                        "Optional trendline type. One of: 'ols' (linear regression), "
+                        "'lowess' (local smoothing), 'rolling', 'ewm', 'expanding'. "
+                        "Requires trendline_options when using 'lowess', 'rolling', or 'ewm'."
+                    ),
                 },
                 "trendline_options": {
                     "type": "array",
-                    "description": """An array containing a dictionary that contains the 'trendline_options' portion of the plotly chart generation.
-                    The 'lowess', 'rolling', and 'ewm' options require trendline_options to be included.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
+                    "description": (
+                        "Required when trendline is 'lowess', 'rolling', or 'ewm'. "
+                        "An array containing a single JSON-formatted dict of trendline options "
+                        "(e.g. [{\"window\": 7}] for a 7-point rolling average)."
+                    ),
+                    "items": {"type": "string"},
                 },
                 "marginal_x": {
                     "type": "string",
-                    "description": """The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the x axis.
-                        This marginal_x value can be one of ['histogram','rug','box','violin'].
-                        Do not send any values outside of this array as the function will fail.
-                        Infer this from the user's message."""
+                    "description": "Optional marginal distribution plot along the x-axis. One of: 'histogram', 'rug', 'box', 'violin'.",
                 },
                 "marginal_y": {
                     "type": "string",
-                    "description": """The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the y axis.
-                        This marginal_y value can be one of ['histogram','rug','box','violin'].
-                        Do not send any values outside of this array as the function will fail.
-                        Infer this from the user's message."""
-                },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
+                    "description": "Optional marginal distribution plot along the y-axis. One of: 'histogram', 'rug', 'box', 'violin'.",
                 },
                 "size": {
                     "type": "string",
-                    "description": """An optional column in our query.csv file that contain a parameter that will define the size of each plot point.
-                    This is useful for a bubble chart where another value in our query can be represented by the size of the plotted point.
-                    Values must be greater than or equal to 0 and so in our query, all values less than 0 should be set equal to zero."""
-                }
+                    "description": "Optional column name whose values control the size of each point (bubble chart). Negative values are clamped to zero.",
+                },
+                "data": _TRACE_STYLE_PARAM,
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["x_column", "y_column"]
+            "required": ["x_column", "y_column"],
         },
     },
     {
         "name": "line_chart_generation_func",
-        "description": """This is a line chart generation tool useful to generate line charts from queried data from our data source that we are querying.
-        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the line_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
-        to it for context if desired.""",
+        "description": (
+            "Generates a Plotly line chart from query.csv data. "
+            "Use for trends over time or any ordered sequence. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
-                "data": {
-                    "type": "array",
-                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
-                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
-                },
                 "x_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contain the x values of the graph."""
+                    "description": "Column name from query.csv for the x-axis (typically a date or ordered index).",
                 },
                 "y_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contain the y values of the graph."""
+                    "description": "Column name from query.csv for the y-axis (numeric values).",
                 },
                 "category": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data."""
+                    "description": "Optional column name used to split the data into multiple colour-coded lines.",
                 },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
-                }
+                "data": _TRACE_STYLE_PARAM,
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["x_column", "y_column", "layout"]
+            "required": ["x_column", "y_column"],
         },
     },
     {
         "name": "bar_chart_generation_func",
-        "description": """This is a bar chart generation tool useful to generate bar charts from queried data from our data source that we are querying.
-        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the bar_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
-        to it for context if desired.""",
+        "description": (
+            "Generates a Plotly bar chart from query.csv data. "
+            "Use for comparing values across categories. Supports grouped/stacked bars via category, "
+            "and faceted subplots via facet_row or facet_col. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
-                "data": {
-                    "type": "array",
-                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
-                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
-                },
                 "x_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contains the x values of the graph."""
+                    "description": "Column name from query.csv for the x-axis (category labels).",
                 },
                 "y_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contains the y values of the graph."""
+                    "description": "Column name from query.csv for the y-axis (numeric values).",
                 },
                 "category": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contains a parameter that will define the category for the data."""
+                    "description": "Optional column name used to colour-code bars into grouped or stacked series.",
                 },
                 "facet_row": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contains a parameter that will define a faceted subplot, where different rows
-                    correspond to different values of the query specified in this parameter."""
+                    "description": "Optional column name. Creates one subplot row per unique value — useful for comparing distributions across a second dimension.",
                 },
                 "facet_col": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contain a parameter that will define the faceted column, corresponding to
-                    different values of our query specified in this parameter."""
+                    "description": "Optional column name. Creates one subplot column per unique value.",
                 },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
-                }
+                "data": _TRACE_STYLE_PARAM,
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["x_column", "y_column", "layout"]
+            "required": ["x_column", "y_column"],
         },
     },
     {
         "name": "pie_chart_generation_func",
-        "description": """This is a pie chart generation tool useful to generate pie charts from queried data from our data source that we are querying.
-        The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the pie_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
-        to it for context if desired.""",
+        "description": (
+            "Generates a Plotly pie chart from query.csv data. "
+            "Use when the user wants to show part-to-whole proportions. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
-                "data": {
-                    "type": "array",
-                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
-                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
-                },
                 "values": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contain the values of the pie chart."""
+                    "description": "Column name from query.csv containing the numeric value for each slice.",
                 },
                 "names": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contain the label or section of each piece of the pie graph and allow us to know what each piece of the pie chart represents."""
+                    "description": "Column name from query.csv containing the label for each slice.",
                 },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
-                }
+                "data": _TRACE_STYLE_PARAM,
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["values", "names", "layout"]
+            "required": ["values", "names"],
         },
     },
     {
         "name": "histogram_generation_func",
-        "description": """This is a histogram generation tool useful to generate histograms from queried data from our data source that we are querying.
-        The data values will come from the columns of our query.csv file but the layout section of the plotly dictionary objects will be generated by you.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the histogram_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
-        to it for context if desired.""",
+        "description": (
+            "Generates a Plotly histogram from query.csv data. "
+            "Use to show the frequency distribution of a numeric column. "
+            "Supports normalisation (percent, probability, density) and aggregation functions per bin. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
-                "data": {
-                    "type": "array",
-                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
-                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
-                },
                 "x_column": {
                     "type": "string",
-                    "description": """The column name in our query.csv file that contains the x values of the histogram.
-                    This would correspond to the counts that would be distributed in the histogram."""
+                    "description": "Column name from query.csv whose values are binned on the x-axis.",
                 },
                 "y_column": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contains the y values of the histogram."""
+                    "description": "Optional column name aggregated per bin via histfunc (e.g. sum of sales per price bucket).",
                 },
                 "histnorm": {
                     "type": "string",
-                    "description": """An optional argument to specify the type of normalization if the default isn't used.
-                    This histnorm value can be one of ['percent','probability','density','probability density'].
-                    Do not send any values outside of this array as the function will fail."""
+                    "description": "Optional normalisation. One of: 'percent', 'probability', 'density', 'probability density'.",
                 },
                 "category": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contains a parameter that will define the category for the data."""
+                    "description": "Optional column name used to overlay multiple colour-coded histograms.",
                 },
                 "histfunc": {
                     "type": "string",
-                    "description": """An optional value that represents the function of data to compute the function which is used on the optional y column.
-                    This histfunc value can be one of ['avg','sum','count'].
-                    Do not send any values outside of this array as the function will fail."""
+                    "description": "Optional aggregation function applied to y_column per bin. One of: 'avg', 'sum', 'count'.",
                 },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
-                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                    "items": {"type": "string"}
-                }
+                "data": _TRACE_STYLE_PARAM,
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["x_column"]
+            "required": ["x_column"],
         },
     },
     {
         "name": "box_chart_generation_func",
-        "description": """This is a box plot generation tool useful for visualising the distribution of a numeric column and identifying outliers.
-        Particularly valuable when comparing distributions across categories (e.g., revenue by country, price by product type).
-        The data values come from the columns of our query.csv file.
-        Returns an iframe string to display inline in the chat window.""",
+        "description": (
+            "Generates a Plotly box plot from query.csv data. "
+            "Use to visualise the distribution of a numeric column and identify outliers. "
+            "Especially useful for comparing distributions across categories. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
                 "y_column": {
                     "type": "string",
-                    "description": "The column name containing the numeric values whose distribution should be shown on the y-axis."
+                    "description": "Column name from query.csv containing the numeric values to distribute on the y-axis.",
                 },
                 "x_column": {
                     "type": "string",
-                    "description": "An optional column name to group the boxes by on the x-axis (e.g., a category or label column). When provided, one box is drawn per unique value."
+                    "description": "Optional column name. Groups data into one box per unique value on the x-axis.",
                 },
                 "category": {
                     "type": "string",
-                    "description": "An optional column name used to colour-code the boxes by a secondary grouping."
+                    "description": "Optional column name used to colour-code boxes by a secondary grouping.",
                 },
-                "layout": {
-                    "type": "array",
-                    "description": """An array containing a dictionary with the 'layout' portion of the plotly chart.
-                    The array must contain a json formatted dictionary with outer brackets included.""",
-                    "items": {"type": "string"}
-                }
+                "layout": _LAYOUT_PARAM,
             },
-            "required": ["y_column"]
+            "required": ["y_column"],
         },
     },
     {
         "name": "correlation_heatmap_func",
-        "description": """This is a correlation heatmap tool that computes the pairwise Pearson correlation between numeric columns
-        in our query.csv file and displays the result as a colour-coded heatmap (blue = positive, red = negative correlation).
-        Use this when the user asks which variables are related, correlated, or associated with each other.
-        Returns an iframe string to display inline in the chat window.""",
+        "description": (
+            "Computes pairwise Pearson correlations between numeric columns in query.csv and renders "
+            "the result as a colour-coded heatmap (blue = positive, red = negative). "
+            "Use when the user asks which variables are related, correlated, or associated with each other. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
                 "columns": {
                     "type": "array",
-                    "description": "An optional list of numeric column names to include in the correlation matrix. If omitted, all numeric columns from the query results are used. Avoid including ID or index columns.",
-                    "items": {"type": "string"}
-                }
+                    "description": "Optional list of numeric column names to include in the matrix. If omitted, all numeric columns from query.csv are used. Avoid ID or index columns.",
+                    "items": {"type": "string"},
+                },
+            },
+            "required": [],
+        },
+    },
+    {
+        "name": "rolling_stats_func",
+        "description": (
+            "Generates a rolling statistics / moving average chart from query.csv data. "
+            "Overlays rolling aggregations (mean, std, min, max) on top of the original series. "
+            "Use when the user asks for a moving average, rolling average, rolling statistics, or wants to smooth a time series. "
+            "Returns an HTML iframe — display it verbatim in the chat."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "x_column": {
+                    "type": "string",
+                    "description": "Column name from query.csv for the x-axis — typically a date or sequential index.",
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": "Column name from query.csv containing the numeric values to compute rolling stats on.",
+                },
+                "window": {
+                    "type": "integer",
+                    "description": "Rolling window size in number of rows. Default 7. Infer from the user's request.",
+                },
+                "stats": {
+                    "type": "array",
+                    "description": "Statistics to overlay. Valid values: 'mean', 'std', 'min', 'max'. Defaults to ['mean'] if omitted.",
+                    "items": {"type": "string"},
+                },
+                "category": {
+                    "type": "string",
+                    "description": "Optional column name to group the data, producing separate rolling stat lines per group.",
+                },
+                "layout": _LAYOUT_PARAM,
             },
-            "required": []
+            "required": ["x_column", "y_column"],
         },
     },
     {
         "name": "table_generation_func",
-        "description": """This is a table generation tool useful to format data as a table from queried data from our data source that we are querying.
-        Takes no parameters as it uses data queried in our query.csv file to build the table.
-        Call this function after running our query and generating query.csv.
-        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
-        from the table_generation_func function in any way and always display the iframe fully to the user in the chat window.""",
+        "description": (
+            "Formats query.csv results as a styled HTML table. "
+            "Use when the user wants to view raw query results in a readable format, "
+            "or when result data is too large to describe in text. Displays up to 200 rows. "
+            "Returns an HTML table — display it verbatim in the chat."
+        ),
         "parameters": {"type": "object", "properties": {}},
     },
 ]

tools/stats_tools.py

@@ -1,50 +1,130 @@
 stats_tool_schemas = [
     {
         "name": "descriptive_stats_func",
-        "description": """This tool computes summary statistics for numeric columns in our query.csv file and returns them as a formatted table.
-        Statistics include count, mean, standard deviation, min, 25th percentile, median, 75th percentile, and max.
-        Use this when the user asks for summary statistics, descriptive statistics, or a statistical overview of the data.""",
+        "description": (
+            "Computes summary statistics for numeric columns in query.csv: "
+            "count, mean, std, min, 25th/50th/75th percentile, and max. "
+            "Use when the user asks for summary statistics, descriptive statistics, or a statistical overview. "
+            "Returns a formatted HTML table."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
                 "columns": {
                     "type": "array",
-                    "description": "An optional list of column names to compute statistics for. If omitted, statistics are computed for all numeric columns in the query results. Avoid including ID or index columns.",
-                    "items": {"type": "string"}
-                }
+                    "description": "Optional list of column names to include. If omitted, all numeric columns from query.csv are used. Avoid ID or index columns.",
+                    "items": {"type": "string"},
+                },
             },
-            "required": []
+            "required": [],
+        },
+    },
+    {
+        "name": "kmeans_clustering_func",
+        "description": (
+            "Runs K-Means clustering on numeric feature columns from query.csv. "
+            "Groups rows into k clusters, displays a scatter plot coloured by cluster assignment, "
+            "and returns a centroid summary table showing the mean of each feature per cluster. "
+            "Use when the user asks to cluster the data, find natural segments or groups, or apply K-Means. "
+            "Returns an HTML iframe and summary table."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "feature_columns": {
+                    "type": "array",
+                    "description": "List of numeric column names from query.csv to use as clustering features.",
+                    "items": {"type": "string"},
+                },
+                "x_column": {
+                    "type": "string",
+                    "description": "Column name from query.csv for the x-axis of the scatter plot. Usually one of the feature columns.",
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": "Column name from query.csv for the y-axis of the scatter plot. Usually one of the feature columns.",
+                },
+                "n_clusters": {
+                    "type": "integer",
+                    "description": "Number of clusters (k). Default 3. Infer from the user's request.",
+                },
+                "layout": {
+                    "type": "array",
+                    "description": "Optional. An array containing a single JSON-formatted Plotly layout dictionary.",
+                    "items": {"type": "string"},
+                },
+            },
+            "required": ["feature_columns", "x_column", "y_column"],
+        },
+    },
+    {
+        "name": "hypothesis_test_func",
+        "description": (
+            "Performs a statistical hypothesis test on query.csv data and returns a formatted results table "
+            "with test statistic, p-value, and significance at α=0.05. "
+            "Supported tests:\n"
+            "- 't_test_independent': compare means of a numeric column across two groups "
+            "(requires group_column; use group_values if the column has more than 2 unique values).\n"
+            "- 't_test_one_sample': test whether a column's mean equals a hypothesized value (requires pop_mean).\n"
+            "- 'chi_square': test independence between two categorical columns (requires column and column2)."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "test_type": {
+                    "type": "string",
+                    "description": "Test to run. One of: 't_test_independent', 't_test_one_sample', 'chi_square'.",
+                },
+                "column": {
+                    "type": "string",
+                    "description": "Primary column for the test. Numeric for t-tests; first categorical column for chi-square.",
+                },
+                "column2": {
+                    "type": "string",
+                    "description": "Second categorical column. Required for 'chi_square'.",
+                },
+                "group_column": {
+                    "type": "string",
+                    "description": "Grouping column. Required for 't_test_independent'. Must have exactly 2 unique values, or specify group_values.",
+                },
+                "group_values": {
+                    "type": "array",
+                    "description": "Exactly 2 group labels to compare. Use when group_column has more than 2 unique values.",
+                    "items": {"type": "string"},
+                },
+                "pop_mean": {
+                    "type": "number",
+                    "description": "Hypothesized population mean (μ₀). Required for 't_test_one_sample'.",
+                },
+            },
+            "required": ["test_type", "column"],
         },
     },
     {
         "name": "regression_func",
-        "description": """This a tool to calculate regressions on our data source that we are querying.
-        We can run queries with our 'sql_query_func' function and they will be available to use in this function via the query.csv file that is generated.
-        Returns a dictionary of values that includes a regression_summary and a regression chart (which is an iframe displaying the
-        linear regression in chart form and should be shown to the user).""",
+        "description": (
+            "Runs an OLS linear regression on query.csv data. "
+            "Use when the user wants to model the relationship between variables, assess predictors, or run a regression. "
+            "Returns a regression summary (coefficients, R², p-values) and a scatter plot with the fitted line as an HTML iframe."
+        ),
         "parameters": {
             "type": "object",
             "properties": {
                 "independent_variables": {
                     "type": "array",
-                    "description": """An array of strings that states the independent variables in our data set which should be column names in our query.csv file that is generated
-                    in the 'sql_query_func' function. This will allow us to identify the data to use for our independent variables.
-                    Infer this from the user's message.""",
-                    "items": {"type": "string"}
+                    "description": "Column names from query.csv to use as independent (predictor) variables.",
+                    "items": {"type": "string"},
                 },
                 "dependent_variable": {
                     "type": "string",
-                    "description": """A string that states the dependent variable in our data set which should be a column name in our query.csv file that is generated
-                    in the 'sql_query_func' function. This will allow us to identify the data to use for our dependent variable.
-                    Infer this from the user's message."""
+                    "description": "Column name from query.csv to use as the dependent (outcome) variable.",
                 },
                 "category": {
                     "type": "string",
-                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data.
-                    Do not send value if no category is needed or specified. This category must be present in our query.csv file to be valid."""
-                }
+                    "description": "Optional column name used to colour-code points and fit separate regression lines per group.",
+                },
             },
-            "required": ["independent_variables", "dependent_variable"]
+            "required": ["independent_variables", "dependent_variable"],
         },
-    }
+    },
 ]