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
This commit is contained in:
Google Team Member
2025-08-29 14:56:47 -07:00
committed by Copybara-Service
parent a17bcbb2aa
commit 2eddc5e4d3
3 changed files with 43 additions and 0 deletions
@@ -1006,3 +1006,25 @@ def test_execute_sql_bq_client_creation(mock_get_bigquery_client):
mock_get_bigquery_client.call_args.kwargs["user_agent"]
== application_name
)
def test_execute_sql_unexpected_project_id():
"""Test execute_sql tool invocation with unexpected project id."""
compute_project_id = "compute_project_id"
tool_call_project_id = "project_id"
query = "SELECT 1"
credentials = mock.create_autospec(Credentials, instance=True)
tool_settings = BigQueryToolConfig(compute_project_id=compute_project_id)
tool_context = mock.create_autospec(ToolContext, instance=True)
result = execute_sql(
tool_call_project_id, query, credentials, tool_settings, tool_context
)
assert result == {
"status": "ERROR",
"error_details": (
f"Cannot execute query in the project {tool_call_project_id}, as the"
" tool is restricted to execute queries only in the project"
f" {compute_project_id}."
),
}