mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
a17bcbb2aa
commit
2eddc5e4d3
@@ -71,6 +71,13 @@ class BigQueryToolConfig(BaseModel):
|
||||
their application/agent for tracking or support purpose, they can set this field.
|
||||
"""
|
||||
|
||||
compute_project_id: Optional[str] = None
|
||||
"""GCP project ID to use for the BigQuery compute operations.
|
||||
|
||||
This can be set as a guardrail to ensure that the tools perform the compute
|
||||
operations (such as query execution) in a specific project.
|
||||
"""
|
||||
|
||||
@field_validator('application_name')
|
||||
@classmethod
|
||||
def validate_application_name(cls, v):
|
||||
|
||||
@@ -78,6 +78,20 @@ def execute_sql(
|
||||
}
|
||||
"""
|
||||
try:
|
||||
# Validate compute project if applicable
|
||||
if (
|
||||
settings.compute_project_id
|
||||
and project_id != settings.compute_project_id
|
||||
):
|
||||
return {
|
||||
"status": "ERROR",
|
||||
"error_details": (
|
||||
f"Cannot execute query in the project {project_id}, as the tool"
|
||||
" is restricted to execute queries only in the project"
|
||||
f" {settings.compute_project_id}."
|
||||
),
|
||||
}
|
||||
|
||||
# Get BigQuery client
|
||||
bq_client = client.get_bigquery_client(
|
||||
project=project_id,
|
||||
|
||||
@@ -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}."
|
||||
),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user