Skip to content

Commit a953807

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
feat: add Bigtable tools
These tools support basic operations to interact with Bigtable table metadata and query results. PiperOrigin-RevId: 796571736
1 parent fa64545 commit a953807

File tree

14 files changed

+1094
-0
lines changed

14 files changed

+1094
-0
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ dependencies = [
3232
"click>=8.1.8, <9.0.0", # For CLI tools
3333
"fastapi>=0.115.0, <1.0.0", # FastAPI framework
3434
"google-api-python-client>=2.157.0, <3.0.0", # Google API client discovery
35+
"google-cloud-bigtable>=2.32.0", # For Bigtable database
3536
"google-cloud-aiplatform[agent_engines]>=1.95.1, <2.0.0", # For VertexAI integrations, e.g. example store.
3637
"google-cloud-secret-manager>=2.22.0, <3.0.0", # Fetching secrets in RestAPI Tool
3738
"google-cloud-spanner>=3.56.0, <4.0.0", # For Spanner database

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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+
"""Bigtable Tools (Experimental).
16+
17+
Bigtable tools under this module are hand crafted and customized while the tools
18+
under google.adk.tools.google_api_tool are auto generated based on API
19+
definition. The rationales to have customized tool are:
20+
21+
1. A dedicated Bigtable toolset to provide an easier, integrated way to interact
22+
with Bigtable for building AI Agent applications quickly.
23+
2. We want to provide extra access guardrails and controls in those tools.
24+
3. Use Bigtable Toolset for more customization and control to interact with
25+
Bigtable tables.
26+
"""
27+
28+
from .bigtable_credentials import BigtableCredentialsConfig
29+
from .bigtable_toolset import BigtableToolset
30+
31+
__all__ = [
32+
"BigtableToolset",
33+
"BigtableCredentialsConfig",
34+
]

‎src/google/adk/tools/bigtable/bigtable_credentials.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
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 ...utils.feature_decorator import experimental
18+
from .._google_credentials import BaseGoogleCredentialsConfig
19+
20+
BIGTABLE_TOKEN_CACHE_KEY = "bigtable_token_cache"
21+
BIGTABLE_DEFAULT_SCOPE = [
22+
"https://www.googleapis.com/auth/bigtable.admin",
23+
"https://www.googleapis.com/auth/bigtable.data",
24+
]
25+
26+
27+
@experimental
28+
class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
29+
"""Bigtable Credentials Configuration for Google API tools (Experimental).
30+
31+
Please do not use this in production, as it may be deprecated later.
32+
"""
33+
34+
def __post_init__(self) -> BigtableCredentialsConfig:
35+
"""Populate default scope if scopes is None."""
36+
super().__post_init__()
37+
38+
if not self.scopes:
39+
self.scopes = BIGTABLE_DEFAULT_SCOPE
40+
41+
# Set the token cache key
42+
self._token_cache_key = BIGTABLE_TOKEN_CACHE_KEY
43+
44+
return self

‎src/google/adk/tools/bigtable/bigtable_toolset.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
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 ...tools.google_tool import GoogleTool
30+
from ...utils.feature_decorator import experimental
31+
from .bigtable_credentials import BigtableCredentialsConfig
32+
from .settings import BigtableToolSettings
33+
34+
DEFAULT_BIGTABLE_TOOL_NAME_PREFIX = "bigtable"
35+
36+
37+
@experimental
38+
class BigtableToolset(BaseToolset):
39+
"""Bigtable Toolset contains tools for interacting with Bigtable data and metadata.
40+
41+
The tool names are:
42+
- bigtable_list_instances
43+
- bigtable_get_instance_info
44+
- bigtable_list_tables
45+
- bigtable_get_table_info
46+
- bigtable_execute_sql
47+
"""
48+
49+
def __init__(
50+
self,
51+
*,
52+
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
53+
credentials_config: Optional[BigtableCredentialsConfig] = None,
54+
bigtable_tool_settings: Optional[BigtableToolSettings] = None,
55+
):
56+
super().__init__(
57+
tool_filter=tool_filter,
58+
tool_name_prefix=DEFAULT_BIGTABLE_TOOL_NAME_PREFIX,
59+
)
60+
self._credentials_config = credentials_config
61+
self._tool_settings = (
62+
bigtable_tool_settings
63+
if bigtable_tool_settings
64+
else BigtableToolSettings()
65+
)
66+
67+
def _is_tool_selected(
68+
self, tool: BaseTool, readonly_context: ReadonlyContext
69+
) -> bool:
70+
if self.tool_filter is None:
71+
return True
72+
73+
if isinstance(self.tool_filter, ToolPredicate):
74+
return self.tool_filter(tool, readonly_context)
75+
76+
if isinstance(self.tool_filter, list):
77+
return tool.name in self.tool_filter
78+
79+
return False
80+
81+
@override
82+
async def get_tools(
83+
self, readonly_context: Optional[ReadonlyContext] = None
84+
) -> List[BaseTool]:
85+
"""Get tools from the toolset."""
86+
all_tools = [
87+
GoogleTool(
88+
func=func,
89+
credentials_config=self._credentials_config,
90+
tool_settings=self._tool_settings,
91+
)
92+
for func in [
93+
metadata_tool.list_instances,
94+
metadata_tool.get_instance_info,
95+
metadata_tool.list_tables,
96+
metadata_tool.get_table_info,
97+
query_tool.execute_sql,
98+
]
99+
]
100+
return [
101+
tool
102+
for tool in all_tools
103+
if self._is_tool_selected(tool, readonly_context)
104+
]

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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.auth.credentials import Credentials
19+
from google.cloud import bigtable
20+
from google.cloud.bigtable import data
21+
22+
from ... import version
23+
24+
USER_AGENT = f"adk-bigtable-tool google-adk/{version.__version__}"
25+
26+
27+
def _get_client_info() -> google.api_core.client_info.ClientInfo:
28+
"""Get client info."""
29+
return google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
30+
31+
32+
def get_bigtable_data_client(
33+
*, project: str, credentials: Credentials
34+
) -> bigtable.BigtableDataClient:
35+
"""Get a Bigtable client."""
36+
37+
bigtable_data_client = data.BigtableDataClient(
38+
project=project, credentials=credentials, client_info=_get_client_info()
39+
)
40+
41+
return bigtable_data_client
42+
43+
44+
def get_bigtable_admin_client(
45+
*, project: str, credentials: Credentials
46+
) -> bigtable.Client:
47+
"""Get a Bigtable client."""
48+
49+
bigtable_admin_client = bigtable.Client(
50+
project=project,
51+
admin=True,
52+
credentials=credentials,
53+
client_info=_get_client_info(),
54+
)
55+
56+
return bigtable_admin_client

‎src/google/adk/tools/bigtable/metadata_tool.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
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 logging
18+
19+
from google.auth.credentials import Credentials
20+
21+
from . import client
22+
23+
24+
def list_instances(project_id: str, credentials: Credentials) -> dict:
25+
"""List Bigtable instance ids in a Google Cloud project.
26+
27+
Args:
28+
project_id (str): The Google Cloud project id.
29+
credentials (Credentials): The credentials to use for the request.
30+
31+
Returns:
32+
dict: Dictionary with a list of the Bigtable instance ids present in the project.
33+
"""
34+
try:
35+
bt_client = client.get_bigtable_admin_client(
36+
project=project_id, credentials=credentials
37+
)
38+
(instances_list, failed_locations_list) = bt_client.list_instances()
39+
if failed_locations_list:
40+
logging.warning(
41+
"Failed to list instances from the following locations: %s",
42+
failed_locations_list,
43+
)
44+
instance_ids = [instance.instance_id for instance in instances_list]
45+
return {"status": "SUCCESS", "results": instance_ids}
46+
except Exception as ex:
47+
return {
48+
"status": "ERROR",
49+
"error_details": str(ex),
50+
}
51+
52+
53+
def get_instance_info(
54+
project_id: str, instance_id: str, credentials: Credentials
55+
) -> dict:
56+
"""Get metadata information about a Bigtable instance.
57+
58+
Args:
59+
project_id (str): The Google Cloud project id containing the instance.
60+
instance_id (str): The Bigtable instance id.
61+
credentials (Credentials): The credentials to use for the request.
62+
63+
Returns:
64+
dict: Dictionary representing the properties of the instance.
65+
"""
66+
try:
67+
bt_client = client.get_bigtable_admin_client(
68+
project=project_id, credentials=credentials
69+
)
70+
instance = bt_client.instance(instance_id)
71+
instance.reload()
72+
instance_info = {
73+
"project_id": project_id,
74+
"instance_id": instance.instance_id,
75+
"display_name": instance.display_name,
76+
"state": instance.state,
77+
"type": instance.type_,
78+
"labels": instance.labels,
79+
}
80+
return {"status": "SUCCESS", "results": instance_info}
81+
except Exception as ex:
82+
return {
83+
"status": "ERROR",
84+
"error_details": str(ex),
85+
}
86+
87+
88+
def list_tables(
89+
project_id: str, instance_id: str, credentials: Credentials
90+
) -> dict:
91+
"""List table ids in a Bigtable instance.
92+
93+
Args:
94+
project_id (str): The Google Cloud project id containing the instance.
95+
instance_id (str): The Bigtable instance id.
96+
credentials (Credentials): The credentials to use for the request.
97+
98+
Returns:
99+
dict: Dictionary with a list of the tables ids present in the instance.
100+
"""
101+
try:
102+
bt_client = client.get_bigtable_admin_client(
103+
project=project_id, credentials=credentials
104+
)
105+
instance = bt_client.instance(instance_id)
106+
tables = instance.list_tables()
107+
table_ids = [table.table_id for table in tables]
108+
return {"status": "SUCCESS", "results": table_ids}
109+
except Exception as ex:
110+
return {
111+
"status": "ERROR",
112+
"error_details": str(ex),
113+
}
114+
115+
116+
def get_table_info(
117+
project_id: str, instance_id: str, table_id: str, credentials: Credentials
118+
) -> dict:
119+
"""Get metadata information about a Bigtable table.
120+
121+
Args:
122+
project_id (str): The Google Cloud project id containing the instance.
123+
instance_id (str): The Bigtable instance id containing the table.
124+
table_id (str): The Bigtable table id.
125+
credentials (Credentials): The credentials to use for the request.
126+
127+
Returns:
128+
dict: Dictionary representing the properties of the table.
129+
"""
130+
try:
131+
bt_client = client.get_bigtable_admin_client(
132+
project=project_id, credentials=credentials
133+
)
134+
instance = bt_client.instance(instance_id)
135+
table = instance.table(table_id)
136+
column_families = table.list_column_families()
137+
table_info = {
138+
"project_id": project_id,
139+
"instance_id": instance.instance_id,
140+
"table_id": table.table_id,
141+
"column_families": list(column_families.keys()),
142+
}
143+
return {"status": "SUCCESS", "results": table_info}
144+
except Exception as ex:
145+
return {
146+
"status": "ERROR",
147+
"error_details": str(ex),
148+
}

0 commit comments

Comments
 (0)