feat: Add Bigquery detect_anomalies tool

This change introduces a new `detect_anomalies` tool in `query_tool.py` which uses BigQuery ML's `CREATE MODEL` with `ARIMA_PLUS` type and `ML.DETECT_ANOMALIES` to detect anomalies. The new function is also added to the `bigquery_toolset`.

PiperOrigin-RevId: 825181489

contributing/samples/bigquery/README.md

@@ -46,6 +46,12 @@ distributed via the `google.adk.tools.bigquery` module. These tools include:
   `CONTRIBUTION_ANALYSIS` model and then querying it with
   `ML.GET_INSIGHTS` to find top contributors for a given metric.
 
+9. `detect_anomalies`
+
+  Perform time series anomaly detection in BigQuery by creating a temporary
+  `ARIMA_PLUS` model and then querying it with
+  `ML.DETECT_ANOMALIES` to detect time series data anomalies.
+
 ## How to use
 
 Set up environment variables in your `.env` file for using

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

@@ -83,6 +83,7 @@ class BigQueryToolset(BaseToolset):
             query_tool.get_execute_sql(self._tool_settings),
             query_tool.forecast,
             query_tool.analyze_contribution,
+            query_tool.detect_anomalies,
             data_insights_tool.ask_data_insights,
         ]
     ]

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

