feat: Add Bigquery Forecast tool

This tool answers questions about structured data in BigQuery using natural language. PiperOrigin-RevId: 805414952

Author: google-genai-bot

contributing/samples/bigquery/README.md

@@ -35,6 +35,11 @@ distributed via the `google.adk.tools.bigquery` module. These tools include:
  the official [Conversational Analytics API documentation](https://cloud.google.com/gemini/docs/conversational-analytics-api/overview)
  for instructions.
 
+1. `forecast`
+
+ Perform time series forecasting using BigQuery's `AI.FORECAST` function,
+ leveraging the TimesFM 2.0 model.
+
 ## How to use
 
 Set up environment variables in your `.env` file for using

src/google/adk/tools/bigquery/bigquery_toolset.py

@@ -81,6 +81,7 @@ async def get_tools(
  metadata_tool.list_dataset_ids,
  metadata_tool.list_table_ids,
  query_tool.get_execute_sql(self._tool_settings),
+ query_tool.forecast,
  data_insights_tool.ask_data_insights,
  ]
  ]

src/google/adk/tools/bigquery/query_tool.py

@@ -18,6 +18,7 @@
 import json
 import types
 from typing import Callable
+from typing import Optional
 
 from google.auth.credentials import Credentials
 from google.cloud import bigquery
@@ -596,3 +597,169 @@ def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
  execute_sql_wrapper.__doc__ = _execute_sql_write_mode.__doc__
 
  return execute_sql_wrapper
+
+
+def forecast(
+ project_id: str,
+ history_data: str,
+ timestamp_col: str,
+ data_col: str,
+ credentials: Credentials,
+ settings: BigQueryToolConfig,
+ tool_context: ToolContext,
+ horizon: int = 10,
+ id_cols: Optional[list[str]] = None,
+) -> dict:
+ """Run a BigQuery AI time series forecast using AI.FORECAST.
+
+ Args:
+ project_id (str): The GCP project id in which the query should be
+ executed.
+ history_data (str): The table id of the BigQuery table containing the
+ history time series data or a query statement that select the history
+ data.
+ timestamp_col (str): The name of the column containing the timestamp for
+ each data point.
+ data_col (str): The name of the column containing the numerical values to
+ be forecasted.
+ credentials (Credentials): The credentials to use for the request.
+ settings (BigQueryToolConfig): The settings for the tool.
+ tool_context (ToolContext): The context for the tool.
+ horizon (int, optional): The number of time steps to forecast into the
+ future. Defaults to 10.
+ id_cols (list, optional): The column names of the id columns to indicate
+ each time series when there are multiple time series in the table. All
+ elements must be strings. Defaults to None.
+
+ Returns:
+ dict: Dictionary representing the result of the forecast. The result
+ contains the forecasted values along with prediction intervals.
+
+ Examples:
+ Forecast daily sales for the next 7 days based on historical data from
+ a BigQuery table:
+
+ >>> forecast(
+ ... project_id="my-gcp-project",
+ ... history_data="my-dataset.my-sales-table",
+ ... timestamp_col="sale_date",
+ ... data_col="daily_sales",
+ ... horizon=7
+ ... )
+ {
+ "status": "SUCCESS",
+ "rows": [
+ {
+ "forecast_timestamp": "2025-01-08T00:00:00",
+ "forecast_value": 12345.67,
+ "confidence_level": 0.95,
+ "prediction_interval_lower_bound": 11000.0,
+ "prediction_interval_upper_bound": 13691.34,
+ "ai_forecast_status": ""
+ },
+ ...
+ ]
+ }
+
+ Forecast multiple time series using a SQL query as input:
+
+ >>> history_query = (
+ ... "SELECT unique_id, timestamp, value "
+ ... "FROM `my-project.my-dataset.my-timeseries-table` "
+ ... "WHERE timestamp > '1980-01-01'"
+ ... )
+ >>> forecast(
+ ... project_id="my-gcp-project",
+ ... history_data=history_query,
+ ... timestamp_col="timestamp",
+ ... data_col="value",
+ ... id_cols=["unique_id"],
+ ... horizon=14
+ ... )
+ {
+ "status": "SUCCESS",
+ "rows": [
+ {
+ "unique_id": "T1",
+ "forecast_timestamp": "1980-08-28T00:00:00",
+ "forecast_value": 1253218.75,
+ "confidence_level": 0.95,
+ "prediction_interval_lower_bound": 274252.51,
+ "prediction_interval_upper_bound": 2232184.99,
+ "ai_forecast_status": ""
+ },
+ ...
+ ]
+ }
+
+ Error Scenarios:
+ When an element in `id_cols` is not a string:
+
+ >>> forecast(
+ ... project_id="my-gcp-project",
+ ... history_data="my-dataset.my-sales-table",
+ ... timestamp_col="sale_date",
+ ... data_col="daily_sales",
+ ... id_cols=["store_id", 123]
+ ... )
+ {
+ "status": "ERROR",
+ "error_details": "All elements in id_cols must be strings."
+ }
+
+ When `history_data` refers to a table that does not exist:
+
+ >>> forecast(
+ ... project_id="my-gcp-project",
+ ... history_data="my-dataset.non-existent-table",
+ ... timestamp_col="sale_date",
+ ... data_col="daily_sales"
+ ... )
+ {
+ "status": "ERROR",
+ "error_details": "Not found: Table
+ my-gcp-project:my-dataset.non-existent-table was not found in
+ location US"
+ }
+ """
+ model = "TimesFM 2.0"
+ confidence_level = 0.95
+ trimmed_upper_history_data = history_data.strip().upper()
+ if trimmed_upper_history_data.startswith(
+ "SELECT"
+ ) or trimmed_upper_history_data.startswith("WITH"):
+ history_data_source = f"({history_data})"
+ else:
+ history_data_source = f"TABLE `{history_data}`"
+
+ if id_cols:
+ if not all(isinstance(item, str) for item in id_cols):
+ return {
+ "status": "ERROR",
+ "error_details": "All elements in id_cols must be strings.",
+ }
+ id_cols_str = "[" + ", ".join([f"'{col}'" for col in id_cols]) + "]"
+
+ query = f"""
+ SELECT * FROM AI.FORECAST(
+ {history_data_source},
+ data_col => '{data_col}',
+ timestamp_col => '{timestamp_col}',
+ model => '{model}',
+ id_cols => {id_cols_str},
+ horizon => {horizon},
+ confidence_level => {confidence_level}
+ )
+ """
+ else:
+ query = f"""
+ SELECT * FROM AI.FORECAST(
+ {history_data_source},
+ data_col => '{data_col}',
+ timestamp_col => '{timestamp_col}',
+ model => '{model}',
+ horizon => {horizon},
+ confidence_level => {confidence_level}
+ )
+ """
+ return execute_sql(project_id, query, credentials, settings, tool_context)

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

