diff --git a/contributing/samples/spanner/agent.py b/contributing/samples/spanner/agent.py index fa3c3a95..d0f570cd 100644 --- a/contributing/samples/spanner/agent.py +++ b/contributing/samples/spanner/agent.py @@ -16,14 +16,21 @@ import os from google.adk.agents.llm_agent import LlmAgent from google.adk.auth.auth_credential import AuthCredentialTypes +from google.adk.tools.google_tool import GoogleTool from google.adk.tools.spanner.settings import Capabilities from google.adk.tools.spanner.settings import SpannerToolSettings from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig from google.adk.tools.spanner.spanner_toolset import SpannerToolset +import google.adk.tools.spanner.utils as spanner_tool_utils +from google.adk.tools.tool_context import ToolContext import google.auth +from google.auth.credentials import Credentials +from google.cloud.spanner_v1 import param_types as spanner_param_types # Define an appropriate credential type -CREDENTIALS_TYPE = AuthCredentialTypes.OAUTH2 +# Set to None to use the application default credentials (ADC) for a quick +# development. +CREDENTIALS_TYPE = None # Define Spanner tool config with read capability set to allowed. @@ -56,10 +63,115 @@ else: credentials=application_default_credentials ) +# Example 1: Use tools from the Spanner toolset. +# For example, data exploration agents help the Spanner database developer or +# data engineer of the organization. spanner_toolset = SpannerToolset( - credentials_config=credentials_config, spanner_tool_settings=tool_settings + credentials_config=credentials_config, + spanner_tool_settings=tool_settings, + # Uncomment to explicitly specify allowed tools. + # tool_filter=["execute_sql", "get_table_schema"], ) + +# Replace the following settings with your specific Spanner database for example +# 2 and 3. +# For example, these settings can also be read from a configuration file or +# environment variables. +_SPANNER_PROJECT_ID = "" +_SPANNER_INSTANCE_ID = "" +_SPANNER_DATABASE_ID = "" + + +# Example 2: Create a customized Spanner query tool with a template SQL query. +# Note that this approach makes it **more vulnerable to SQL injection**. This +# might be suitable for some specific use cases, and **adding additional checks +# or callbacks** is recommended. +def count_rows_in_table( + table_name: str, + credentials: Credentials, + settings: SpannerToolSettings, + tool_context: ToolContext, +): + """Counts the total number of rows for a specified table. + + Args: + table_name: The name of the table for which to count rows. + + Returns: + The total number of rows in the table. + """ + + # Example of adding additional checks: + # if table_name not in ["table1", "table2"]: + # raise ValueError("Table name is not allowed.") + + sql_template = f"SELECT COUNT(*) FROM {table_name}" + + return spanner_tool_utils.execute_sql( + project_id=_SPANNER_PROJECT_ID, + instance_id=_SPANNER_INSTANCE_ID, + database_id=_SPANNER_DATABASE_ID, + query=sql_template, + credentials=credentials, + settings=settings, + tool_context=tool_context, + ) + + +# Example 3: Create a customized Spanner query tool with a template +# parameterized SQL query. +# For example, it could query data that all authenticated users of the system +# have access to. This can also work for searching public knowledge bases, such +# as company policies and FAQs. +def search_hotels( + location_name: str, + credentials: Credentials, + settings: SpannerToolSettings, + tool_context: ToolContext, +): + """Search hotels for a specific location. + + This function takes a geographical location name and returns a list of hotels + in that area, including key details for each. + + Args: + location_name (str): The geographical location (e.g., city or town) for the + hotel search. + Example: + { + "location_name": "Seattle" + } + Example: + { + "location_name": "New York" + } + Example: + { + "location_name": "Los Angeles" + } + + Returns: + The hotels name, rating and description. + """ + + sql_template = """ + SELECT name, rating, description FROM hotels + WHERE location_name = @location_name + """ + return spanner_tool_utils.execute_sql( + project_id=_SPANNER_PROJECT_ID, + instance_id=_SPANNER_INSTANCE_ID, + database_id=_SPANNER_DATABASE_ID, + query=sql_template, + credentials=credentials, + settings=settings, + tool_context=tool_context, + params={"location_name": location_name}, + params_types={"location_name": spanner_param_types.STRING}, + ) + + # The variable name `root_agent` determines what your root agent is for the # debug CLI root_agent = LlmAgent( @@ -73,5 +185,19 @@ root_agent = LlmAgent( You are a data agent with access to several Spanner tools. Make use of those tools to answer the user's questions. """, - tools=[spanner_toolset], + tools=[ + # Use tools from Spanner toolset. + spanner_toolset, + # Or, uncomment to use customized Spanner tools. + # GoogleTool( + # func=count_rows_in_table, + # credentials_config=credentials_config, + # tool_settings=tool_settings, + # ), + # GoogleTool( + # func=search_hotels, + # credentials_config=credentials_config, + # tool_settings=tool_settings, + # ), + ], ) diff --git a/src/google/adk/tools/spanner/query_tool.py b/src/google/adk/tools/spanner/query_tool.py index e317a0ce..f6fa5a34 100644 --- a/src/google/adk/tools/spanner/query_tool.py +++ b/src/google/adk/tools/spanner/query_tool.py @@ -14,17 +14,12 @@ from __future__ import annotations -import json - from google.auth.credentials import Credentials -from google.cloud.spanner_admin_database_v1.types import DatabaseDialect -from . import client +from . import utils from ..tool_context import ToolContext from .settings import SpannerToolSettings -DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS = 50 - def execute_sql( project_id: str, @@ -68,47 +63,12 @@ def execute_sql( Note: This is running with Read-Only Transaction for query that only read data. """ - - try: - # Get Spanner client - spanner_client = client.get_spanner_client( - project=project_id, credentials=credentials - ) - instance = spanner_client.instance(instance_id) - database = instance.database(database_id) - - if database.database_dialect == DatabaseDialect.POSTGRESQL: - return { - "status": "ERROR", - "error_details": "PostgreSQL dialect is not supported.", - } - - with database.snapshot() as snapshot: - result_set = snapshot.execute_sql(query) - rows = [] - counter = ( - settings.max_executed_query_result_rows - if settings and settings.max_executed_query_result_rows > 0 - else DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS - ) - for row in result_set: - try: - # if the json serialization of the row succeeds, use it as is - json.dumps(row) - except: - row = str(row) - - rows.append(row) - counter -= 1 - if counter <= 0: - break - - result = {"status": "SUCCESS", "rows": rows} - if counter <= 0: - result["result_is_likely_truncated"] = True - return result - except Exception as ex: - return { - "status": "ERROR", - "error_details": str(ex), - } + return utils.execute_sql( + project_id, + instance_id, + database_id, + query, + credentials, + settings, + tool_context, + ) diff --git a/src/google/adk/tools/spanner/utils.py b/src/google/adk/tools/spanner/utils.py new file mode 100644 index 00000000..64c03859 --- /dev/null +++ b/src/google/adk/tools/spanner/utils.py @@ -0,0 +1,107 @@ +# 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 json +from typing import Optional + +from google.auth.credentials import Credentials +from google.cloud.spanner_admin_database_v1.types import DatabaseDialect + +from . import client +from ..tool_context import ToolContext +from .settings import SpannerToolSettings + +DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS = 50 + + +def execute_sql( + project_id: str, + instance_id: str, + database_id: str, + query: str, + credentials: Credentials, + settings: SpannerToolSettings, + tool_context: ToolContext, + params: Optional[dict] = None, + params_types: Optional[dict] = None, +) -> dict: + """Utility function to run a Spanner Read-Only query in the spanner database and return the result. + + Args: + project_id (str): The GCP project id in which the spanner database + resides. + instance_id (str): The instance id of the spanner database. + database_id (str): The database id of the spanner database. + query (str): The Spanner SQL query to be executed. + credentials (Credentials): The credentials to use for the request. + settings (SpannerToolSettings): The settings for the tool. + tool_context (ToolContext): The context for the tool. + params (dict): values for parameter replacement. Keys must match the + names used in ``query``. + params_types (dict): maps explicit types for one or more param values. + + Returns: + dict: Dictionary with the result of the query. + 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. + """ + + try: + # Get Spanner client + spanner_client = client.get_spanner_client( + project=project_id, credentials=credentials + ) + instance = spanner_client.instance(instance_id) + database = instance.database(database_id) + + if database.database_dialect == DatabaseDialect.POSTGRESQL: + return { + "status": "ERROR", + "error_details": "PostgreSQL dialect is not supported.", + } + + with database.snapshot() as snapshot: + result_set = snapshot.execute_sql( + sql=query, params=params, param_types=params_types + ) + rows = [] + counter = ( + settings.max_executed_query_result_rows + if settings and settings.max_executed_query_result_rows > 0 + else DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS + ) + for row in result_set: + try: + # if the json serialization of the row succeeds, use it as is + json.dumps(row) + except: + row = str(row) + + rows.append(row) + counter -= 1 + if counter <= 0: + break + + result = {"status": "SUCCESS", "rows": rows} + if counter <= 0: + result["result_is_likely_truncated"] = True + return result + except Exception as ex: + return { + "status": "ERROR", + "error_details": str(ex), + }