encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: allow setting agent/application name for BigQuery tools

Browse Source 
This will allow tracking of tool usage per agent/application.

PiperOrigin-RevId: 800607186
This commit is contained in:
Google Team Member
2025-08-28 14:10:01 -07:00
committed by Copybara-Service
parent f4a8df0ba2
commit 11a2ffe35a
9 changed files with 259 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/tools/bigquery/test_bigquery_query_tool.py
+23
View File
@@ -983,3 +983,26 @@ def test_execute_sql_result_dtype(
# Test the tool worked without invoking default auth
result = execute_sql(project, query, credentials, tool_settings, tool_context)
assert result == {"status": "SUCCESS", "rows": tool_result_rows}
@mock.patch(
"google.adk.tools.bigquery.client.get_bigquery_client", autospec=True
)
def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
"""Test BigQuery client creation params during execute_sql tool invocation."""
project = "my_project_id"
query = "SELECT 1"
credentials = mock.create_autospec(Credentials, instance=True)
application_name = "my-agent"
tool_settings = BigQueryToolConfig(application_name=application_name)
tool_context = mock.create_autospec(ToolContext, instance=True)
execute_sql(project, query, credentials, tool_settings, tool_context)
mock_get_bigquery_client.assert_called_once()
assert len(mock_get_bigquery_client.call_args.kwargs) == 3
assert mock_get_bigquery_client.call_args.kwargs["project"] == project
assert mock_get_bigquery_client.call_args.kwargs["credentials"] == credentials
assert (
mock_get_bigquery_client.call_args.kwargs["user_agent"]
== application_name
)