@@ -1092,3 +1092,214 @@ def analyze_contribution(
     settings.write_mode == original_write_mode
 
   return result
+
+
+def detect_anomalies(
+    project_id: str,
+    history_data: str,
+    times_series_timestamp_col: str,
+    times_series_data_col: str,
+    horizon: Optional[int] = 10,
+    times_series_id_cols: Optional[list[str]] = None,
+    anomaly_prob_threshold: Optional[float] = 0.95,
+    *,
+    credentials: Credentials,
+    settings: BigQueryToolConfig,
+    tool_context: ToolContext,
+) -> dict:
+  """Run a BigQuery time series ARIMA_PLUS model training and anomaly detection using CREATE MODEL and ML.DETECT_ANOMALIES clauses.
+
+  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.
+      times_series_timestamp_col (str): The name of the colum containing the
+        timestamp for each data point.
+      times_series_data_col (str): The name of the column containing the
+        numerical values to be forecasted and anomaly detected.
+      horizon (int, optional): The number of time steps to forecast into the
+        future. Defaults to 10.
+      times_series_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.
+      anomaly_prob_threshold (float, optional): The probability threshold to
+        determine if a data point is an anomaly. Defaults to 0.95.
+      credentials (Credentials): The credentials to use for the request.
+      settings (BigQueryToolConfig): The settings for the tool.
+      tool_context (ToolContext): The context for the tool.
+
+  Returns:
+      dict: Dictionary representing the result of the anomaly detection. The
+            result contains the boolean value if the data point is anomaly or
+            not, lower bound, upper bound and anomaly probability for each data
+            point and also the probability of whether the data point is anomaly
+            or not.
+
+  Examples:
+      Detect Anomalies daily sales based on historical data from a BigQuery
+      table:
+
+          >>> detect_anomalies(
+          ...     project_id="my-gcp-project",
+          ...     history_data="my-dataset.my-sales-table",
+          ...     times_series_timestamp_col="sale_date",
+          ...     times_series_data_col="daily_sales"
+          ... )
+          {
+            "status": "SUCCESS",
+            "rows": [
+              {
+                "ts_timestamp": "2021-01-01 00:00:01 UTC",
+                "ts_data": 125.3,
+                "is_anomaly": TRUE,
+                "lower_bound": 129.5,
+                "upper_bound": 133.6 ,
+                "anomaly_probability": 0.93
+              },
+              ...
+            ]
+          }
+
+      Detect Anomalies on 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'"
+          ... )
+          >>> detect_anomalies(
+          ...     project_id="my-gcp-project",
+          ...     history_data=history_query,
+          ...     times_series_timestamp_col="timestamp",
+          ...     times_series_data_col="value",
+          ...     times_series_id_cols=["unique_id"]
+          ... )
+          {
+            "status": "SUCCESS",
+            "rows": [
+              {
+                "unique_id": "T1",
+                "ts_timestamp": "2021-01-01 00:00:01 UTC",
+                "ts_data": 125.3,
+                "is_anomaly": TRUE,
+                "lower_bound": 129.5,
+                "upper_bound": 133.6 ,
+                "anomaly_probability": 0.93
+              },
+              ...
+            ]
+          }
+
+      Error Scenarios:
+          When an element in `times_series_id_cols` is not a string:
+
+          >>> detect_anomalies(
+          ...     project_id="my-gcp-project",
+          ...     history_data="my-dataset.my-sales-table",
+          ...     times_series_timestamp_col="sale_date",
+          ...     times_series_data_col="daily_sales",
+          ...     times_series_id_cols=["store_id", 123]
+          ... )
+          {
+            "status": "ERROR",
+            "error_details": "All elements in times_series_id_cols must be
+            strings."
+          }
+
+          When `history_data` refers to a table that does not exist:
+
+          >>> detect_anomalies(
+          ...     project_id="my-gcp-project",
+          ...     history_data="my-dataset.non-existent-table",
+          ...     times_series_timestamp_col="sale_date",
+          ...     times_series_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"
+          }
+  """
+  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"SELECT * FROM `{history_data}`"
+
+  options = [
+      "MODEL_TYPE = 'ARIMA_PLUS'",
+      f"TIME_SERIES_TIMESTAMP_COL = '{times_series_timestamp_col}'",
+      f"TIME_SERIES_DATA_COL = '{times_series_data_col}'",
+      f"HORIZON = {horizon}",
+  ]
+
+  if times_series_id_cols:
+    if not all(isinstance(item, str) for item in times_series_id_cols):
+      return {
+          "status": "ERROR",
+          "error_details": (
+              "All elements in times_series_id_cols must be strings."
+          ),
+      }
+    times_series_id_cols_str = (
+        "[" + ", ".join([f"'{col}'" for col in times_series_id_cols]) + "]"
+    )
+    options.append(f"TIME_SERIES_ID_COL = {times_series_id_cols_str}")
+
+  options_str = ", ".join(options)
+
+  model_name = f"detect_anomalies_model_{str(uuid.uuid4()).replace('-', '_')}"
+
+  create_model_query = f"""
+  CREATE TEMP MODEL {model_name}
+    OPTIONS ({options_str})
+  AS {history_data_source}
+  """
+
+  anomaly_detection_query = f"""
+  SELECT * FROM ML.DETECT_ANOMALIES(MODEL {model_name}, STRUCT({anomaly_prob_threshold} AS anomaly_prob_threshold))
+  """
+
+  # Create a session and run the create model query.
+  original_write_mode = settings.write_mode
+  try:
+    if settings.write_mode == WriteMode.BLOCKED:
+      raise ValueError("anomaly detection is not allowed in this session.")
+    elif original_write_mode != WriteMode.PROTECTED:
+      # Running create temp model requires a session. So we set the write mode
+      # to PROTECTED to run the create model query and job query in the same
+      # session.
+      settings.write_mode = WriteMode.PROTECTED
+
+    result = execute_sql(
+        project_id,
+        create_model_query,
+        credentials,
+        settings,
+        tool_context,
+    )
+    if result["status"] != "SUCCESS":
+      return result
+
+    result = execute_sql(
+        project_id,
+        anomaly_detection_query,
+        credentials,
+        settings,
+        tool_context,
+    )
+  except Exception as ex:  # pylint: disable=broad-except
+    return {
+        "status": "ERROR",
+        "error_details": f"Error during anomaly detection: {str(ex)}",
+    }
+  finally:
+    # Restore the original write mode.
+    settings.write_mode == original_write_mode
+
+  return result

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 analyze_contribution
+from google.adk.tools.bigquery.query_tool import detect_anomalies
 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
@@ -1401,3 +1402,132 @@ def test_analyze_contribution_with_invalid_dimension_id_cols():
       "All elements in dimension_id_cols must be strings."
       in result["error_details"]
   )
