diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index 23382579..6d35af0c 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/src/google/adk/tools/bigquery/query_tool.py @@ -1136,7 +1136,7 @@ def detect_anomalies( history_data: str, times_series_timestamp_col: str, times_series_data_col: str, - horizon: Optional[int] = 10, + horizon: Optional[int] = 1000, target_data: Optional[str] = None, times_series_id_cols: Optional[list[str]] = None, anomaly_prob_threshold: Optional[float] = 0.95, @@ -1158,7 +1158,7 @@ def detect_anomalies( 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. + future. Defaults to 1000. 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. @@ -1301,9 +1301,14 @@ def detect_anomalies( OPTIONS ({options_str}) AS {history_data_source} """ + order_by_id_cols = ( + ", ".join(col for col in times_series_id_cols) + ", " + if times_series_id_cols + else "" + ) 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)) ORDER BY {order_by_id_cols}{times_series_timestamp_col} """ if target_data: trimmed_upper_target_data = target_data.strip().upper() @@ -1312,10 +1317,10 @@ def detect_anomalies( ) or trimmed_upper_target_data.startswith("WITH"): target_data_source = f"({target_data})" else: - target_data_source = f"SELECT * FROM `{target_data}`" + 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}) + SELECT * FROM ML.DETECT_ANOMALIES(MODEL {model_name}, STRUCT({anomaly_prob_threshold} AS anomaly_prob_threshold), {target_data_source}) ORDER BY {order_by_id_cols}{times_series_timestamp_col} """ # Create a session and run the create model query. diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index bba8d4aa..d4906e3d 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py @@ -1436,12 +1436,12 @@ def test_detect_anomalies_with_table_id(mock_uuid, mock_execute_sql): 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) + OPTIONS (MODEL_TYPE = 'ARIMA_PLUS', TIME_SERIES_TIMESTAMP_COL = 'ts_timestamp', TIME_SERIES_DATA_COL = 'ts_data', HORIZON = 1000) 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)) + SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.95 AS anomaly_prob_threshold)) ORDER BY ts_timestamp """ assert mock_execute_sql.call_count == 2 @@ -1497,7 +1497,7 @@ def test_detect_anomalies_with_custom_params(mock_uuid, mock_execute_sql): """ expected_anomaly_detection_query = """ - SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.8 AS anomaly_prob_threshold)) + SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.8 AS anomaly_prob_threshold)) ORDER BY dim1, dim2, ts_timestamp """ assert mock_execute_sql.call_count == 2 @@ -1555,7 +1555,61 @@ def test_detect_anomalies_on_target_table(mock_uuid, mock_execute_sql): """ 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`)) + SELECT * FROM ML.DETECT_ANOMALIES(MODEL detect_anomalies_model_test_uuid, STRUCT(0.8 AS anomaly_prob_threshold), (SELECT * FROM `test-dataset.target-table`)) ORDER BY dim1, dim2, ts_timestamp + """ + + assert mock_execute_sql.call_count == 2 + mock_execute_sql.assert_any_call( + project_id="test-project", + query=expected_create_model_query, + credentials=mock_credentials, + settings=mock_settings, + tool_context=mock_tool_context, + caller_id="detect_anomalies", + ) + mock_execute_sql.assert_any_call( + project_id="test-project", + query=expected_anomaly_detection_query, + credentials=mock_credentials, + settings=mock_settings, + tool_context=mock_tool_context, + caller_id="detect_anomalies", + ) + + +# 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_str_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", + target_data="test-dataset.target-table", + 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 = 1000) + 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), (SELECT * FROM `test-dataset.target-table`)) ORDER BY ts_timestamp """ assert mock_execute_sql.call_count == 2