8000 feat: add BigQuery first-party tools. · SVendittelli/adk-python@d6c6bb4 · GitHub
[go: up one dir, main page]
Skip to content

Commit d6c6bb4

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: add BigQuery first-party tools.
These tools support getting BigQuery dataset/table metadata and query results. PiperOrigin-RevId: 764139132
1 parent 46282ee commit d6c6bb4

File tree

11 files changed

+748
-11
lines changed

11 files changed

+748
-11
lines changed

‎contributing/samples/bigquery/README.md

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# BigQuery Tools Sample
2+
3+
## Introduction
4+
5+
This sample agent demonstrates the BigQuery first-party tools in ADK,
6+
distributed via the `google.adk.tools.bigquery` module. These tools include:
7+
8+
1. `list_dataset_ids`
9+
10+
Fetches BigQuery dataset ids present in a GCP project.
11+
12+
1. `get_dataset_info`
13+
14+
Fetches metadata about a BigQuery dataset.
15+
16+
1. `list_table_ids`
17+
18+
Fetches table ids present in a BigQuery dataset.
19+
20+
1. `get_table_info`
21+
22+
Fetches metadata about a BigQuery table.
23+
24+
1. `execute_sql`
25+
26+
Runs a SQL query in BigQuery.
27+
28+
## How to use
29+
30+
Set up environment variables in your `.env` file for using
31+
[Google AI Studio](https://google.github.io/adk-docs/get-started/quickstart/#gemini---google-ai-studio)
32+
or
33+
[Google Clo 8000 ud Vertex AI](https://google.github.io/adk-docs/get-started/quickstart/#gemini---google-cloud-vertex-ai)
34+
for the LLM service for your agent. For example, for using Google AI Studio you
35+
would set:
36+
37+
* GOOGLE_GENAI_USE_VERTEXAI=FALSE
38+
* GOOGLE_API_KEY={your api key}
39+
40+
### With Application Default Credentials
41+
42+
This mode is useful for quick development when the agent builder is the only
43+
user interacting with the agent. The tools are initialized with the default
44+
credentials present on the machine running the agent.
45+
46+
1. Create application default credentials on the machine where the agent would
47+
be running by following https://cloud.google.com/docs/authentication/provide-credentials-adc.
48+
49+
1. Set `RUN_WITH_ADC=True` in `agent.py` and run the agent
50+
51+
### With Interactive OAuth
52+
53+
1. Follow
54+
https://developers.google.com/identity/protocols/oauth2#1.-obtain-oauth-2.0-credentials-from-the-dynamic_data.setvar.console_name.
55+
to get your client id and client secret. Be sure to choose "web" as your client
56+
type.
57+
58+
1. Follow https://developers.google.com/workspace/guides/configure-oauth-consent to add scope "https://www.googleapis.com/auth/bigquery".
59+
60+
1. Follow https://developers.google.com/identity/protocols/oauth2/web-server#creatingcred to add http://localhost/dev-ui/ to "Authorized redirect URIs".
61+
62+
Note: localhost here is just a hostname that you use to access the dev ui,
63+
replace it with the actual hostname you use to access the dev ui.
64+
65+
1. For 1st run, allow popup for localhost in Chrome.
66+
67+
1. Configure your `.env` file to add two more variables before running the agent:
68+
69+
* OAUTH_CLIENT_ID={your client id}
70+
* OAUTH_CLIENT_SECRET={your client secret}
71+
72+
Note: don't create a separate .env, instead put it to the same .env file that
73+
stores your Vertex AI or Dev ML credentials
74+
75+
1. Set `RUN_WITH_ADC=False` in `agent.py` and run the agent
76+
77+
## Sample prompts
78+
79+
* which weather datasets exist in bigquery public data?
80+
* tell me more about noaa_lightning
81+
* which tables exist in the ml_datasets dataset?
82+
* show more details about the penguins table
83+
* compute penguins population per island.

‎contributing/samples/bigquery/__init__.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from . import agent

‎contributing/samples/bigquery/agent.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
from google.adk.agents import llm_agent
18+
from google.adk.tools.bigquery import BigQueryCredentialsConfig
19+
from google.adk.tools.bigquery import BigQueryToolset
20+
import google.auth
21+
22+
23+
RUN_WITH_ADC = False
24+
25+
26+
if RUN_WITH_ADC:
27+
# Initialize the tools to use the application default credentials.
28+
application_default_credentials, _ = google.auth.default()
29+
credentials_config = BigQueryCredentialsConfig(
30+
credentials=application_default_credentials
31+
)
32+
else:
33+
# Initiaze the tools to do interactive OAuth
34+
# The environment variables OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET
35+
# must be set
36+
credentials_config = BigQueryCredentialsConfig(
37+
client_id=os.getenv("OAUTH_CLIENT_ID"),
38+
client_secret=os.getenv("OAUTH_CLIENT_SECRET"),
39+
scopes=["https://www.googleapis.com/auth/bigquery"],
40+
)
41+
42+
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)
43+
44+
# The variable name `root_agent` determines what your root agent is for the
45+
# debug CLI
46+
root_agent = llm_agent.Agent(
47+
model="gemini-2.0-flash",
48+
name="hello_agent",
49+
description=(
50+
"Agent to answer questions about BigQuery data and models and execute"
51+
" SQL queries."
52+
),
53+
instruction="""\
54+
You are a data science agent with access to several BigQuery tools.
55+
Make use of those tools to answer the user's questions.
56+
""",
57+
tools=[bigquery_toolset],
58+
)

‎src/google/adk/tools/bigquery/__init__.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
"""BigQuery Tools. (Experimental)
15+
"""BigQuery Tools (Experimental).
1616
1717
BigQuery Tools under this module are hand crafted and customized while the tools
1818
under google.adk.tools.google_api_tool are auto generated based on API
@@ -26,3 +26,13 @@
2626
4. We want to provide extra access guardrails in those tools. For example,
2727
execute_sql can't arbitrarily mutate existing data.
2828
"""
29+
30+
from .bigquery_credentials import BigQueryCredentialsConfig
31+
from .bigquery_tool import BigQueryTool
32+
from .bigquery_toolset import BigQueryToolset
33+
34+
__all__ = [
35+
"BigQueryTool",
36+
"BigQueryToolset",
37+
"BigQueryCredentialsConfig",
38+
]

‎src/google/adk/tools/bigquery/bigquery_credentials.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
from __future__ import annotations
1616

17+
import json
1718
from typing import List
1819
from typing import Optional
1920

@@ -121,7 +122,7 @@ async def get_valid_credentials(
121122
creds_json = tool_context.state.get(BIGQUERY_TOKEN_CACHE_KEY, None)
122123
creds = (
123124
Credentials.from_authorized_user_info(
124-
creds_json, self.credentials_config.scopes
125+
json.loads(creds_json), self.credentials_config.scopes
125126
)
126127
if creds_json
127128
else None

‎src/google/adk/tools/bigquery/bigquery_toolset.py

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from typing import List
18+
from typing import Optional
19+
from typing import Union
20+
21+
from google.adk.agents.readonly_context import ReadonlyContext
22+
from typing_extensions import override
23+
24+
from . import metadata_tool
25+
from . import query_tool
26+
from ...tools.base_tool import BaseTool
27+
from ...tools.base_toolset import BaseToolset
28+
from ...tools.base_toolset import ToolPredicate
29+
from .bigquery_credentials import BigQueryCredentialsConfig
30+
from .bigquery_tool import BigQueryTool
31+
32+
33+
class BigQueryToolset(BaseToolset):
34+
"""BigQuery Toolset contains tools for interacting with BigQuery data and metadata."""
35+
36+
def __init__(
37+
self,
38+
*,
39+
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
40+
credentials_config: Optional[BigQueryCredentialsConfig] = None,
41+
):
42+
self._credentials_config = credentials_config
43+
self.tool_filter = tool_filter
44+
45+
def _is_tool_selected(
46+
self, tool: BaseTool, readonly_context: ReadonlyContext
47+
) -> bool:
48+
if self.tool_filter is None:
49+
return True
50+
51+
if isinstance(self.tool_filter, ToolPredicate):
52+
return self.tool_filter(tool, readonly_context)
53+
54+
if isinstance(self.tool_filter, list):
55+
return tool.name in self.tool_filter
56+
57+
return False
58+
59+
@override
60+
async def get_tools(
61+
self, readonly_context: Optional[ReadonlyContext] = None
62+
) -> List[BaseTool]:
63+
"""Get tools from the toolset."""
64+
all_tools = [
65+
BigQueryTool(
66+
func=func,
67+
credentials=self._credentials_config,
68+
)
69+
for func in [
70+
metadata_tool.get_dataset_info,
71+
metadata_tool.get_table_info,
72+
metadata_tool.list_dataset_ids,
73+
metadata_tool.list_table_ids,
74+
query_tool.execute_sql,
75+
]
76+
]
77+
78+
return [
79+
tool
80+
for tool in all_tools
81+
if self._is_tool_selected(tool, readonly_context)
82+
]
83+
84+
@override
85+
async def close(self):
86+
pass

‎src/google/adk/tools/bigquery/client.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import google.api_core.client_info
18+
from google.cloud import bigquery
19+
from google.oauth2.credentials import Credentials
20+
21+
USER_AGENT = "adk-bigquery-tool"
22+
23+
24+
def get_bigquery_client(*, credentials: Credentials) -> bigquery.Client:
25+
"""Get a BigQuery client."""
26+
27+
client_info = google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
28+
29+
bigquery_client = bigquery.Client(
30+
credentials=credentials, client_info=client_info
31+
)
32+
33+
return bigquery_client

0 commit comments

Comments
 (0)
0