@@ -29,6 +29,7 @@
 from google.adk.tools.bigquery.config import BigQueryToolConfig
 from google.adk.tools.bigquery.config import WriteMode
 from google.adk.tools.bigquery.query_tool import execute_sql
+from google.adk.tools.bigquery.query_tool import forecast
 from google.adk.tools.tool_context import ToolContext
 from google.auth.exceptions import DefaultCredentialsError
 from google.cloud import bigquery
@@ -1028,3 +1029,103 @@ def test_execute_sql_unexpected_project_id():
  f" {compute_project_id}."
  ),
  }
+
+
+# AI.Forecast calls execute_sql with a specific query statement. We need to
+# test that the query is properly constructed and call execute_sql with the
+# correct parameters exactly once.
+@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+def test_forecast_with_table_id(mock_execute_sql):
+ mock_credentials = mock.MagicMock(spec=Credentials)
+ mock_settings = BigQueryToolConfig()
+ mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+
+ forecast(
+ project_id="test-project",
+ history_data="test-dataset.test-table",
+ timestamp_col="ts_col",
+ data_col="data_col",
+ credentials=mock_credentials,
+ settings=mock_settings,
+ tool_context=mock_tool_context,
+ horizon=20,
+ id_cols=["id1", "id2"],
+ )
+
+ expected_query = """
+ SELECT * FROM AI.FORECAST(
+ TABLE `test-dataset.test-table`,
+ data_col => 'data_col',
+ timestamp_col => 'ts_col',
+ model => 'TimesFM 2.0',
+ id_cols => ['id1', 'id2'],
+ horizon => 20,
+ confidence_level => 0.95
+ )
+ """
+ mock_execute_sql.assert_called_once_with(
+ "test-project",
+ expected_query,
+ mock_credentials,
+ mock_settings,
+ mock_tool_context,
+ )
+
+
+# AI.Forecast calls execute_sql with a specific query statement. We need to
+# test that the query is properly constructed and call execute_sql with the
+# correct parameters exactly once.
+@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+def test_forecast_with_query_statement(mock_execute_sql):
+ mock_credentials = mock.MagicMock(spec=Credentials)
+ mock_settings = BigQueryToolConfig()
+ mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+
+ history_data_query = "SELECT * FROM `test-dataset.test-table`"
+ forecast(
+ project_id="test-project",
+ history_data=history_data_query,
+ timestamp_col="ts_col",
+ data_col="data_col",
+ credentials=mock_credentials,
+ settings=mock_settings,
+ tool_context=mock_tool_context,
+ )
+
+ expected_query = f"""
+ SELECT * FROM AI.FORECAST(
+ ({history_data_query}),
+ data_col => 'data_col',
+ timestamp_col => 'ts_col',
+ model => 'TimesFM 2.0',
+ horizon => 10,
+ confidence_level => 0.95
+ )
+ """
+ mock_execute_sql.assert_called_once_with(
+ "test-project",
+ expected_query,
+ mock_credentials,
+ mock_settings,
+ mock_tool_context,
+ )
+
+
+def test_forecast_with_invalid_id_cols():
+ mock_credentials = mock.MagicMock(spec=Credentials)
+ mock_settings = BigQueryToolConfig()
+ mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+
+ result = forecast(
+ project_id="test-project",
+ history_data="test-dataset.test-table",
+ timestamp_col="ts_col",
+ data_col="data_col",
+ credentials=mock_credentials,
+ settings=mock_settings,
+ tool_context=mock_tool_context,
+ id_cols=["id1", 123],
+ )
+
+ assert result["status"] == "ERROR"
+ assert "All elements in id_cols must be strings." in result["error_details"]

tests/unittests/tools/bigquery/test_bigquery_toolset.py

@@ -41,7 +41,7 @@ async def test_bigquery_toolset_tools_default():
  tools = await toolset.get_tools()
  assert tools is not None
 
- assert len(tools) == 6
+ assert len(tools) == 7
  assert all([isinstance(tool, GoogleTool) for tool in tools])
 
  expected_tool_names = set([
@@ -51,6 +51,7 @@ async def test_bigquery_toolset_tools_default():
  "get_table_info",
  "execute_sql",
  "ask_data_insights",
+ "forecast",
  ])
  actual_tool_names = set([tool.name for tool in tools])
  assert actual_tool_names == expected_tool_names