feat: Add Bigquery Forecast tool

This tool answers questions about structured data in BigQuery using natural language.

PiperOrigin-RevId: 805414952

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 @@ class BigQueryToolset(BaseToolset):
             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 functools
 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 import BigQueryToolset
 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