fix: stop updating write mode in the global settings during tool execution

Two tools - detect_anomalies and analyze_contribution are modifying the settings passed to them, which is not right as the settings are held and passed by the top level, which means several tools share the same settings.

PiperOrigin-RevId: 832081738
This commit is contained in:
Google Team Member
2025-11-13 18:02:19 -08:00
committed by Copybara-Service
parent 23ad40bad2
commit 5adbf95a0a
2 changed files with 102 additions and 20 deletions
+18 -20
View File
@@ -1093,21 +1093,23 @@ def analyze_contribution(
"""
# Create a session and run the create model query.
original_write_mode = settings.write_mode
try:
if original_write_mode == WriteMode.BLOCKED:
execute_sql_settings = settings
if execute_sql_settings.write_mode == WriteMode.BLOCKED:
raise ValueError("analyze_contribution is not allowed in this session.")
elif original_write_mode != WriteMode.PROTECTED:
elif execute_sql_settings.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
execute_sql_settings = settings.model_copy(
update={"write_mode": WriteMode.PROTECTED}
)
result = _execute_sql(
project_id=project_id,
query=create_model_query,
credentials=credentials,
settings=settings,
settings=execute_sql_settings,
tool_context=tool_context,
caller_id="analyze_contribution",
)
@@ -1118,18 +1120,15 @@ def analyze_contribution(
project_id=project_id,
query=get_insights_query,
credentials=credentials,
settings=settings,
settings=execute_sql_settings,
tool_context=tool_context,
caller_id="analyze_contribution",
)
except Exception as ex: # pylint: disable=broad-except
return {
"status": "ERROR",
"error_details": f"Error during analyze_contribution: {str(ex)}",
"error_details": f"Error during analyze_contribution: {repr(ex)}",
}
finally:
# Restore the original write mode.
settings.write_mode == original_write_mode
return result
@@ -1327,21 +1326,23 @@ def detect_anomalies(
"""
# Create a session and run the create model query.
original_write_mode = settings.write_mode
try:
if settings.write_mode == WriteMode.BLOCKED:
execute_sql_settings = settings
if execute_sql_settings.write_mode == WriteMode.BLOCKED:
raise ValueError("anomaly detection is not allowed in this session.")
elif original_write_mode != WriteMode.PROTECTED:
elif execute_sql_settings.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
execute_sql_settings = settings.model_copy(
update={"write_mode": WriteMode.PROTECTED}
)
result = _execute_sql(
project_id=project_id,
query=create_model_query,
credentials=credentials,
settings=settings,
settings=execute_sql_settings,
tool_context=tool_context,
caller_id="detect_anomalies",
)
@@ -1352,17 +1353,14 @@ def detect_anomalies(
project_id=project_id,
query=anomaly_detection_query,
credentials=credentials,
settings=settings,
settings=execute_sql_settings,
tool_context=tool_context,
caller_id="detect_anomalies",
)
except Exception as ex: # pylint: disable=broad-except
return {
"status": "ERROR",
"error_details": f"Error during anomaly detection: {str(ex)}",
"error_details": f"Error during anomaly detection: {repr(ex)}",
}
finally:
# Restore the original write mode.
settings.write_mode == original_write_mode
return result
@@ -1849,3 +1849,87 @@ def test_execute_sql_maximum_bytes_billed_config():
bq_client.query_and_wait.assert_called_once()
call_args = bq_client.query_and_wait.call_args
assert call_args.kwargs["job_config"].maximum_bytes_billed == 11_000_000
@pytest.mark.parametrize(
("tool_call",),
[
pytest.param(
lambda settings, tool_context: execute_sql(
project_id="test-project",
query="SELECT * FROM `test-dataset.test-table`",
credentials=mock.create_autospec(Credentials, instance=True),
settings=settings,
tool_context=tool_context,
),
id="execute-sql",
),
pytest.param(
lambda settings, tool_context: forecast(
project_id="test-project",
history_data="SELECT * FROM `test-dataset.test-table`",
timestamp_col="ts_col",
data_col="data_col",
credentials=mock.create_autospec(Credentials, instance=True),
settings=settings,
tool_context=tool_context,
),
id="forecast",
),
pytest.param(
lambda settings, tool_context: analyze_contribution(
project_id="test-project",
input_data="test-dataset.test-table",
dimension_id_cols=["dim1", "dim2"],
contribution_metric="SUM(metric)",
is_test_col="is_test",
credentials=mock.create_autospec(Credentials, instance=True),
settings=settings,
tool_context=tool_context,
),
id="analyze-contribution",
),
pytest.param(
lambda settings, tool_context: detect_anomalies(
project_id="test-project",
history_data="SELECT * FROM `test-dataset.test-table`",
times_series_timestamp_col="ts_timestamp",
times_series_data_col="ts_data",
credentials=mock.create_autospec(Credentials, instance=True),
settings=settings,
tool_context=tool_context,
),
id="detect-anomalies",
),
],
)
def test_tool_call_doesnt_change_global_settings(tool_call):
"""Test query tools don't change global settings."""
settings = BigQueryToolConfig(write_mode=WriteMode.ALLOWED)
tool_context = mock.create_autospec(ToolContext, instance=True)
tool_context.state.get.return_value = (
"test-bq-session-id",
"_anonymous_dataset",
)
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
# The mock instance
bq_client = Client.return_value
# Simulate the result of query API
query_job = mock.create_autospec(bigquery.QueryJob)
query_job.destination.dataset_id = "_anonymous_dataset"
bq_client.query.return_value = query_job
bq_client.query_and_wait.return_value = []
# Test settings write mode before
assert settings.write_mode == WriteMode.ALLOWED
# Call the tool
result = tool_call(settings, tool_context)
# Test successfull executeion of the tool
assert result == {"status": "SUCCESS", "rows": []}
# Test settings write mode after
assert settings.write_mode == WriteMode.ALLOWED