Skip to content

Commit 2eddc5e

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: allow setting compute project for BigQuery tools
This will allow restricting BigQuery SQL executions to the specified project. The agent/LLM should resolve the `project_id` param for tools like `execute_sql` and sometimes they can resolve it to an unexpected value due to hallucination or ambiguity. This guardrail will protect against that situation. PiperOrigin-RevId: 801039685
1 parent a17bcbb commit 2eddc5e

File tree

3 files changed

+43
-0
lines changed

3 files changed

+43
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,13 @@ class BigQueryToolConfig(BaseModel):
7171
their application/agent for tracking or support purpose, they can set this field.
7272
"""
7373

74+
compute_project_id: Optional[str] = None
75+
"""GCP project ID to use for the BigQuery compute operations.
76+
77+
This can be set as a guardrail to ensure that the tools perform the compute
78+
operations (such as query execution) in a specific project.
79+
"""
80+
7481
@field_validator('application_name')
7582
@classmethod
7683
def validate_application_name(cls, v):

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,20 @@ def execute_sql(
7878
}
7979
"""
8080
try:
81+
# Validate compute project if applicable
82+
if (
83+
settings.compute_project_id
84+
and project_id != settings.compute_project_id
85+
):
86+
return {
87+
"status": "ERROR",
88+
"error_details": (
89+
f"Cannot execute query in the project {project_id}, as the tool"
90+
" is restricted to execute queries only in the project"
91+
f" {settings.compute_project_id}."
92+
),
93+
}
94+
8195
# Get BigQuery client
8296
bq_client = client.get_bigquery_client(
8397
project=project_id,

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1006,3 +1006,25 @@ def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
10061006
mock_get_bigquery_client.call_args.kwargs["user_agent"]
10071007
== application_name
10081008
)
1009+
1010+
1011+
def test_execute_sql_unexpected_project_id():
1012+
"""Test execute_sql tool invocation with unexpected project id."""
1013+
compute_project_id = "compute_project_id"
1014+
tool_call_project_id = "project_id"
1015+
query = "SELECT 1"
1016+
credentials = mock.create_autospec(Credentials, instance=True)
1017+
tool_settings = BigQueryToolConfig(compute_project_id=compute_project_id)
1018+
tool_context = mock.create_autospec(ToolContext, instance=True)
1019+
1020+
result = execute_sql(
1021+
tool_call_project_id, query, credentials, tool_settings, tool_context
1022+
)
1023+
assert result == {
1024+
"status": "ERROR",
1025+
"error_details": (
1026+
f"Cannot execute query in the project {tool_call_project_id}, as the"
1027+
" tool is restricted to execute queries only in the project"
1028+
f" {compute_project_id}."
1029+
),
1030+
}

0 commit comments

Comments
 (0)