+
+
+# detect_anomalies calls execute_sql twice. We need to test that
+# the queries are properly constructed and call execute_sql with the correct
+# parameters exactly twice.
+@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("uuid.uuid4", autospec=True)
+def test_detect_anomalies_with_table_id(mock_uuid, mock_execute_sql):
+  """Test time series anomaly detection tool invocation with a table id."""
+  mock_credentials = mock.MagicMock(spec=Credentials)
+  mock_settings = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
+  mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+  mock_uuid.return_value = "test_uuid"
+  mock_execute_sql.return_value = {"status": "SUCCESS"}
+
+  history_data_query = "SELECT * FROM `test-dataset.test-table`"
+  detect_anomalies(
+      project_id="test-project",
+      history_data=history_data_query,
+      times_series_timestamp_col="ts_timestamp",
+      times_series_data_col="ts_data",
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+  )
+
+  expected_create_model_query = """
+  CREATE TEMP MODEL detect_anomalies_model_test_uuid
+    OPTIONS (MODEL_TYPE = 'ARIMA_PLUS', TIME_SERIES_TIMESTAMP_COL = 'ts_timestamp', TIME_SERIES_DATA_COL = 'ts_data', HORIZON = 10)
+  AS (SELECT * FROM `test-dataset.test-table`)
+  """
+
+  expected_anomaly_detection_query = """
+  SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.95 AS anomaly_prob_threshold))
+  """
+
+  assert mock_execute_sql.call_count == 2
+  mock_execute_sql.assert_any_call(
+      "test-project",
+      expected_create_model_query,
+      mock_credentials,
+      mock_settings,
+      mock_tool_context,
+  )
+  mock_execute_sql.assert_any_call(
+      "test-project",
+      expected_anomaly_detection_query,
+      mock_credentials,
+      mock_settings,
+      mock_tool_context,
+  )
+
+
+# detect_anomalies calls execute_sql twice. We need to test that
+# the queries are properly constructed and call execute_sql with the correct
+# parameters exactly twice.
+@mock.patch("google.adk.tools.bigquery.query_tool.execute_sql", autospec=True)
+@mock.patch("uuid.uuid4", autospec=True)
+def test_detect_anomalies_with_custom_params(mock_uuid, mock_execute_sql):
+  """Test time series anomaly detection tool invocation with a table id."""
+  mock_credentials = mock.MagicMock(spec=Credentials)
+  mock_settings = BigQueryToolConfig(write_mode=WriteMode.PROTECTED)
+  mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+  mock_uuid.return_value = "test_uuid"
+  mock_execute_sql.return_value = {"status": "SUCCESS"}
+
+  history_data_query = "SELECT * FROM `test-dataset.test-table`"
+  detect_anomalies(
+      project_id="test-project",
+      history_data=history_data_query,
+      times_series_timestamp_col="ts_timestamp",
+      times_series_data_col="ts_data",
+      times_series_id_cols=["dim1", "dim2"],
+      horizon=20,
+      anomaly_prob_threshold=0.8,
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+  )
+
+  expected_create_model_query = """
+  CREATE TEMP MODEL detect_anomalies_model_test_uuid
+    OPTIONS (MODEL_TYPE = 'ARIMA_PLUS', TIME_SERIES_TIMESTAMP_COL = 'ts_timestamp', TIME_SERIES_DATA_COL = 'ts_data', HORIZON = 20, TIME_SERIES_ID_COL = ['dim1', 'dim2'])
+  AS (SELECT * FROM `test-dataset.test-table`)
+  """
+
+  expected_anomaly_detection_query = """
+  SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.8 AS anomaly_prob_threshold))
+  """
+
+  assert mock_execute_sql.call_count == 2
+  mock_execute_sql.assert_any_call(
+      "test-project",
+      expected_create_model_query,
+      mock_credentials,
+      mock_settings,
+      mock_tool_context,
+  )
+  mock_execute_sql.assert_any_call(
+      "test-project",
+      expected_anomaly_detection_query,
+      mock_credentials,
+      mock_settings,
+      mock_tool_context,
+  )
+
+
+def test_detect_anomalies__with_invalid_id_cols():
+  """Test time series anomaly detection tool invocation with invalid times_series_id_cols."""
+  mock_credentials = mock.MagicMock(spec=Credentials)
+  mock_settings = BigQueryToolConfig()
+  mock_tool_context = mock.create_autospec(ToolContext, instance=True)
+
+  result = detect_anomalies(
+      project_id="test-project",
+      history_data="test-dataset.test-table",
+      times_series_timestamp_col="ts_timestamp",
+      times_series_data_col="ts_data",
+      times_series_id_cols=["dim1", 123],
+      credentials=mock_credentials,
+      settings=mock_settings,
+      tool_context=mock_tool_context,
+  )
+
+  assert result["status"] == "ERROR"
+  assert (
+      "All elements in times_series_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) == 8
+  assert len(tools) == 9
   assert all([isinstance(tool, GoogleTool) for tool in tools])
 
   expected_tool_names = set([
@@ -53,6 +53,7 @@ async def test_bigquery_toolset_tools_default():
       "ask_data_insights",
       "forecast",
       "analyze_contribution",
+      "detect_anomalies",
   ])
   actual_tool_names = set([tool.name for tool in tools])
   assert actual_tool_names == expected_tool_names