diff --git a/contributing/samples/bigquery/README.md b/contributing/samples/bigquery/README.md index c1d2b161..ea6c70a2 100644 --- a/contributing/samples/bigquery/README.md +++ b/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 diff --git a/src/google/adk/tools/bigquery/bigquery_toolset.py b/src/google/adk/tools/bigquery/bigquery_toolset.py index f4fb35f7..627c0bbc 100644 --- a/src/google/adk/tools/bigquery/bigquery_toolset.py +++ b/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, ] ] diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index a068d93c..f4a82871 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/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) diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index 21a14f11..48b6e392 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/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"] diff --git a/tests/unittests/tools/bigquery/test_bigquery_toolset.py b/tests/unittests/tools/bigquery/test_bigquery_toolset.py index c2cc0eef..74e913f5 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_toolset.py +++ b/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