Skip to content

Commit f13a11e

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Propagate application_name set for the BigQuery Tools as BigQuery job labels
This change will help the tools user identify per agent job usage in BQ console and INFORMATION_SCHEMA views. This change fulfills the feature request #3582. Here is a demo after change: screen/C6YB4ge2FM2ZREi. PiperOrigin-RevId: 834480140
1 parent 131d39c commit f13a11e

File tree

3 files changed

+88
-6
lines changed

3 files changed

+88
-6
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,9 @@ class BigQueryToolConfig(BaseModel):
8080
8181
By default, no particular application name will be set in the BigQuery
8282
interaction. But if the tool user (agent builder) wants to differentiate
83-
their application/agent for tracking or support purpose, they can set this field.
83+
their application/agent for tracking or support purpose, they can set this
84+
field. If set, this value will be added to the user_agent in BigQuery API calls, and also to the BigQuery job labels with the key
85+
"adk-bigquery-application-name".
8486
"""
8587

8688
compute_project_id: Optional[str] = None

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,8 @@ def _execute_sql(
7171
bq_job_labels = {}
7272
if caller_id:
7373
bq_job_labels["adk-bigquery-tool"] = caller_id
74+
if settings and settings.application_name:
75+
bq_job_labels["adk-bigquery-application-name"] = settings.application_name
7476

7577
if not settings or settings.write_mode == WriteMode.BLOCKED:
7678
dry_run_query_job = bq_client.query(

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

Lines changed: 83 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1673,7 +1673,9 @@ def test_execute_sql_job_labels(
16731673
query = "SELECT 123 AS num"
16741674
statement_type = "SELECT"
16751675
credentials = mock.create_autospec(Credentials, instance=True)
1676-
tool_settings = BigQueryToolConfig(write_mode=write_mode)
1676+
tool_settings = BigQueryToolConfig(
1677+
write_mode=write_mode, application_name="test-app"
1678+
)
16771679
tool_context = mock.create_autospec(ToolContext, instance=True)
16781680
tool_context.state.get.return_value = None
16791681

@@ -1702,12 +1704,13 @@ def test_execute_sql_job_labels(
17021704
for call_args in call_args_list:
17031705
_, mock_kwargs = call_args
17041706
assert mock_kwargs["job_config"].labels == {
1705-
"adk-bigquery-tool": "execute_sql"
1707+
"adk-bigquery-tool": "execute_sql",
1708+
"adk-bigquery-application-name": "test-app",
17061709
}
17071710

17081711

17091712
@pytest.mark.parametrize(
1710-
("tool_call", "expected_label"),
1713+
("tool_call", "expected_tool_label"),
17111714
[
17121715
pytest.param(
17131716
lambda tool_context: query_tool.forecast(
@@ -1751,7 +1754,7 @@ def test_execute_sql_job_labels(
17511754
),
17521755
],
17531756
)
1754-
def test_ml_tool_job_labels(tool_call, expected_label):
1757+
def test_ml_tool_job_labels(tool_call, expected_tool_label):
17551758
"""Test ML tools for job label."""
17561759

17571760
with mock.patch.object(bigquery, "Client", autospec=True) as Client:
@@ -1768,10 +1771,85 @@ def test_ml_tool_job_labels(tool_call, expected_label):
17681771
for call_args in call_args_list:
17691772
_, mock_kwargs = call_args
17701773
assert mock_kwargs["job_config"].labels == {
1771-
"adk-bigquery-tool": expected_label
1774+
"adk-bigquery-tool": expected_tool_label
17721775
}
17731776

17741777

1778+
@pytest.mark.parametrize(
1779+
("tool_call", "expected_tool_label"),
1780+
[
1781+
pytest.param(
1782+
lambda tool_context: query_tool.forecast(
1783+
project_id="test-project",
1784+
history_data="SELECT * FROM `test-dataset.test-table`",
1785+
timestamp_col="ts_col",
1786+
data_col="data_col",
1787+
credentials=mock.create_autospec(Credentials, instance=True),
1788+
settings=BigQueryToolConfig(
1789+
write_mode=WriteMode.ALLOWED, application_name="test-app"
1790+
),
1791+
tool_context=tool_context,
1792+
),
1793+
"forecast",
1794+
id="forecast-app-name",
1795+
),
1796+
pytest.param(
1797+
lambda tool_context: query_tool.analyze_contribution(
1798+
project_id="test-project",
1799+
input_data="test-dataset.test-table",
1800+
dimension_id_cols=["dim1", "dim2"],
1801+
contribution_metric="SUM(metric)",
1802+
is_test_col="is_test",
1803+
credentials=mock.create_autospec(Credentials, instance=True),
1804+
settings=BigQueryToolConfig(
1805+
write_mode=WriteMode.ALLOWED, application_name="test-app"
1806+
),
1807+
tool_context=tool_context,
1808+
),
1809+
"analyze_contribution",
1810+
id="analyze-contribution-app-name",
1811+
),
1812+
pytest.param(
1813+
lambda tool_context: query_tool.detect_anomalies(
1814+
project_id="test-project",
1815+
history_data="SELECT * FROM `test-dataset.test-table`",
1816+
times_series_timestamp_col="ts_timestamp",
1817+
times_series_data_col="ts_data",
1818+
credentials=mock.create_autospec(Credentials, instance=True),
1819+
settings=BigQueryToolConfig(
1820+
write_mode=WriteMode.ALLOWED, application_name="test-app"
1821+
),
1822+
tool_context=tool_context,
1823+
),
1824+
"detect_anomalies",
1825+
id="detect-anomalies-app-name",
1826+
),
1827+
],
1828+
)
1829+
def test_ml_tool_job_labels_w_application_name(tool_call, expected_tool_label):
1830+
"""Test ML tools for job label with application name."""
1831+
1832+
with mock.patch.object(bigquery, "Client", autospec=True) as Client:
1833+
bq_client = Client.return_value
1834+
1835+
tool_context = mock.create_autospec(ToolContext, instance=True)
1836+
tool_context.state.get.return_value = None
1837+
tool_call(tool_context)
1838+
1839+
expected_labels = {
1840+
"adk-bigquery-tool": expected_tool_label,
1841+
"adk-bigquery-application-name": "test-app",
1842+
}
1843+
1844+
for call_args_list in [
1845+
bq_client.query.call_args_list,
1846+
bq_client.query_and_wait.call_args_list,
1847+
]:
1848+
for call_args in call_args_list:
1849+
_, mock_kwargs = call_args
1850+
assert mock_kwargs["job_config"].labels == expected_labels
1851+
1852+
17751853
def test_execute_sql_max_rows_config():
17761854
"""Test execute_sql tool respects max_query_result_rows from config."""
17771855
project = "my_project"

0 commit comments

Comments
 (0)