feat: add Bigtable tools

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
   "click>=8.1.8, <9.0.0",                           # For CLI tools
   "fastapi>=0.115.0, <1.0.0",                       # FastAPI framework
   "google-api-python-client>=2.157.0, <3.0.0",      # Google API client discovery
+  "google-cloud-bigtable>=2.32.0",                  # For Bigtable database
   "google-cloud-aiplatform[agent_engines]>=1.95.1, <2.0.0", # For VertexAI integrations, e.g. example store.
   "google-cloud-secret-manager>=2.22.0, <3.0.0",    # Fetching secrets in RestAPI Tool
   "google-cloud-spanner>=3.56.0, <4.0.0",           # For Spanner database

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

@@ -0,0 +1,34 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Bigtable Tools (Experimental).
+
+Bigtable tools under this module are hand crafted and customized while the tools
+under google.adk.tools.google_api_tool are auto generated based on API
+definition. The rationales to have customized tool are:
+
+1. A dedicated Bigtable toolset to provide an easier, integrated way to interact
+with Bigtable for building AI Agent applications quickly.
+2. We want to provide extra access guardrails and controls in those tools.
+3. Use Bigtable Toolset for more customization and control to interact with
+Bigtable tables.
+"""
+
+from .bigtable_credentials import BigtableCredentialsConfig
+from .bigtable_toolset import BigtableToolset
+
+__all__ = [
+    "BigtableToolset",
+    "BigtableCredentialsConfig",
+]

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

@@ -0,0 +1,44 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from ...utils.feature_decorator import experimental
+from .._google_credentials import BaseGoogleCredentialsConfig
+
+BIGTABLE_TOKEN_CACHE_KEY = "bigtable_token_cache"
+BIGTABLE_DEFAULT_SCOPE = [
+    "https://www.googleapis.com/auth/bigtable.admin",
+    "https://www.googleapis.com/auth/bigtable.data",
+]
+
+
+@experimental
+class BigtableCredentialsConfig(BaseGoogleCredentialsConfig):
+  """Bigtable Credentials Configuration for Google API tools (Experimental).
+
+  Please do not use this in production, as it may be deprecated later.
+  """
+
+  def __post_init__(self) -> BigtableCredentialsConfig:
+    """Populate default scope if scopes is None."""
+    super().__post_init__()
+
+    if not self.scopes:
+      self.scopes = BIGTABLE_DEFAULT_SCOPE
+
+    # Set the token cache key
+    self._token_cache_key = BIGTABLE_TOKEN_CACHE_KEY
+
+    return self

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

@@ -0,0 +1,104 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import List
+from typing import Optional
+from typing import Union
+
+from google.adk.agents.readonly_context import ReadonlyContext
+from typing_extensions import override
+
+from . import metadata_tool
+from . import query_tool
+from ...tools.base_tool import BaseTool
+from ...tools.base_toolset import BaseToolset
+from ...tools.base_toolset import ToolPredicate
+from ...tools.google_tool import GoogleTool
+from ...utils.feature_decorator import experimental
+from .bigtable_credentials import BigtableCredentialsConfig
+from .settings import BigtableToolSettings
+
+DEFAULT_BIGTABLE_TOOL_NAME_PREFIX = "bigtable"
+
+
+@experimental
+class BigtableToolset(BaseToolset):
+  """Bigtable Toolset contains tools for interacting with Bigtable data and metadata.
+
+  The tool names are:
+    - bigtable_list_instances
+    - bigtable_get_instance_info
+    - bigtable_list_tables
+    - bigtable_get_table_info
+    - bigtable_execute_sql
+  """
+
+  def __init__(
+      self,
+      *,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+      credentials_config: Optional[BigtableCredentialsConfig] = None,
+      bigtable_tool_settings: Optional[BigtableToolSettings] = None,
+  ):
+    super().__init__(
+        tool_filter=tool_filter,
+        tool_name_prefix=DEFAULT_BIGTABLE_TOOL_NAME_PREFIX,
+    )
+    self._credentials_config = credentials_config
+    self._tool_settings = (
+        bigtable_tool_settings
+        if bigtable_tool_settings
+        else BigtableToolSettings()
+    )
+
+  def _is_tool_selected(
+      self, tool: BaseTool, readonly_context: ReadonlyContext
+  ) -> bool:
+    if self.tool_filter is None:
+      return True
+
+    if isinstance(self.tool_filter, ToolPredicate):
+      return self.tool_filter(tool, readonly_context)
+
+    if isinstance(self.tool_filter, list):
+      return tool.name in self.tool_filter
+
+    return False
+
+  @override
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> List[BaseTool]:
+    """Get tools from the toolset."""
+    all_tools = [
+        GoogleTool(
+            func=func,
+            credentials_config=self._credentials_config,
+            tool_settings=self._tool_settings,
+        )
+        for func in [
+            metadata_tool.list_instances,
+            metadata_tool.get_instance_info,
+            metadata_tool.list_tables,
+            metadata_tool.get_table_info,
+            query_tool.execute_sql,
+        ]
+    ]
+    return [
+        tool
+        for tool in all_tools
+        if self._is_tool_selected(tool, readonly_context)
+    ]

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

@@ -0,0 +1,56 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import google.api_core.client_info
+from google.auth.credentials import Credentials
+from google.cloud import bigtable
+from google.cloud.bigtable import data
+
+from ... import version
+
+USER_AGENT = f"adk-bigtable-tool google-adk/{version.__version__}"
+
+
+def _get_client_info() -> google.api_core.client_info.ClientInfo:
+  """Get client info."""
+  return google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
+
+
+def get_bigtable_data_client(
+    *, project: str, credentials: Credentials
+) -> bigtable.BigtableDataClient:
+  """Get a Bigtable client."""
+
+  bigtable_data_client = data.BigtableDataClient(
+      project=project, credentials=credentials, client_info=_get_client_info()
+  )
+
+  return bigtable_data_client
+
+
+def get_bigtable_admin_client(
+    *, project: str, credentials: Credentials
+) -> bigtable.Client:
+  """Get a Bigtable client."""
+
+  bigtable_admin_client = bigtable.Client(
+      project=project,
+      admin=True,
+      credentials=credentials,
+      client_info=_get_client_info(),
+  )
+
+  return bigtable_admin_client

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

@@ -0,0 +1,148 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+
+from google.auth.credentials import Credentials
+
+from . import client
+
+
+def list_instances(project_id: str, credentials: Credentials) -> dict:
+  """List Bigtable instance ids in a Google Cloud project.
+
+  Args:
+      project_id (str): The Google Cloud project id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary with a list of the Bigtable instance ids present in the project.
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    (instances_list, failed_locations_list) = bt_client.list_instances()
+    if failed_locations_list:
+      logging.warning(
+          "Failed to list instances from the following locations: %s",
+          failed_locations_list,
+      )
+    instance_ids = [instance.instance_id for instance in instances_list]
+    return {"status": "SUCCESS", "results": instance_ids}
+  except Exception as ex:
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }
+
+
+def get_instance_info(
+    project_id: str, instance_id: str, credentials: Credentials
+) -> dict:
+  """Get metadata information about a Bigtable instance.
+
+  Args:
+      project_id (str): The Google Cloud project id containing the instance.
+      instance_id (str): The Bigtable instance id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary representing the properties of the instance.
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    instance = bt_client.instance(instance_id)
+    instance.reload()
+    instance_info = {
+        "project_id": project_id,
+        "instance_id": instance.instance_id,
+        "display_name": instance.display_name,
+        "state": instance.state,
+        "type": instance.type_,
+        "labels": instance.labels,
+    }
+    return {"status": "SUCCESS", "results": instance_info}
+  except Exception as ex:
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }
+
+
+def list_tables(
+    project_id: str, instance_id: str, credentials: Credentials
+) -> dict:
+  """List table ids in a Bigtable instance.
+
+  Args:
+      project_id (str): The Google Cloud project id containing the instance.
+      instance_id (str): The Bigtable instance id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary with a list of the tables ids present in the instance.
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    instance = bt_client.instance(instance_id)
+    tables = instance.list_tables()
+    table_ids = [table.table_id for table in tables]
+    return {"status": "SUCCESS", "results": table_ids}
+  except Exception as ex:
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }
+
+
+def get_table_info(
+    project_id: str, instance_id: str, table_id: str, credentials: Credentials
+) -> dict:
+  """Get metadata information about a Bigtable table.
+
+  Args:
+      project_id (str): The Google Cloud project id containing the instance.
+      instance_id (str): The Bigtable instance id containing the table.
+      table_id (str): The Bigtable table id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary representing the properties of the table.
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    instance = bt_client.instance(instance_id)
+    table = instance.table(table_id)
+    column_families = table.list_column_families()
+    table_info = {
+        "project_id": project_id,
+        "instance_id": instance.instance_id,
+        "table_id": table.table_id,
+        "column_families": list(column_families.keys()),
+    }
+    return {"status": "SUCCESS", "results": table_info}
+  except Exception as ex:
+    return {
+        "status": "ERROR",
+        "error_details": str(ex),
+    }