feat: Extend Bigquery detect_anomalies tool to support future data anomaly detection

ARIMA supports both historical data and future data anomaly detection. This CL add how the tool support future table anomaly detection.

PiperOrigin-RevId: 827803748
This commit is contained in:
Google Team Member
2025-11-03 23:02:02 -08:00
committed by Copybara-Service
parent d2888a3766
commit 38ea749c9c
2 changed files with 72 additions and 0 deletions
@@ -1100,6 +1100,7 @@ def detect_anomalies(
times_series_timestamp_col: str, times_series_timestamp_col: str,
times_series_data_col: str, times_series_data_col: str,
horizon: Optional[int] = 10, horizon: Optional[int] = 10,
target_data: Optional[str] = None,
times_series_id_cols: Optional[list[str]] = None, times_series_id_cols: Optional[list[str]] = None,
anomaly_prob_threshold: Optional[float] = 0.95, anomaly_prob_threshold: Optional[float] = 0.95,
*, *,
@@ -1121,6 +1122,9 @@ def detect_anomalies(
numerical values to be forecasted and anomaly detected. numerical values to be forecasted and anomaly detected.
horizon (int, optional): The number of time steps to forecast into the horizon (int, optional): The number of time steps to forecast into the
future. Defaults to 10. future. Defaults to 10.
target_data (str, optional): The table id of the BigQuery table containing
the target time series data or a query statement that select the target
data.
times_series_id_cols (list, optional): The column names of the id columns 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 to indicate each time series when there are multiple time series in the
table. All elements must be strings. Defaults to None. table. All elements must be strings. Defaults to None.
@@ -1264,6 +1268,18 @@ def detect_anomalies(
anomaly_detection_query = f""" anomaly_detection_query = f"""
SELECT * FROM ML.DETECT_ANOMALIES(MODEL {model_name}, STRUCT({anomaly_prob_threshold} AS anomaly_prob_threshold)) SELECT * FROM ML.DETECT_ANOMALIES(MODEL {model_name}, STRUCT({anomaly_prob_threshold} AS anomaly_prob_threshold))
""" """
if target_data:
trimmed_upper_target_data = target_data.strip().upper()
if trimmed_upper_target_data.startswith(
"SELECT"
) or trimmed_upper_target_data.startswith("WITH"):
target_data_source = f"({target_data})"
else:
target_data_source = f"SELECT * FROM `{target_data}`"
anomaly_detection_query = f"""
SELECT * FROM ML.DETECT_ANOMALIES(MODEL {model_name}, STRUCT({anomaly_prob_threshold} AS anomaly_prob_threshold), {target_data_source})
"""
# Create a session and run the create model query. # Create a session and run the create model query.
original_write_mode = settings.write_mode original_write_mode = settings.write_mode
@@ -1509,6 +1509,62 @@ def test_detect_anomalies_with_custom_params(mock_uuid, mock_execute_sql):
) )
# 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_on_target_table(mock_uuid, mock_execute_sql):
"""Test time series anomaly detection tool with target data is provided."""
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.history-table`"
target_data_query = "SELECT * FROM `test-dataset.target-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,
target_data=target_data_query,
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.history-table`)
"""
expected_anomaly_detection_query = """
SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.8 AS anomaly_prob_threshold), (SELECT * FROM `test-dataset.target-table`))
"""
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(): def test_detect_anomalies__with_invalid_id_cols():
"""Test time series anomaly detection tool invocation with invalid times_series_id_cols.""" """Test time series anomaly detection tool invocation with invalid times_series_id_cols."""
mock_credentials = mock.MagicMock(spec=Credentials) mock_credentials = mock.MagicMock(spec=Credentials)