Skip to content

Commit 97a5759

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: add Bigtable first-party tools
These tools support basic operations to interact with Bigtable table metadata and query results. PiperOrigin-RevId: 787191308
1 parent a74d334 commit 97a5759

14 files changed

Lines changed: 1068 additions & 0 deletions

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
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
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+
"""Tools for interacting with Bigtable."""
16+
17+
from .bigtable_credentials import BigtableCredentialsConfig
18+
from .bigtable_toolset import BigtableToolset
19+
20+
__all__ = [
21+
"BigtableToolset",
22+
"BigtableCredentialsConfig",
23+
]
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
Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
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+
33+
34+
@experimental
35+
class BigtableToolset(BaseToolset):
36+
"""Bigtable Toolset contains tools for interacting with Bigtable data and metadata."""
37+
38+
def __init__(
39+
self,
40+
*,
41+
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
42+
credentials_config: Optional[BigtableCredentialsConfig] = None,
43+
):
44+
super().__init__(tool_filter=tool_filter)
45+
self._credentials_config = credentials_config
46+
47+
def _is_tool_selected(
48+
self, tool: BaseTool, readonly_context: ReadonlyContext
49+
) -> bool:
50+
if self.tool_filter is None:
51+
return True
52+
53+
if isinstance(self.tool_filter, ToolPredicate):
54+
return self.tool_filter(tool, readonly_context)
55+
56+
if isinstance(self.tool_filter, list):
57+
return tool.name in self.tool_filter
58+
59+
return False
60+
61+
@override
62+
async def get_tools(
63+
self, readonly_context: Optional[ReadonlyContext] = None
64+
) -> List[BaseTool]:
65+
"""Get tools from the toolset."""
66+
all_tools = [
67+
GoogleTool(
68+
func=func,
69+
credentials_config=self._credentials_config,
70+
)
71+
for func in [
72+
metadata_tool.list_instances,
73+
metadata_tool.get_instance_info,
74+
metadata_tool.list_tables,
75+
metadata_tool.get_table_info,
76+
query_tool.execute_sql,
77+
]
78+
]
79+
return [
80+
tool
81+
for tool in all_tools
82+
if self._is_tool_selected(tool, readonly_context)
83+
]
84+
85+
@override
86+
async def close(self):
87+
pass
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
Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
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 google.auth.credentials import Credentials
18+
19+
from . import client
20+
21+
22+
def list_instances(project_id: str, credentials: Credentials) -> dict:
23+
"""List Bigtable instance ids in a Google Cloud project.
24+
25+
Args:
26+
project_id (str): The Google Cloud project id.
27+
credentials (Credentials): The credentials to use for the request.
28+
29+
Returns:
30+
dict: Dictionary with a list of the Bigtable instance ids present in the project.
31+
"""
32+
try:
33+
bt_client = client.get_bigtable_admin_client(
34+
project=project_id, credentials=credentials
35+
)
36+
(instances_list, failed_locations_list) = bt_client.list_instances()
37+
instance_ids = [instance.instance_id for instance in instances_list]
38+
return {"status": "SUCCESS", "results": instance_ids}
39+
except Exception as ex:
40+
return {
41+
"status": "ERROR",
42+
"error_details": str(ex),
43+
}
44+
45+
46+
def get_instance_info(
47+
project_id: str, instance_id: str, credentials: Credentials
48+
) -> dict:
49+
"""Get metadata information about a Bigtable instance.
50+
51+
Args:
52+
project_id (str): The Google Cloud project id containing the instance.
53+
instance_id (str): The Bigtable instance id.
54+
credentials (Credentials): The credentials to use for the request.
55+
56+
Returns:
57+
dict: Dictionary representing the properties of the instance.
58+
"""
59+
try:
60+
bt_client = client.get_bigtable_admin_client(
61+
project=project_id, credentials=credentials
62+
)
63+
instance = bt_client.instance(instance_id)
64+
instance.reload()
65+
instance_info = {
66+
"project_id": project_id,
67+
"instance_id": instance.instance_id,
68+
"display_name": instance.display_name,
69+
"state": instance.state,
70+
"type": instance.type_,
71+
"labels": instance.labels,
72+
}
73+
return {"status": "SUCCESS", "results": instance_info}
74+
except Exception as ex:
75+
return {
76+
"status": "ERROR",
77+
"error_details": str(ex),
78+
}
79+
80+
81+
def list_tables(
82+
project_id: str, instance_id: str, credentials: Credentials
83+
) -> dict:
84+
"""List table ids in a Bigtable instance.
85+
86+
Args:
87+
project_id (str): The Google Cloud project id containing the instance.
88+
instance_id (str): The Bigtable instance id.
89+
credentials (Credentials): The credentials to use for the request.
90+
91+
Returns:
92+
dict: Dictionary with a list of the tables ids present in the instance.
93+
"""
94+
try:
95+
bt_client = client.get_bigtable_admin_client(
96+
project=project_id, credentials=credentials
97+
)
98+
instance = bt_client.instance(instance_id)
99+
tables = instance.list_tables()
100+
table_ids = [table.table_id for table in tables]
101+
return {"status": "SUCCESS", "results": table_ids}
102+
except Exception as ex:
103+
return {
104+
"status": "ERROR",
105+
"error_details": str(ex),
106+
}
107+
108+
109+
def get_table_info(
110+
project_id: str, instance_id: str, table_id: str, credentials: Credentials
111+
) -> dict:
112+
"""Get metadata information about a Bigtable table.
113+
114+
Args:
115+
project_id (str): The Google Cloud project id containing the instance.
116+
instance_id (str): The Bigtable instance id containing the table.
117+
table_id (str): The Bigtable table id.
118+
credentials (Credentials): The credentials to use for the request.
119+
120+
Returns:
121+
dict: Dictionary representing the properties of the table.
122+
"""
123+
try:
124+
bt_client = client.get_bigtable_admin_client(
125+
project=project_id, credentials=credentials
126+
)
127+
instance = bt_client.instance(instance_id)
128+
table = instance.table(table_id)
129+
column_families = table.list_column_families()
130+
table_info = {
131+
"project_id": project_id,
132+
"instance_id": instance.instance_id,
133+
"table_id": table.table_id,
134+
"column_families": list(column_families.keys()),
135+
}
136+
return {"status": "SUCCESS", "results": table_info}
137+
except Exception as ex:
138+
return {
139+
"status": "ERROR",
140+
"error_details": str(ex),
141+
}

0 commit comments

Comments
 (0)