mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
23ad40bad2
commit
5adbf95a0a
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user