mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: add Bigtable tools
These tools support basic operations to interact with Bigtable table metadata and query results. PiperOrigin-RevId: 796571736
This commit is contained in:
committed by
Copybara-Service
parent
fa64545a9d
commit
a953807cce
@@ -32,6 +32,7 @@ dependencies = [
|
|||||||
"click>=8.1.8, <9.0.0", # For CLI tools
|
"click>=8.1.8, <9.0.0", # For CLI tools
|
||||||
"fastapi>=0.115.0, <1.0.0", # FastAPI framework
|
"fastapi>=0.115.0, <1.0.0", # FastAPI framework
|
||||||
"google-api-python-client>=2.157.0, <3.0.0", # Google API client discovery
|
"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-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-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
|
"google-cloud-spanner>=3.56.0, <4.0.0", # For Spanner database
|
||||||
|
|||||||
@@ -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",
|
||||||
|
]
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
|
]
|
||||||
@@ -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
|
||||||
@@ -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),
|
||||||
|
}
|
||||||
@@ -0,0 +1,119 @@
|
|||||||
|
# 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
|
||||||
|
|
||||||
|
"""Tool to execute SQL queries against Bigtable."""
|
||||||
|
import json
|
||||||
|
from typing import Any
|
||||||
|
from typing import Dict
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
from google.cloud import bigtable
|
||||||
|
|
||||||
|
from . import client
|
||||||
|
from ..tool_context import ToolContext
|
||||||
|
from .settings import BigtableToolSettings
|
||||||
|
|
||||||
|
DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS = 50
|
||||||
|
|
||||||
|
|
||||||
|
def execute_sql(
|
||||||
|
project_id: str,
|
||||||
|
instance_id: str,
|
||||||
|
query: str,
|
||||||
|
credentials: Credentials,
|
||||||
|
settings: BigtableToolSettings,
|
||||||
|
tool_context: ToolContext,
|
||||||
|
) -> dict:
|
||||||
|
"""Execute a GoogleSQL query from a Bigtable table.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
project_id (str): The GCP project id in which the query should be
|
||||||
|
executed.
|
||||||
|
instance_id (str): The instance id of the Bigtable database.
|
||||||
|
query (str): The Bigtable SQL query to be executed.
|
||||||
|
credentials (Credentials): The credentials to use for the request.
|
||||||
|
settings (BigtableToolSettings): The configuration for the tool.
|
||||||
|
tool_context (ToolContext): The context for the tool.
|
||||||
|
Returns:
|
||||||
|
dict: Dictionary containing the status and the rows read.
|
||||||
|
If the result contains the key "result_is_likely_truncated" with
|
||||||
|
value True, it means that there may be additional rows matching the
|
||||||
|
query not returned in the result.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
Fetch data or insights from a table:
|
||||||
|
|
||||||
|
>>> execute_sql("my_project", "my_instance",
|
||||||
|
... "SELECT * from mytable", credentials, config, tool_context)
|
||||||
|
{
|
||||||
|
"status": "SUCCESS",
|
||||||
|
"rows": [
|
||||||
|
{
|
||||||
|
"user_id": 1,
|
||||||
|
"user_name": "Alice"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
del tool_context # Unused for now
|
||||||
|
|
||||||
|
try:
|
||||||
|
bt_client = client.get_bigtable_data_client(
|
||||||
|
project=project_id, credentials=credentials
|
||||||
|
)
|
||||||
|
eqi = bt_client.execute_query(
|
||||||
|
query=query,
|
||||||
|
instance_id=instance_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
rows: List[Dict[str, Any]] = []
|
||||||
|
max_rows = (
|
||||||
|
settings.max_query_result_rows
|
||||||
|
if settings and settings.max_query_result_rows > 0
|
||||||
|
else DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS
|
||||||
|
)
|
||||||
|
counter = max_rows
|
||||||
|
truncated = False
|
||||||
|
try:
|
||||||
|
for row in eqi:
|
||||||
|
if counter <= 0:
|
||||||
|
truncated = True
|
||||||
|
break
|
||||||
|
row_values = {}
|
||||||
|
for key, val in dict(row.fields).items():
|
||||||
|
try:
|
||||||
|
# if the json serialization of the value succeeds, use it as is
|
||||||
|
json.dumps(val)
|
||||||
|
except:
|
||||||
|
val = str(val)
|
||||||
|
row_values[key] = val
|
||||||
|
rows.append(row_values)
|
||||||
|
counter -= 1
|
||||||
|
finally:
|
||||||
|
eqi.close()
|
||||||
|
|
||||||
|
result = {"status": "SUCCESS", "rows": rows}
|
||||||
|
if truncated:
|
||||||
|
result["result_is_likely_truncated"] = True
|
||||||
|
return result
|
||||||
|
|
||||||
|
except Exception as ex:
|
||||||
|
print(ex)
|
||||||
|
return {
|
||||||
|
"status": "ERROR",
|
||||||
|
"error_details": str(ex),
|
||||||
|
}
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
# 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 pydantic import BaseModel
|
||||||
|
|
||||||
|
from ...utils.feature_decorator import experimental
|
||||||
|
|
||||||
|
|
||||||
|
@experimental('Tool settings defaults may have breaking change in the future.')
|
||||||
|
class BigtableToolSettings(BaseModel):
|
||||||
|
"""Settings for Bigtable tools."""
|
||||||
|
|
||||||
|
max_query_result_rows: int = 50
|
||||||
|
"""Maximum number of rows to return from a query result."""
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
# 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.
|
||||||
@@ -0,0 +1,91 @@
|
|||||||
|
# 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 unittest import mock
|
||||||
|
|
||||||
|
from google.adk.tools.bigtable.bigtable_credentials import BIGTABLE_DEFAULT_SCOPE
|
||||||
|
from google.adk.tools.bigtable.bigtable_credentials import BigtableCredentialsConfig
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
import google.oauth2.credentials
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
class TestBigtableCredentials:
|
||||||
|
"""Test suite for Bigtable credentials configuration validation.
|
||||||
|
|
||||||
|
This class tests the credential configuration logic that ensures
|
||||||
|
either existing credentials or client ID/secret pairs are provided.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_bigtable_credentials_config_client_id_secret(self):
|
||||||
|
"""Test BigtableCredentialsConfig with client_id and client_secret.
|
||||||
|
|
||||||
|
Ensures that when client_id and client_secret are provided, the config
|
||||||
|
object is created with the correct attributes.
|
||||||
|
"""
|
||||||
|
config = BigtableCredentialsConfig(client_id="abc", client_secret="def")
|
||||||
|
assert config.client_id == "abc"
|
||||||
|
assert config.client_secret == "def"
|
||||||
|
assert config.scopes == BIGTABLE_DEFAULT_SCOPE
|
||||||
|
assert config.credentials is None
|
||||||
|
|
||||||
|
def test_bigtable_credentials_config_existing_creds(self):
|
||||||
|
"""Test BigtableCredentialsConfig with existing generic credentials.
|
||||||
|
|
||||||
|
Ensures that when a generic Credentials object is provided, it is
|
||||||
|
stored correctly.
|
||||||
|
"""
|
||||||
|
mock_creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
config = BigtableCredentialsConfig(credentials=mock_creds)
|
||||||
|
assert config.credentials == mock_creds
|
||||||
|
assert config.client_id is None
|
||||||
|
assert config.client_secret is None
|
||||||
|
|
||||||
|
def test_bigtable_credentials_config_oauth2_creds(self):
|
||||||
|
"""Test BigtableCredentialsConfig with existing OAuth2 credentials.
|
||||||
|
|
||||||
|
Ensures that when a google.oauth2.credentials.Credentials object is
|
||||||
|
provided, the client_id, client_secret, and scopes are extracted
|
||||||
|
from the credentials object.
|
||||||
|
"""
|
||||||
|
mock_creds = mock.create_autospec(
|
||||||
|
google.oauth2.credentials.Credentials, instance=True
|
||||||
|
)
|
||||||
|
mock_creds.client_id = "oauth_client_id"
|
||||||
|
mock_creds.client_secret = "oauth_client_secret"
|
||||||
|
mock_creds.scopes = ["fake_scope"]
|
||||||
|
config = BigtableCredentialsConfig(credentials=mock_creds)
|
||||||
|
assert config.client_id == "oauth_client_id"
|
||||||
|
assert config.client_secret == "oauth_client_secret"
|
||||||
|
assert config.scopes == ["fake_scope"]
|
||||||
|
|
||||||
|
def test_bigtable_credentials_config_validation_errors(self):
|
||||||
|
"""Test BigtableCredentialsConfig validation errors.
|
||||||
|
|
||||||
|
Ensures that ValueError is raised under the following conditions:
|
||||||
|
- No arguments are provided.
|
||||||
|
- Only client_id is provided.
|
||||||
|
- Both credentials and client_id/client_secret are provided.
|
||||||
|
"""
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
BigtableCredentialsConfig()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
BigtableCredentialsConfig(client_id="abc")
|
||||||
|
|
||||||
|
mock_creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
BigtableCredentialsConfig(
|
||||||
|
credentials=mock_creds, client_id="abc", client_secret="def"
|
||||||
|
)
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
# 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.
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
from google.adk.tools.bigtable import metadata_tool
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_instances():
|
||||||
|
"""Test list_instances function."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_admin_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_instance = mock.MagicMock()
|
||||||
|
mock_instance.instance_id = "test-instance"
|
||||||
|
mock_client.list_instances.return_value = ([mock_instance], [])
|
||||||
|
|
||||||
|
creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
result = metadata_tool.list_instances("test-project", creds)
|
||||||
|
assert result == {"status": "SUCCESS", "results": ["test-instance"]}
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_instances_failed_locations():
|
||||||
|
"""Test list_instances function when some locations fail."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_admin_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
with mock.patch.object(logging, "warning") as mock_warning:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_instance = mock.MagicMock()
|
||||||
|
mock_instance.instance_id = "test-instance"
|
||||||
|
failed_locations = ["us-west1-a"]
|
||||||
|
mock_client.list_instances.return_value = (
|
||||||
|
[mock_instance],
|
||||||
|
failed_locations,
|
||||||
|
)
|
||||||
|
|
||||||
|
creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
result = metadata_tool.list_instances("test-project", creds)
|
||||||
|
assert result == {"status": "SUCCESS", "results": ["test-instance"]}
|
||||||
|
mock_warning.assert_called_once_with(
|
||||||
|
"Failed to list instances from the following locations: %s",
|
||||||
|
failed_locations,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_instance_info():
|
||||||
|
"""Test get_instance_info function."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_admin_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_instance = mock.MagicMock()
|
||||||
|
mock_client.instance.return_value = mock_instance
|
||||||
|
mock_instance.instance_id = "test-instance"
|
||||||
|
mock_instance.display_name = "Test Instance"
|
||||||
|
mock_instance.state = "READY"
|
||||||
|
mock_instance.type_ = "PRODUCTION"
|
||||||
|
mock_instance.labels = {"env": "test"}
|
||||||
|
|
||||||
|
creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
result = metadata_tool.get_instance_info(
|
||||||
|
"test-project", "test-instance", creds
|
||||||
|
)
|
||||||
|
expected_result = {
|
||||||
|
"project_id": "test-project",
|
||||||
|
"instance_id": "test-instance",
|
||||||
|
"display_name": "Test Instance",
|
||||||
|
"state": "READY",
|
||||||
|
"type": "PRODUCTION",
|
||||||
|
"labels": {"env": "test"},
|
||||||
|
}
|
||||||
|
assert result == {"status": "SUCCESS", "results": expected_result}
|
||||||
|
mock_instance.reload.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_tables():
|
||||||
|
"""Test list_tables function."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_admin_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_instance = mock.MagicMock()
|
||||||
|
mock_client.instance.return_value = mock_instance
|
||||||
|
mock_table = mock.MagicMock()
|
||||||
|
mock_table.table_id = "test-table"
|
||||||
|
mock_instance.list_tables.return_value = [mock_table]
|
||||||
|
|
||||||
|
creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
result = metadata_tool.list_tables("test-project", "test-instance", creds)
|
||||||
|
assert result == {"status": "SUCCESS", "results": ["test-table"]}
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_table_info():
|
||||||
|
"""Test get_table_info function."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_admin_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_instance = mock.MagicMock()
|
||||||
|
mock_client.instance.return_value = mock_instance
|
||||||
|
mock_table = mock.MagicMock()
|
||||||
|
mock_instance.table.return_value = mock_table
|
||||||
|
mock_table.table_id = "test-table"
|
||||||
|
mock_instance.instance_id = "test-instance"
|
||||||
|
mock_table.list_column_families.return_value = {"cf1": mock.MagicMock()}
|
||||||
|
|
||||||
|
creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
result = metadata_tool.get_table_info(
|
||||||
|
"test-project", "test-instance", "test-table", creds
|
||||||
|
)
|
||||||
|
expected_result = {
|
||||||
|
"project_id": "test-project",
|
||||||
|
"instance_id": "test-instance",
|
||||||
|
"table_id": "test-table",
|
||||||
|
"column_families": ["cf1"],
|
||||||
|
}
|
||||||
|
assert result == {"status": "SUCCESS", "results": expected_result}
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
# 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 Optional
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
from google.adk.tools.base_tool import BaseTool
|
||||||
|
from google.adk.tools.bigtable import BigtableCredentialsConfig
|
||||||
|
from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset
|
||||||
|
from google.adk.tools.bigtable.query_tool import execute_sql
|
||||||
|
from google.adk.tools.bigtable.settings import BigtableToolSettings
|
||||||
|
from google.adk.tools.tool_context import ToolContext
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
from google.cloud import bigtable
|
||||||
|
from google.cloud.bigtable.data.execute_query import ExecuteQueryIterator
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_sql_basic():
|
||||||
|
"""Test execute_sql tool basic functionality."""
|
||||||
|
project = "my_project"
|
||||||
|
instance_id = "my_instance"
|
||||||
|
query = "SELECT * FROM my_table"
|
||||||
|
credentials = mock.create_autospec(Credentials, instance=True)
|
||||||
|
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||||
|
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_data_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True)
|
||||||
|
mock_client.execute_query.return_value = mock_iterator
|
||||||
|
|
||||||
|
# Mock row data
|
||||||
|
mock_row = mock.MagicMock()
|
||||||
|
mock_row.fields = {"col1": "val1", "col2": 123}
|
||||||
|
mock_iterator.__iter__.return_value = [mock_row]
|
||||||
|
|
||||||
|
result = execute_sql(
|
||||||
|
project_id=project,
|
||||||
|
instance_id=instance_id,
|
||||||
|
credentials=credentials,
|
||||||
|
query=query,
|
||||||
|
settings=BigtableToolSettings(),
|
||||||
|
tool_context=tool_context,
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_rows = [{"col1": "val1", "col2": 123}]
|
||||||
|
assert result == {"status": "SUCCESS", "rows": expected_rows}
|
||||||
|
mock_client.execute_query.assert_called_once_with(
|
||||||
|
query=query, instance_id=instance_id
|
||||||
|
)
|
||||||
|
mock_iterator.close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_sql_truncated():
|
||||||
|
"""Test execute_sql tool truncation functionality."""
|
||||||
|
project = "my_project"
|
||||||
|
instance_id = "my_instance"
|
||||||
|
query = "SELECT * FROM my_table"
|
||||||
|
credentials = mock.create_autospec(Credentials, instance=True)
|
||||||
|
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||||
|
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_data_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_iterator = mock.create_autospec(ExecuteQueryIterator, instance=True)
|
||||||
|
mock_client.execute_query.return_value = mock_iterator
|
||||||
|
|
||||||
|
# Mock row data
|
||||||
|
mock_row1 = mock.MagicMock()
|
||||||
|
mock_row1.fields = {"col1": "val1"}
|
||||||
|
mock_row2 = mock.MagicMock()
|
||||||
|
mock_row2.fields = {"col1": "val2"}
|
||||||
|
mock_iterator.__iter__.return_value = [mock_row1, mock_row2]
|
||||||
|
|
||||||
|
result = execute_sql(
|
||||||
|
project_id=project,
|
||||||
|
instance_id=instance_id,
|
||||||
|
credentials=credentials,
|
||||||
|
query=query,
|
||||||
|
settings=BigtableToolSettings(max_query_result_rows=1),
|
||||||
|
tool_context=tool_context,
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_rows = [{"col1": "val1"}]
|
||||||
|
assert result == {
|
||||||
|
"status": "SUCCESS",
|
||||||
|
"rows": expected_rows,
|
||||||
|
"result_is_likely_truncated": True,
|
||||||
|
}
|
||||||
|
mock_client.execute_query.assert_called_once_with(
|
||||||
|
query=query, instance_id=instance_id
|
||||||
|
)
|
||||||
|
mock_iterator.close.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_sql_error():
|
||||||
|
"""Test execute_sql tool error handling."""
|
||||||
|
project = "my_project"
|
||||||
|
instance_id = "my_instance"
|
||||||
|
query = "SELECT * FROM my_table"
|
||||||
|
credentials = mock.create_autospec(Credentials, instance=True)
|
||||||
|
tool_context = mock.create_autospec(ToolContext, instance=True)
|
||||||
|
|
||||||
|
with mock.patch(
|
||||||
|
"google.adk.tools.bigtable.client.get_bigtable_data_client"
|
||||||
|
) as mock_get_client:
|
||||||
|
mock_client = mock.MagicMock()
|
||||||
|
mock_get_client.return_value = mock_client
|
||||||
|
mock_client.execute_query.side_effect = Exception("Test error")
|
||||||
|
|
||||||
|
result = execute_sql(
|
||||||
|
project_id=project,
|
||||||
|
instance_id=instance_id,
|
||||||
|
credentials=credentials,
|
||||||
|
query=query,
|
||||||
|
settings=BigtableToolSettings(),
|
||||||
|
tool_context=tool_context,
|
||||||
|
)
|
||||||
|
assert result == {"status": "ERROR", "error_details": "Test error"}
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
# 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 unittest import mock
|
||||||
|
|
||||||
|
from google.adk.tools.bigtable import BigtableCredentialsConfig
|
||||||
|
from google.adk.tools.bigtable import metadata_tool
|
||||||
|
from google.adk.tools.bigtable import query_tool
|
||||||
|
from google.adk.tools.bigtable.bigtable_toolset import BigtableToolset
|
||||||
|
from google.adk.tools.bigtable.bigtable_toolset import DEFAULT_BIGTABLE_TOOL_NAME_PREFIX
|
||||||
|
from google.adk.tools.google_tool import GoogleTool
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_bigtable_toolset_name_prefix():
|
||||||
|
"""Test Bigtable toolset name prefix."""
|
||||||
|
credentials_config = BigtableCredentialsConfig(
|
||||||
|
client_id="abc", client_secret="def"
|
||||||
|
)
|
||||||
|
toolset = BigtableToolset(credentials_config=credentials_config)
|
||||||
|
assert toolset.tool_name_prefix == DEFAULT_BIGTABLE_TOOL_NAME_PREFIX
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bigtable_toolset_tools_default():
|
||||||
|
"""Test default Bigtable toolset."""
|
||||||
|
credentials_config = BigtableCredentialsConfig(
|
||||||
|
client_id="abc", client_secret="def"
|
||||||
|
)
|
||||||
|
toolset = BigtableToolset(credentials_config=credentials_config)
|
||||||
|
|
||||||
|
tools = await toolset.get_tools()
|
||||||
|
assert tools is not None
|
||||||
|
|
||||||
|
assert len(tools) == 5
|
||||||
|
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||||
|
|
||||||
|
expected_tool_names = set([
|
||||||
|
"list_instances",
|
||||||
|
"get_instance_info",
|
||||||
|
"list_tables",
|
||||||
|
"get_table_info",
|
||||||
|
"execute_sql",
|
||||||
|
])
|
||||||
|
actual_tool_names = set([tool.name for tool in tools])
|
||||||
|
assert actual_tool_names == expected_tool_names
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"selected_tools",
|
||||||
|
[
|
||||||
|
pytest.param([], id="None"),
|
||||||
|
pytest.param(
|
||||||
|
["list_instances", "get_instance_info"], id="instance-metadata"
|
||||||
|
),
|
||||||
|
pytest.param(["list_tables", "get_table_info"], id="table-metadata"),
|
||||||
|
pytest.param(["execute_sql"], id="query"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bigtable_toolset_tools_selective(selected_tools):
|
||||||
|
"""Test Bigtable toolset with filter.
|
||||||
|
|
||||||
|
This test verifies the behavior of the Bigtable toolset when filter is
|
||||||
|
specified. A use case for this would be when the agent builder wants to
|
||||||
|
use only a subset of the tools provided by the toolset.
|
||||||
|
"""
|
||||||
|
credentials_config = BigtableCredentialsConfig(
|
||||||
|
client_id="abc", client_secret="def"
|
||||||
|
)
|
||||||
|
toolset = BigtableToolset(
|
||||||
|
credentials_config=credentials_config, tool_filter=selected_tools
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await toolset.get_tools()
|
||||||
|
assert tools is not None
|
||||||
|
|
||||||
|
assert len(tools) == len(selected_tools)
|
||||||
|
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||||
|
|
||||||
|
expected_tool_names = set(selected_tools)
|
||||||
|
actual_tool_names = set([tool.name for tool in tools])
|
||||||
|
assert actual_tool_names == expected_tool_names
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("selected_tools", "returned_tools"),
|
||||||
|
[
|
||||||
|
pytest.param(["unknown"], [], id="all-unknown"),
|
||||||
|
pytest.param(
|
||||||
|
["unknown", "execute_sql"],
|
||||||
|
["execute_sql"],
|
||||||
|
id="mixed-known-unknown",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bigtable_toolset_unknown_tool(selected_tools, returned_tools):
|
||||||
|
"""Test Bigtable toolset with filter.
|
||||||
|
|
||||||
|
This test verifies the behavior of the Bigtable toolset when filter is
|
||||||
|
specified with an unknown tool.
|
||||||
|
"""
|
||||||
|
credentials_config = BigtableCredentialsConfig(
|
||||||
|
client_id="abc", client_secret="def"
|
||||||
|
)
|
||||||
|
|
||||||
|
toolset = BigtableToolset(
|
||||||
|
credentials_config=credentials_config, tool_filter=selected_tools
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await toolset.get_tools()
|
||||||
|
assert tools is not None
|
||||||
|
|
||||||
|
assert len(tools) == len(returned_tools)
|
||||||
|
assert all([isinstance(tool, GoogleTool) for tool in tools])
|
||||||
|
|
||||||
|
expected_tool_names = set(returned_tools)
|
||||||
|
actual_tool_names = set([tool.name for tool in tools])
|
||||||
|
assert actual_tool_names == expected_tool_names
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
# 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 unittest import mock
|
||||||
|
|
||||||
|
from google.adk.tools.bigtable import client
|
||||||
|
from google.auth.credentials import Credentials
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_bigtable_data_client():
|
||||||
|
"""Test get_bigtable_client function."""
|
||||||
|
with mock.patch(
|
||||||
|
"google.cloud.bigtable.data.BigtableDataClient"
|
||||||
|
) as MockBigtableDataClient:
|
||||||
|
mock_creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
client.get_bigtable_data_client(
|
||||||
|
project="test-project", credentials=mock_creds
|
||||||
|
)
|
||||||
|
MockBigtableDataClient.assert_called_once_with(
|
||||||
|
project="test-project",
|
||||||
|
credentials=mock_creds,
|
||||||
|
client_info=mock.ANY,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_bigtable_admin_client():
|
||||||
|
"""Test get_bigtable_admin_client function."""
|
||||||
|
with mock.patch("google.cloud.bigtable.Client") as BigtableDataClient:
|
||||||
|
mock_creds = mock.create_autospec(Credentials, instance=True)
|
||||||
|
client.get_bigtable_admin_client(
|
||||||
|
project="test-project", credentials=mock_creds
|
||||||
|
)
|
||||||
|
# Admin client is a BigtableDataClient created with admin=True.
|
||||||
|
BigtableDataClient.assert_called_once_with(
|
||||||
|
project="test-project",
|
||||||
|
admin=True,
|
||||||
|
credentials=mock_creds,
|
||||||
|
client_info=mock.ANY,
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user